cyclegan-project/dataset.py at main · Hasith2005/cyclegan-project · GitHub

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
from torch.utils.data import Dataset
import os
import numpy as np
import torch
import math
# audio_A = non-lofi audio
# audio_B = lofi audio

class AudioDataset(Dataset):
    def __init__(self,audios_A_path, audios_B_path):
        super().__init__()
        self.audios_A_path = audios_A_path # path to audios A
        self.audios_B_path = audios_B_path # path to audios B
        self.audios_A_contents=os.listdir(self.audios_A_path)
        self.audios_B_contents=os.listdir(self.audios_B_path)
        self.len_A = len(os.listdir(self.audios_A_path))
        self.len_B = len(os.listdir(self.audios_B_path))
        self.len_ds= min(self.len_A,self.len_B)   # fine because the datasets are nearly the same size.
    def __len__(self):
        return self.len_ds


    def __getitem__(self, i):
        SEGMENT_LEN = 1292
        audio_A = np.load(os.path.join(self.audios_A_path, self.audios_A_contents[i]))
        audio_B = np.load(os.path.join(self.audios_B_path, self.audios_B_contents[i]))

        # Match lengths
        min_AB = min(audio_A.shape[1], audio_B.shape[1])
        audio_A = audio_A[:, :min_AB]
        audio_B = audio_B[:, :min_AB]
        start = np.random.randint(0, min_AB - SEGMENT_LEN + 1)  # +1 to be inclusive
        audio_A = audio_A[:, start:start+SEGMENT_LEN]
        audio_B = audio_B[:, start:start+SEGMENT_LEN]

        return torch.tensor(audio_A), torch.tensor(audio_B)