|
import os
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy as np
|
|
import pedalboard as pb
|
|
import torch
|
|
import torchaudio as ta
|
|
from torch.utils import data
|
|
|
|
from models.bandit.core.data._types import AudioDict, DataDict
|
|
|
|
|
|
class BaseSourceSeparationDataset(data.Dataset, ABC):
|
|
def __init__(
|
|
self, split: str,
|
|
stems: List[str],
|
|
files: List[str],
|
|
data_path: str,
|
|
fs: int,
|
|
npy_memmap: bool,
|
|
recompute_mixture: bool
|
|
):
|
|
self.split = split
|
|
self.stems = stems
|
|
self.stems_no_mixture = [s for s in stems if s != "mixture"]
|
|
self.files = files
|
|
self.data_path = data_path
|
|
self.fs = fs
|
|
self.npy_memmap = npy_memmap
|
|
self.recompute_mixture = recompute_mixture
|
|
|
|
@abstractmethod
|
|
def get_stem(
|
|
self,
|
|
*,
|
|
stem: str,
|
|
identifier: Dict[str, Any]
|
|
) -> torch.Tensor:
|
|
raise NotImplementedError
|
|
|
|
def _get_audio(self, stems, identifier: Dict[str, Any]):
|
|
audio = {}
|
|
for stem in stems:
|
|
audio[stem] = self.get_stem(stem=stem, identifier=identifier)
|
|
|
|
return audio
|
|
|
|
def get_audio(self, identifier: Dict[str, Any]) -> AudioDict:
|
|
|
|
if self.recompute_mixture:
|
|
audio = self._get_audio(
|
|
self.stems_no_mixture,
|
|
identifier=identifier
|
|
)
|
|
audio["mixture"] = self.compute_mixture(audio)
|
|
return audio
|
|
else:
|
|
return self._get_audio(self.stems, identifier=identifier)
|
|
|
|
@abstractmethod
|
|
def get_identifier(self, index: int) -> Dict[str, Any]:
|
|
pass
|
|
|
|
def compute_mixture(self, audio: AudioDict) -> torch.Tensor:
|
|
|
|
return sum(
|
|
audio[stem] for stem in audio if stem != "mixture"
|
|
)
|
|
|
