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()