Refactor frame classification models to use single WindowedFramesDatapipe #574
Description
I think VocalDataset can be rewritten to be more general, and a lot of the logic moved into transforms.
This gives us more flexibility while also making the code more concise.
E.g., the following much simpler version of VocalDataset could be combined with the right transforms to give us what we have now and optionally work with other things, e.g. a model that uses audio as input. The transform should include loading audio, spectrogram files, etc. This would also make it easier to move to DataPipes should we decide to do so.
from typing import Callable, Optional, Sequence
import pandas as pd
# TODO: use vocles
from ...typing import PathLike
RETURNS_COLUMNS_MAP = {
'spect': 'spect_path',
'audio': 'audio_path',
'annot': 'annot_path',
}
VOCAL_DATASET_ITEM_KEYS = list(RETURNS_COLUMNS_MAP.keys())
class VocalDataset:
"""Class representing a dataset of vocalizations,
that can include audio, spectrograms, and annotations."
def __init__(self,
csv_path: PathLike,
returns: Sequence[str] = ('spect', 'annot'),
transforms : Optional[Callable] = None,
):
self.voc_df = pd.read_csv(csv_path)
if not all([return_ in ('audio', 'spect', 'annot')
for return_ in returns]):
raise ValueError(
f"Values for 'returns' must all be in: {{'audio', 'spect', 'annot'}} "
f"but got '{returns}'"
)
self.returns = returns
self.transforms = transforms
def __getitem__(self, idx):
voc_row = self.voc_df.iloc[idx, :]
item = {
key: voc_row[RETURNS_COLUMNS_MAP[key]] if key in self.returns else None
for key, val in RETURNS_COLUMNS_MAP.items()
}
if self.transforms:
item = self.transforms(item)
return item