File size: 8,851 Bytes
5085882
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
import json
import random
from tqdm import tqdm
import torch
import decord

decord.bridge.set_bridge("torch")
import torchaudio
from math import ceil
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np


class AudioVisualDataset(Dataset):
    """Can sample data from audio-visual databases
    Params:
    min_video_frames: used to drop short video clips
    video_resize: resize for CLIP processing
    sampling_rate: audio sampling rate
    max_clip_len: max length (seconds) of audiovisual clip to be sampled
    num_sample_frames: number of image frames to be uniformly sampled from video
    """

    def __init__(
        self,
        datafiles=[
            "/mnt/bn/data-xubo/dataset/audioset_videos/datafiles/audioset_balanced_train.json",
        ],
        min_video_frames=30,
        video_resize=[224, 224],
        sampling_rate=16000,
        sample_av_clip=True,
        max_clip_len=10,
        num_sample_frames=10,
        # hyparameters used for SpecAug
        freqm=48,
        timem=192,
        return_label=False,
    ):
        all_data_json = []
        for datafile in datafiles:
            with open(datafile, "r") as fp:
                data_json = json.load(fp)["data"]
                all_data_json.extend(data_json)

        # drop short video clips
        self.all_data_json = [
            data
            for data in all_data_json
            if int(data["video_shape"][0]) >= min_video_frames
        ]

        self.max_clip_len = max_clip_len
        self.video_resize = video_resize
        self.sampling_rate = sampling_rate
        self.sample_av_clip = sample_av_clip
        self.num_sample_frames = num_sample_frames
        self.corresponding_audio_len = self.sampling_rate * self.max_clip_len

        # hyparameters used for AudioMAE
        self.freqm = freqm
        self.timem = timem
        self.norm_mean = -4.2677393
        self.norm_std = 4.5689974
        self.melbins = 128
        self.TARGET_LEN = 1024

        self.return_label = return_label
        if self.return_label:
            self.audioset_label2idx = self._prepare_audioset()

    def __len__(self):
        return len(self.all_data_json)

    def _read_audio_video(self, index):
        try:
            video_path = self.all_data_json[index]["mp4"]
            # read audio
            ar = decord.AudioReader(
                video_path, sample_rate=self.sampling_rate, mono=True
            )
            # read video frames
            vr = decord.VideoReader(
                video_path,
                height=self.video_resize[0],
                width=self.video_resize[1],
            )

            labels = self.all_data_json[index]["labels"]
            return vr, ar, labels

        except Exception as e:
            print(f"error: {e} occurs, when loading {video_path}")
            random_index = random.randint(0, len(self.all_data_json) - 1)
            return self._read_audio_video(index=random_index)

    def _prepare_audioset(self):
        df1 = pd.read_csv(
            "/mnt/bn/lqhaoheliu/datasets/audioset/metadata/class_labels_indices.csv",
            delimiter=",",
            skiprows=0,
        )
        label_set = df1.to_numpy()
        code2id = {}
        for i in range(len(label_set)):
            code2id[label_set[i][1]] = label_set[i][0]
        return code2id

    def __getitem__(self, index):
        # read audio and video
        vr, ar, labels = self._read_audio_video(index)

        # create a audio tensor
        audio_data = ar[:]  # [1, samples]
        audio_len = audio_data.shape[1] / self.sampling_rate
        audio_data = audio_data.squeeze(0)  # [samples]

        # create a video tensor
        full_vid_length = len(vr)
        video_rate = ceil(vr.get_avg_fps())
        samples_per_frame = float(self.sampling_rate) / video_rate
        start_frame = 0

        # sample video clip
        if audio_len > self.max_clip_len and self.sample_av_clip:
            start_frame = random.randint(
                0, max(full_vid_length - video_rate * self.max_clip_len, 0)
            )
        end_frame = min(start_frame + video_rate * self.max_clip_len, full_vid_length)
        video_data = vr.get_batch(range(start_frame, end_frame))

        # sample audio clip
        if audio_len > self.max_clip_len and self.sample_av_clip:
            # corresponding_audio_len = int(video_data.size()[0] * samples_per_frame)
            corresponding_audio_start = int(start_frame * samples_per_frame)
            audio_data = audio_data[corresponding_audio_start:]

        # cut or pad audio clip with respect to the sampled video clip
        if audio_data.shape[0] < self.corresponding_audio_len:
            zero_data = torch.zeros(self.corresponding_audio_len)
            zero_data[: audio_data.shape[0]] = audio_data
            audio_data = zero_data
        elif audio_data.shape[0] > self.corresponding_audio_len:
            audio_data = audio_data[: self.corresponding_audio_len]

        # uniformly sample image frames from video [tentative solution]
        interval = video_data.shape[0] // self.num_sample_frames
        video_data = video_data[::interval][: self.num_sample_frames]

        assert (
            video_data.shape[0] == self.num_sample_frames
        ), f"number of sampled image frames is {video_data.shape[0]}"

        assert (
            audio_data.shape[0] == self.corresponding_audio_len
        ), f"number of audio samples is {audio_data.shape[0]}"

        # video transformation
        video_data = video_data / 255.0
        video_data = video_data.permute(0, 3, 1, 2)  # [N, H, W, C] -> [N, C, H, W]

        # calculate mel fbank of waveform for audio encoder
        audio_data = audio_data.unsqueeze(0)  # [1, samples]
        audio_data = audio_data - audio_data.mean()
        fbank = torchaudio.compliance.kaldi.fbank(
            audio_data,
            htk_compat=True,
            sample_frequency=self.sampling_rate,
            use_energy=False,
            window_type="hanning",
            num_mel_bins=self.melbins,
            dither=0.0,
            frame_shift=10,
        )
        # cut and pad
        n_frames = fbank.shape[0]
        p = self.TARGET_LEN - n_frames
        if p > 0:
            m = torch.nn.ZeroPad2d((0, 0, 0, p))
            fbank = m(fbank)
        elif p < 0:
            fbank = fbank[0 : self.TARGET_LEN, :]

        # SpecAug for training (not for eval)
        freqm = torchaudio.transforms.FrequencyMasking(self.freqm)
        timem = torchaudio.transforms.TimeMasking(self.timem)
        fbank = fbank.transpose(0, 1).unsqueeze(0)  # 1, 128, 1024 (...,freq,time)
        if self.freqm != 0:
            fbank = freqm(fbank)
        if self.timem != 0:
            fbank = timem(fbank)  # (..., freq, time)
        fbank = torch.transpose(fbank.squeeze(), 0, 1)  # time, freq
        fbank = (fbank - self.norm_mean) / (self.norm_std * 2)
        fbank = fbank.unsqueeze(0)

        if self.return_label:
            # get audioset lebel indexes
            label_indices = np.zeros(527)

            for label_str in labels.split(","):
                label_indices[int(self.audioset_label2idx[label_str])] = 1.0

            label_indices = torch.FloatTensor(label_indices)

            data_dict = {
                "labels": label_indices,
                "images": video_data,
                "fbank": fbank,
                # 'modality': 'audio_visual'
            }

        else:
            data_dict = {
                "images": video_data,
                "fbank": fbank,
                # 'modality': 'audio_visual'
            }

        return data_dict


def collate_fn(list_data_dict):
    r"""Collate mini-batch data to inputs and targets for training.

    Args:
        list_data_dict: e.g., [
            {'vocals': (channels_num, segment_samples),
             'accompaniment': (channels_num, segment_samples),
             'mixture': (channels_num, segment_samples)
            },
            {'vocals': (channels_num, segment_samples),
             'accompaniment': (channels_num, segment_samples),
             'mixture': (channels_num, segment_samples)
            },
            ...]

    Returns:
        data_dict: e.g. {
            'vocals': (batch_size, channels_num, segment_samples),
            'accompaniment': (batch_size, channels_num, segment_samples),
            'mixture': (batch_size, channels_num, segment_samples)
            }
    """

    data_dict = {}
    for key in list_data_dict[0].keys():
        # for key in ['waveform']:
        # try:
        data_dict[key] = [data_dict[key] for data_dict in list_data_dict]
        # except:
        #     from IPython import embed; embed(using=False); os._exit(0)

        data_dict[key] = torch.stack(data_dict[key])

    return data_dict