get_transcript.py

from torch.utils.data import DataLoader, Dataset

import argparse
import os
import numpy as np
import random
import torch


class MyDataset(Dataset):
    def __init__(self, dataset_name, root, split_name, sample_rate, sample_type):
        assert dataset_name in ('breakfast', 'hollywood', 'crosstask', '50salads', 'gtea')
        self.dataset_name = dataset_name
        self.root = os.path.join(root, dataset_name)
        self.sample_rate = sample_rate
        self.sample_type = sample_type
        self.max_len = 0
        self.video_lst, self.gts, self.trans, self.n_cls = self.load_data(split_name)

        self.bg_cls = 0  # background class id
        if dataset_name == 'crosstask':
            self.feat_dim = 3200
        else:
            self.feat_dim = 2048

    def load_data(self, split_name):
        # load video names
        samples = []
        with open(os.path.join(self.root, 'splits', split_name), 'r') as f:
            for line in f:
                line = line.strip()
                line = os.path.splitext(line)
                samples.append(line[0])

        # load label2idx mapping
        label2idx = {}
        with open(os.path.join(self.root, 'mapping.txt'), 'r') as f:
            for line in f:
                line = line.strip().split()
                label2idx[line[1]] = int(line[0])

        # read labels and transcripts
        gts, trans = [], []
        if not os.path.exists(os.path.join(self.root, 'transcripts')):
            os.mkdir(os.path.join(self.root, 'transcripts'))
        self.create_transcript(samples, label2idx)
        for name in samples:
            with open(os.path.join(self.root, 'groundTruth', name + '.txt'), 'r') as f:
                self.max_len = max(self.max_len, len(f.readlines()))
                gt = [label2idx[line.strip()] for line in f]
            with open(os.path.join(self.root, 'transcripts', name + '.txt'), 'r') as f:
                tr = [label2idx[line.strip()] for line in f]
            gts.append(gt)
            trans.append(tr)
        return samples, gts, trans, len(label2idx)

    def create_transcript(self, samples, label2idx):
        for name in samples:
            with open(os.path.join(self.root, 'groundTruth', name + '.txt'), 'r') as f:
                gt = [line.strip() for line in f]
            gt = self.deduplicate_keep_order(gt)
            file_path = os.path.join(self.root, 'transcripts', name + '.txt')
            self.write_list_to_txt(file_path, gt)

    @staticmethod
    def deduplicate_keep_order(lst):
        if not lst:  
            return []

        result = [lst[0]]
        for item in lst[1:]:
            if item != result[-1]:
                result.append(item)
        return result

    @staticmethod
    def write_list_to_txt(file_path, data_list):
        """
        Write list content to TXT file line by line

        Parameters:
            file_path (str): Path of the file to be created
            data_list (list): List data to be written
        """
        try:
            with open(file_path, 'w', encoding='utf-8') as file:
                for item in data_list:
                    file.write(f"{item}\n")

            print(f"Successfully wrote {len(data_list)} items to file: {file_path}")
            return True

        except Exception as e:
            print(f"Error writing to file: {e}")
            return False

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

    def __getitem__(self, idx):
        feat = np.load(os.path.join(self.root, 'features', self.video_lst[idx] + '.npy'))  # (t, c)
        if feat.dtype == np.float64:
            feat = feat.astype(np.float32)
        if self.dataset_name == 'breakfast' or self.dataset_name == 'crosstask':
            feat = feat.T  # (t, c)
        gt = np.array(self.gts[idx])
        tr = np.array(self.trans[idx])
        if self.dataset_name == 'hollywood':
            diff = feat.shape[0] - gt.shape[0]
            if diff > 0:
                feat = feat[:gt.shape[0]]
            elif diff < 0:
                gt = gt[:feat.shape[0]]
        assert feat.shape[0] == gt.shape[0]

        raw_gt = gt.copy()
        raw_len = gt.shape[0]

        sampled_ts = self.sampling_fun(feat.shape[0], self.sample_rate, self.sample_type)
        feat, gt = feat[sampled_ts], gt[sampled_ts]

        vid_la = set(tr)
        multihot = np.zeros(self.n_cls)
        for la in vid_la:
            multihot[la] = 1

        ret = {
            'name': self.video_lst[idx],
            'feat': feat,
            'gt': gt,
            'transcript': tr,
            'multi_hot': multihot,

            'raw_gt': raw_gt,
            'raw_len': raw_len,
        }
        return ret

    def sampling_fun(self, T, GAP, sample_type):
        '''
        ref: DPDTW (CVPR21)
        '''
        start_idxes = list(range(0, T, GAP))
        N = len(start_idxes)
        idxes = start_idxes + [T]

        sample_ts = []
        for i in range(N):
            start_i = idxes[i]
            end_i = idxes[i + 1] - 1
            assert start_i <= end_i, (start_i, end_i)
            if sample_type == 'mid':
                sample_ts.append(int((start_i + end_i) / 2))
            elif sample_type == 'rand':
                sample_ts.append(random.randint(start_i, end_i))
            else:
                raise ValueError('Unknown sample method: {}'.format(sample_type))
        return sample_ts


def collate_fn(sample):
    max_len = max([s["feat"].shape[0] for s in sample])

    name_lst, feat_lst, gt_lst, mask_lst = [], [], [], []
    for s in sample:
        name_lst.append(s['name'])
        feat, gt = s['feat'], s['gt']
        t = feat.shape[0]
        pad_t = max_len - t
        feat = np.pad(feat, ((0, pad_t), (0, 0)), mode='constant', constant_values=0)
        gt = np.pad(gt, (0, pad_t), mode='constant', constant_values=0)
        feat, gt = torch.from_numpy(feat), torch.from_numpy(gt)
        mask = torch.zeros(max_len)
        mask[:t] = torch.ones(t)
        feat_lst.append(feat)
        gt_lst.append(gt)
        mask_lst.append(mask.bool())

    feat_lst = torch.stack(feat_lst, dim=0)  # (b, t, c)
    gt_lst = torch.stack(gt_lst, dim=0)  # (b, t)
    mask_lst = torch.stack(mask_lst, dim=0)  # (b, t)
    tr_lst = [torch.LongTensor(s['transcript']) for s in sample]
    mh_lst = [torch.from_numpy(s['multi_hot']).int() for s in sample]
    mh_lst = torch.stack(mh_lst, dim=0)  # (b, cls)

    raw_gt_lst = [torch.LongTensor(s['raw_gt']) for s in sample]
    raw_len_lst = torch.LongTensor([s['raw_len'] for s in sample])

    ret = {
        'name': name_lst,
        'feat': feat_lst,
        'gt': gt_lst,
        'transcript': tr_lst,
        'mask': mask_lst,
        'multi_hot': mh_lst,

        'raw_gt': raw_gt_lst,
        'raw_len': raw_len_lst,
    }
    return ret


def get_dataloader(data, root, split=1, sample_rate=1, sample_type="mid", batch_size=32, num_workers=8):
    train_data = MyDataset(
        data,
        root,
        f"train.split{split}.bundle",
        sample_rate,
        sample_type,
    )
    train_loader = DataLoader(
        train_data,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=num_workers,
    )
    print(f"Train dataset length: {len(train_data)}")
    print(f"Train dataset max length: {train_data.max_len}")

    """
    test data
    """
    test_data = MyDataset(
        data,
        root,
        f"test.split{split}.bundle",
        sample_rate,
        sample_type,
    )
    test_loader = DataLoader(
        test_data,
        batch_size=batch_size,
        collate_fn=collate_fn,
        shuffle=False,
        num_workers=num_workers,
    )
    print(f"Test dataset length: {len(test_data)}")
    print(f"Test dataset max length: {test_data.max_len}")
    return train_loader, test_loader


def run(data, root, split=1, sample_rate=1, sample_type="mid", batch_size=32, num_workers=8):
    train_loader, test_loader = get_dataloader(
        data=data,
        root=root,
        split=split,
        sample_rate=sample_rate,
        sample_type=sample_type,
        batch_size=batch_size,
        num_workers=num_workers,
    )
    return train_loader, test_loader


def build_parser():
    parser = argparse.ArgumentParser(description="Create transcripts and build train/test dataloaders.")
    parser.add_argument("--dataset", type=str, default="gtea", choices=["breakfast", "hollywood", "crosstask", "50salads", "gtea"])
    parser.add_argument("--root", type=str, default="data", help="Dataset root directory")
    parser.add_argument("--sample-rate", type=int, default=1)
    parser.add_argument("--sample-type", type=str, default="mid", choices=["mid", "rand"])
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--num-workers", type=int, default=8)
    return parser


def main():
    args = build_parser().parse_args()
    run(
        data=args.dataset,
        root=args.root,
        sample_rate=args.sample_rate,
        sample_type=args.sample_type,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
    )


if __name__ == '__main__':
    main()