FunASR/funasr/models/eend/eend_ola_dataloader.py

import logging

import kaldiio
import numpy as np
import torch
from torch.utils.data import DataLoader
from torch.utils.data import Dataset


def custom_collate(batch):
    keys, speech, speaker_labels, orders = zip(*batch)
    speech = [torch.from_numpy(np.copy(sph)).to(torch.float32) for sph in speech]
    speaker_labels = [torch.from_numpy(np.copy(spk)).to(torch.float32) for spk in speaker_labels]
    orders = [torch.from_numpy(np.copy(o)).to(torch.int64) for o in orders]
    batch = dict(speech=speech,
                 speaker_labels=speaker_labels,
                 orders=orders)

    return keys, batch


class EENDOLADataset(Dataset):
    def __init__(
            self,
            data_file,
    ):
        self.data_file = data_file
        with open(data_file) as f:
            lines = f.readlines()
        self.samples = [line.strip().split() for line in lines]
        logging.info("total samples: {}".format(len(self.samples)))

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

    def __getitem__(self, idx):
        key, speech_path, speaker_label_path = self.samples[idx]
        speech = kaldiio.load_mat(speech_path)
        speaker_label = kaldiio.load_mat(speaker_label_path).reshape(speech.shape[0], -1)

        order = np.arange(speech.shape[0])
        np.random.shuffle(order)

        return key, speech, speaker_label, order


class EENDOLADataLoader():
    def __init__(self, data_file, batch_size, shuffle=True, num_workers=8):
        dataset = EENDOLADataset(data_file)
        self.data_loader = DataLoader(dataset,
                                      batch_size=batch_size,
                                      collate_fn=custom_collate,
                                      shuffle=shuffle,
                                      num_workers=num_workers)

    def build_iter(self, epoch):
        return self.data_loader