segment_finder.py 4.32 KB
Newer Older
Matteo's avatar
update    
Matteo committed
1
2
import os
import tempfile
Matteo's avatar
update  
Matteo committed
3
4
5
from uuid import uuid4

from mpai_cae_arp.audio import AudioWave, Noise
Matteo's avatar
update    
Matteo committed
6
from mpai_cae_arp.files import File, FileType
Matteo's avatar
update  
Matteo committed
7
from mpai_cae_arp.types.irregularity import Irregularity, IrregularityFile, Source
Matteo's avatar
update    
Matteo committed
8
from mpai_cae_arp.time import frames_to_seconds, seconds_to_frames, seconds_to_string, time_to_seconds
Matteo's avatar
update  
Matteo committed
9

Matteo's avatar
update    
Matteo committed
10
11
temp_dir = tempfile.gettempdir()
TMP_CHANNELS_MAP = os.path.join(temp_dir, "channels_map.json")
Matteo's avatar
update  
Matteo committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

def calculate_offset(audio: AudioWave, video: AudioWave) -> float:
    """
    Calculates the offset between two audio files based on their cross-correlation.

    Parameters
    ----------
    audio : AudioWave
        The audio file to be used as reference.
    video : AudioWave
        The audio file to be used as target.
    
    Returns
    -------
    float
    """

Matteo's avatar
update    
Matteo committed
29
30
31
    # corr = np.correlate(audio.array, video.array, mode="full")
    # lags = np.arange(-len(audio.array) + 1, len(video.array))
    # lag_idx = np.argmax(np.abs(corr))
Matteo's avatar
update  
Matteo committed
32

Matteo's avatar
update    
Matteo committed
33
34
    # return lags[lag_idx] / audio.samplerate
    return 150
Matteo's avatar
update  
Matteo committed
35
36


Matteo's avatar
update    
Matteo committed
37
def get_irregularities_from_audio(audio_src: AudioWave) -> list[Irregularity]:
Matteo's avatar
update  
Matteo committed
38
    input_channels: list[AudioWave] = []
Matteo's avatar
update    
Matteo committed
39
40
41
42
43
44

    if audio_src.channels > 1:
        for channel in range(audio_src.channels):
            input_channels.append(audio_src.get_channel(channel))
    else:
        input_channels.append(audio_src)
Matteo's avatar
update    
Matteo committed
45
46

    channels_map = {}
Matteo's avatar
update  
Matteo committed
47
48

    irreg_list: list[Irregularity] = []
Matteo's avatar
update    
Matteo committed
49
    for idx, audio in enumerate(input_channels):
Matteo's avatar
update  
Matteo committed
50
51
52
53
54
55
        for _, noise_list in audio.get_silence_slices([
            Noise("A", -50, -63),
            Noise("B", -63, -69),
            Noise("C", -69, -72)],
            length=500).items():
            for start, _ in noise_list:
Matteo's avatar
update    
Matteo committed
56
                id = uuid4()
Matteo's avatar
update  
Matteo committed
57
58
                irreg_list.append(
                    Irregularity(
Matteo's avatar
update    
Matteo committed
59
                        irregularity_ID=id,
Matteo's avatar
update  
Matteo committed
60
                        source=Source.AUDIO,
Matteo's avatar
update    
Matteo committed
61
                        time_label= seconds_to_string(frames_to_seconds(start, audio.samplerate))
Matteo's avatar
update  
Matteo committed
62
63
                    )
                )
Matteo's avatar
update    
Matteo committed
64
                channels_map[str(id)] = idx
Matteo's avatar
update    
Matteo committed
65
66

    File(TMP_CHANNELS_MAP, FileType.JSON).write_content(channels_map)
Matteo's avatar
update  
Matteo committed
67
68
69

    return irreg_list

Matteo's avatar
update    
Matteo committed
70
71
72
73
74
75

def create_irreg_file(audio_src: str, video_src: str) -> IrregularityFile:

    audio = AudioWave.from_file(audio_src, bufferize=True)
    
    offset = calculate_offset(audio, video_src)
Matteo's avatar
update    
Matteo committed
76
77
78
79
80
    irregularities = get_irregularities_from_audio(audio)

    irregularities.sort(key=lambda x: time_to_seconds(x.time_label))
    
    return IrregularityFile(irregularities=irregularities, offset=offset)
Matteo's avatar
update    
Matteo committed
81
82
83
84


def merge_irreg_files(
    file1: IrregularityFile,
Matteo's avatar
update    
Matteo committed
85
86
87
88
89
90
91
92
93
94
    file2: IrregularityFile
) -> IrregularityFile:

    match file1.offset, file2.offset:
        case None, _:
            offset=file2.offset
        case _, None:
            offset=file1.offset
        case _, _:
            offset=max(file1.offset, file2.offset)
Matteo's avatar
update    
Matteo committed
95

Matteo's avatar
update    
Matteo committed
96
97
98
99
100
    irregularities = file1.irregularities + file2.irregularities
    irregularities.sort(key=lambda x: time_to_seconds(x.time_label))

    new_file = IrregularityFile(
        irregularities=irregularities, offset=offset)
Matteo's avatar
update    
Matteo committed
101
102
103
104
105

    return new_file


def extract_audio_irregularities(
Matteo's avatar
update    
Matteo committed
106
    audio_src: str,
Matteo's avatar
update    
Matteo committed
107
    irreg_file: IrregularityFile,
Matteo's avatar
update    
Matteo committed
108
    path: str
Matteo's avatar
update    
Matteo committed
109
) -> IrregularityFile:
Matteo's avatar
update    
Matteo committed
110

Matteo's avatar
update    
Matteo committed
111
    channels_map = File(TMP_CHANNELS_MAP, FileType.JSON).get_content()
Matteo's avatar
update    
Matteo committed
112
113
114
    os.makedirs(f"{path}/AudioBlocks", exist_ok=True)

    audio = AudioWave.from_file(audio_src, bufferize=True)
Matteo's avatar
update    
Matteo committed
115
    for irreg in irreg_file.irregularities:
Matteo's avatar
update    
Matteo committed
116
117
        if channels_map.get(str(irreg.irregularity_ID)) is None:
            audio[seconds_to_frames(
Matteo's avatar
update    
Matteo committed
118
                        time_to_seconds(irreg.time_label), audio.samplerate
Matteo's avatar
update    
Matteo committed
119
                    ):seconds_to_frames(
Matteo's avatar
update    
Matteo committed
120
                        time_to_seconds(irreg.time_label), audio.samplerate)+audio.samplerate//2]\
Matteo's avatar
update    
Matteo committed
121
122
123
124
125
126
127
128
                .save(f"{path}/AudioBlocks/{irreg.irregularity_ID}.wav")
        else:
            audio.get_channel(channels_map[str(irreg.irregularity_ID)])[
                    seconds_to_frames(
                        time_to_seconds(irreg.time_label), audio.samplerate
                    ):seconds_to_frames(
                        time_to_seconds(irreg.time_label), audio.samplerate)+audio.samplerate//2]\
                .save(f"{path}/AudioBlocks/{irreg.irregularity_ID}.wav")
Matteo's avatar
update    
Matteo committed
129
        irreg.audio_block_URI = f"{path}/AudioBlocks/{irreg.irregularity_ID}.wav"
Matteo's avatar
update    
Matteo committed
130
    os.remove(TMP_CHANNELS_MAP)
Matteo's avatar
update    
Matteo committed
131
132

    return irreg_file