mlsa_neural_vocoder/dataset.py at main · DwangoMediaVillage/mlsa_neural_vocoder · 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
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
import os
from typing import List, TypedDict
import numpy as np
import torch
from torch.utils.data import Dataset

from dataclasses import dataclass

from config import PreProcessConfig
from utils.pad import pad_1D, pad_2D


@dataclass
class SynthDatasetItem:
    id: str
    f0: np.ndarray
    mcep: np.ndarray
    apdc: np.ndarray
    wav: np.ndarray


class SynthDatasetBatchNumpy(TypedDict):
    ids: list[str]
    frame_f0s: np.ndarray
    mceps: np.ndarray
    apdcs: np.ndarray
    mcep_lens: np.ndarray
    wavs: np.ndarray


class SynthDatasetBatchTorch(TypedDict):
    ids: list[str]
    frame_f0s: torch.Tensor
    mceps: torch.Tensor
    apdcs: torch.Tensor
    mcep_lens: torch.LongTensor
    wavs: torch.Tensor


class SynthDataset(Dataset):
    def __init__(
        self,
        config: PreProcessConfig,
        filename: str,
        max_len: int,
        convert_torch: bool = True,
        torch_device: str = "cpu",
    ):
        self.config = config
        self.data_path = config.path.preprocessed_path
        self.max_len = max_len
        self.convert_torch = convert_torch
        self.torch_device = torch_device

        self.basename: List[str] = []

        with open(os.path.join(self.data_path, filename), "r", encoding="utf-8") as f:
            for line in f.readlines():
                n = line.strip("\n")
                self.basename.append(n)

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

    def __getitem__(self, idx):
        basename = self.basename[idx]

        pitch_path = os.path.join(
            self.data_path,
            "pitch",
            "pitch-{}.npy".format(basename),
        )
        pitch = np.load(pitch_path)

        wav_path = os.path.join(
            self.data_path,
            "wav",
            "wav-{}.npy".format(basename),
        )
        wav = np.load(wav_path)

        mcep_path = os.path.join(
            self.data_path,
            "mcep",
            "mcep-{}.npy".format(basename),
        )
        mcep = np.load(mcep_path)

        apdc_path = os.path.join(
            self.data_path,
            "apdc",
            "apdc-{}.npy".format(basename),
        )
        apdc = np.load(apdc_path)

        sample = SynthDatasetItem(
            id=basename,
            f0=pitch,
            mcep=mcep,
            apdc=apdc,
            wav=wav,
        )

        return sample

    def collate_fn(self, batch: List[SynthDatasetItem]):
        ids = [sample.id for sample in batch]
        frame_f0s = [sample.f0 for sample in batch]
        mceps = [sample.mcep for sample in batch]
        apdcs = [sample.apdc for sample in batch]
        wavs = [sample.wav for sample in batch]

        hop_length = self.config.stft.hop_length

        for sample_index in range(len(frame_f0s)):
            f0_item = frame_f0s[sample_index]
            if len(f0_item) > self.max_len:
                i = np.random.random_integers(0, len(f0_item) - self.max_len)
                frame_f0s[sample_index] = f0_item[i: i + self.max_len]
                mceps[sample_index] = mceps[sample_index][i: i + self.max_len]
                apdcs[sample_index] = apdcs[sample_index][i: i + self.max_len]
                wavs[sample_index] = wavs[sample_index][i * hop_length: (i + self.max_len) * hop_length]

        mcep_lens = np.array(
            [mcep.shape[0] for mcep in mceps], dtype=np.int64
        )

        frame_f0s = pad_1D(frame_f0s).astype(np.float32)
        mceps = pad_2D(mceps).astype(np.float32)
        apdcs = pad_2D(apdcs).astype(np.float32)
        wavs = pad_1D(wavs)

        res = {
            "ids": ids,
            "frame_f0s": frame_f0s,
            "mceps": mceps,
            "apdcs": apdcs,
            "mcep_lens": mcep_lens,
            "wavs": wavs,
        }

        if self.convert_torch:
            res = {k: torch.from_numpy(v) if k != "ids" else v for k, v in res.items()}

        return res