data_util/dcase2016task2.py

import json
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import pandas as pd
import soundfile as sf
import torch
from intervaltree import IntervalTree
from torch.utils.data import Dataset


class FixCropDataset(Dataset):
    """
    Read in a JSON file and return audio and audio filenames
    """

    def __init__(self, data: Dict,
                 audio_dir: Path,
                 sample_rate: int,
                 label_fps: int,
                 label_to_idx: Dict,
                 nlabels: int):
        self.clip_len = 120
        self.target_len = 10
        self.pieces_per_clip = self.clip_len // self.target_len
        self.filenames = list(data.keys())
        self.audio_dir = audio_dir
        assert self.audio_dir.is_dir(), f"{audio_dir} is not a directory"
        self.sample_rate = sample_rate
        # all files are 120 seconds long, split them into 12 x 10 second pieces
        self.pieces = []
        self.labels = []
        self.timestamps = []
        for filename in self.filenames:
            self.pieces += [(filename, i) for i in range(self.pieces_per_clip)]
            labels = data[filename]
            frame_len = 1000 / label_fps
            timestamps = np.arange(label_fps * self.clip_len) * frame_len + 0.5 * frame_len
            timestamp_labels = get_labels_for_timestamps(labels, timestamps)
            ys = []
            for timestamp_label in timestamp_labels:
                timestamp_label_idxs = [label_to_idx[str(event)] for event in timestamp_label]
                y_timestamp = label_to_binary_vector(timestamp_label_idxs, nlabels)
                ys.append(y_timestamp)
            ys = torch.stack(ys)
            frames_per_clip = ys.size(0) // self.pieces_per_clip
            self.labels += [ys[frames_per_clip * i: frames_per_clip * (i + 1)] for i in range(self.pieces_per_clip)]
            self.timestamps += [timestamps[frames_per_clip * i: frames_per_clip * (i + 1)] for i in
                                range(self.pieces_per_clip)]

        assert len(self.labels) == len(self.pieces) == len(self.filenames) * self.pieces_per_clip

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

    def __getitem__(self, idx):
        filename = self.pieces[idx][0]
        piece = self.pieces[idx][1]
        audio_path = self.audio_dir.joinpath(filename)
        audio, sr = sf.read(str(audio_path), dtype=np.float32)
        assert sr == self.sample_rate
        start = self.sample_rate * piece * self.target_len
        end = start + self.sample_rate * self.target_len
        audio = audio[start:end]
        return audio, self.labels[idx].transpose(0, 1), filename, self.timestamps[idx]


class RandomCropDataset(Dataset):
    """
    Read in a JSON file and return audio and audio filenames
    """

    def __init__(self, data: Dict,
                 audio_dir: Path,
                 sample_rate: int,
                 label_fps: int,
                 label_to_idx: Dict,
                 nlabels: int):
        self.clip_len = 120
        self.target_len = 10
        self.pieces_per_clip = self.clip_len // self.target_len
        self.filenames = list(data.keys())
        self.audio_dir = audio_dir
        assert self.audio_dir.is_dir(), f"{audio_dir} is not a directory"
        self.sample_rate = sample_rate
        self.label_fps = label_fps
        # all files are 120 seconds long, randomly crop 10 seconds snippets
        self.labels = []
        self.timestamps = []
        for filename in self.filenames:
            labels = data[filename]
            frame_len = 1000 / label_fps
            timestamps = np.arange(label_fps * self.clip_len) * frame_len + 0.5 * frame_len
            timestamp_labels = get_labels_for_timestamps(labels, timestamps)
            ys = []
            for timestamp_label in timestamp_labels:
                timestamp_label_idxs = [label_to_idx[str(event)] for event in timestamp_label]
                y_timestamp = label_to_binary_vector(timestamp_label_idxs, nlabels)
                ys.append(y_timestamp)
            ys = torch.stack(ys)
            self.labels.append(ys)
            self.timestamps.append(timestamps)

        assert len(self.labels) == len(self.filenames)

    def __len__(self):
        return len(self.filenames) * self.clip_len // self.target_len

    def __getitem__(self, idx):
        idx = idx % len(self.filenames)
        filename = self.filenames[idx]
        audio_path = self.audio_dir.joinpath(filename)
        audio, sr = sf.read(str(audio_path), dtype=np.float32)
        assert sr == self.sample_rate

        # crop random 10 seconds piece
        labels_to_pick = self.target_len * self.label_fps
        max_offset = len(self.labels[idx]) - labels_to_pick + 1
        offset = torch.randint(max_offset, (1,)).item()
        labels = self.labels[idx][offset:offset + labels_to_pick]
        scale = self.sample_rate // self.label_fps
        audio = audio[offset * scale:offset * scale + labels_to_pick * scale]
        timestamps = self.timestamps[idx][offset:offset + labels_to_pick]
        return audio, labels.transpose(0, 1), filename, timestamps


def get_training_dataset(
        task_path,
        sample_rate=16000,
        label_fps=25,
        wavmix_p=0.0,
        random_crop=True
):
    task_path = Path(task_path)

    label_vocab, nlabels = label_vocab_nlabels(task_path)
    label_to_idx = label_vocab_as_dict(label_vocab, key="label", value="idx")

    train_fold = task_path.joinpath("train.json")
    audio_dir = task_path.joinpath(str(sample_rate), "train")
    train_fold_data = json.load(train_fold.open())
    if random_crop:
        dataset = RandomCropDataset(train_fold_data, audio_dir, sample_rate, label_fps, label_to_idx, nlabels)
    else:
        dataset = FixCropDataset(train_fold_data, audio_dir, sample_rate, label_fps, label_to_idx, nlabels)
    if wavmix_p > 0:
        dataset = MixupDataset(dataset, rate=wavmix_p)
    return dataset


def get_validation_dataset(
        task_path,
        sample_rate=16000,
        label_fps=25,
):
    task_path = Path(task_path)

    label_vocab, nlabels = label_vocab_nlabels(task_path)
    label_to_idx = label_vocab_as_dict(label_vocab, key="label", value="idx")

    valid_fold = task_path.joinpath("valid.json")
    audio_dir = task_path.joinpath(str(sample_rate), "valid")
    valid_fold_data = json.load(valid_fold.open())
    dataset = FixCropDataset(valid_fold_data, audio_dir, sample_rate, label_fps, label_to_idx, nlabels)
    return dataset


def get_test_dataset(
        task_path,
        sample_rate=16000,
        label_fps=25,
):
    task_path = Path(task_path)

    label_vocab, nlabels = label_vocab_nlabels(task_path)
    label_to_idx = label_vocab_as_dict(label_vocab, key="label", value="idx")

    test_fold = task_path.joinpath("test.json")
    audio_dir = task_path.joinpath(str(sample_rate), "test")
    test_fold_data = json.load(test_fold.open())
    dataset = FixCropDataset(test_fold_data, audio_dir, sample_rate, label_fps, label_to_idx, nlabels)
    return dataset


def get_labels_for_timestamps(labels: List, timestamps: np.ndarray) -> List:
    # A list of labels present at each timestamp
    tree = IntervalTree()
    # Add all events to the label tree
    for event in labels:
        # We add 0.0001 so that the end also includes the event
        tree.addi(event["start"], event["end"] + 0.0001, event["label"])

    timestamp_labels = []
    # Update the binary vector of labels with intervals for each timestamp
    for j, t in enumerate(timestamps):
        interval_labels: List[str] = [interval.data for interval in tree[t]]
        timestamp_labels.append(interval_labels)
        # If we want to store the timestamp too
        # labels_for_sound.append([float(t), interval_labels])

    assert len(timestamp_labels) == len(timestamps)
    return timestamp_labels


def label_vocab_nlabels(task_path: Path) -> Tuple[pd.DataFrame, int]:
    label_vocab = pd.read_csv(task_path.joinpath("labelvocabulary.csv"))

    nlabels = len(label_vocab)
    assert nlabels == label_vocab["idx"].max() + 1
    return (label_vocab, nlabels)


def label_vocab_as_dict(df: pd.DataFrame, key: str, value: str) -> Dict:
    """
    Returns a dictionary of the label vocabulary mapping the label column to
    the idx column. key sets whether the label or idx is the key in the dict. The
    other column will be the value.
    """
    if key == "label":
        # Make sure the key is a string
        df["label"] = df["label"].astype(str)
        value = "idx"
    else:
        assert key == "idx", "key argument must be either 'label' or 'idx'"
        value = "label"
    return df.set_index(key).to_dict()[value]


def label_to_binary_vector(label: List, num_labels: int) -> torch.Tensor:
    """
    Converts a list of labels into a binary vector
    Args:
        label: list of integer labels
        num_labels: total number of labels

    Returns:
        A float Tensor that is multi-hot binary vector
    """
    # Lame special case for multilabel with no labels
    if len(label) == 0:
        # BCEWithLogitsLoss wants float not long targets
        binary_labels = torch.zeros((num_labels,), dtype=torch.float)
    else:
        binary_labels = torch.zeros((num_labels,)).scatter(0, torch.tensor(label), 1.0)

    # Validate the binary vector we just created
    assert set(torch.where(binary_labels == 1.0)[0].numpy()) == set(label)
    return binary_labels


class MixupDataset(Dataset):
    """ Mixing Up wave forms
    """

    def __init__(self, dataset, beta=0.2, rate=0.5):
        self.beta = beta
        self.rate = rate
        self.dataset = dataset
        print(f"Mixing up waveforms from dataset of len {len(dataset)}")

    def __getitem__(self, index):
        if torch.rand(1) < self.rate:
            batch1 = self.dataset[index]
            idx2 = torch.randint(len(self.dataset), (1,)).item()
            batch2 = self.dataset[idx2]
            x1, x2 = batch1[0], batch2[0]
            y1, y2 = batch1[1], batch2[1]
            l = np.random.beta(self.beta, self.beta)
            l = max(l, 1. - l)
            x1 = x1 - x1.mean()
            x2 = x2 - x2.mean()
            x = (x1 * l + x2 * (1. - l))
            x = x - x.mean()
            y = (y1 * l + y2 * (1. - l))
            return x, y, batch1[2], batch1[3]
        return self.dataset[index]

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