Add Thumos dataset

config.py

@@ -115,6 +115,9 @@ def get_args_parser():
     parser.add_argument("--margin", default=1.0, type=float)
 
     # dataset parameters
+    parser.add_argument(
+        "--dataset", default="thumos", type=str, help="thumos or tvseries"
+    )
     parser.add_argument("--dataset_file", type=str, default="data/data_info_new.json")
     parser.add_argument("--frozen_weights", type=str, default=None)
     parser.add_argument(

dataset.py

@@ -1,9 +1,135 @@
+from os import minor
 import os.path as osp
-import pickle
+import pickle5 as pickle
 import torch
 import torch.utils.data as data
 import numpy as np
-from ipdb import set_trace
+from pathlib import Path
+
+
+class TRNTVSeriesDataLayer(data.Dataset):
+    def __init__(self, args, phase="train"):
+        self.pickle_root = "data/"
+        self.sessions = getattr(args, phase + "_session_set")  # video name
+        self.enc_steps = args.enc_layers
+        self.numclass = args.numclass
+        self.dec_steps = args.query_num
+        self.training = phase == "train"
+        self.feature_pretrain = args.feature
+        assert (Path(self.pickle_root) / self.feature_pretrain).suffix in {
+            ".pickle",
+            ".pkl",
+        }
+        self.inputs = []
+
+        # Load anno
+        target_all = pickle.load(
+            open(
+                osp.join(self.pickle_root, "tvseries_anno.pickle"),
+                "rb",
+            )
+        )
+
+        # Load features
+        assert osp.exists(osp.join(self.pickle_root, self.feature_pretrain))
+        self.feature_All = pickle.load(
+            open(
+                osp.join(self.pickle_root, self.feature_pretrain),
+                "rb",
+            )
+        )
+
+        # Ensure that features and annotataions have same length
+        for k, v in self.feature_All.items():
+            min_len = min(v["rgb"].shape[0], v["flow"].shape[0])
+            if v["rgb"].shape[0] != min_len:
+                v["rgb"] = v["rgb"][:min_len]
+
+            if v["flow"].shape[0] != min_len:
+                v["flow"] = v["flow"][:min_len]
+
+            if target_all[k].shape[0] != min_len:
+                target_all[k] = target_all[k][:min_len]
+
+        for session in self.sessions:
+            target = target_all[session]
+
+            seed = np.random.randint(self.enc_steps) if self.training else 0
+            for start, end in zip(
+                range(seed, target.shape[0], 1),
+                range(seed + self.enc_steps, target.shape[0] - self.dec_steps, 1),
+            ):
+                enc_target = target[start:end]
+                dec_target = target[end : end + self.dec_steps]
+                distance_target, class_h_target = self.get_distance_target(
+                    target[start:end]
+                )
+                self.inputs.append(
+                    [
+                        session,
+                        start,
+                        end,
+                        enc_target,
+                        distance_target,
+                        class_h_target,
+                        dec_target,
+                    ]
+                )
+
+        print(f"Loaded {self.feature_pretrain}")
+
+    def get_dec_target(self, target_vector):
+        target_matrix = np.zeros(
+            (self.enc_steps, self.dec_steps, target_vector.shape[-1])
+        )
+        for i in range(self.enc_steps):
+            for j in range(self.dec_steps):
+                # 0 -> [1, 2, 3]
+                # target_matrix[i,j] = target_vector[i+j+1,:]
+                # 0 -> [0, 1, 2]
+                target_matrix[i, j] = target_vector[i + j, :]
+        return target_matrix
+
+    def get_distance_target(self, target_vector):
+        target_matrix = np.zeros(self.enc_steps - 1)
+        target_argmax = target_vector[self.enc_steps - 1].argmax()
+        for i in range(self.enc_steps - 1):
+            if target_vector[i].argmax() == target_argmax:
+                target_matrix[i] = 1.0
+        return target_matrix, target_vector[self.enc_steps - 1]
+
+    def __getitem__(self, index):
+        """self.inputs.append([
+            session, start, end, enc_target, distance_target, class_h_target
+        ])"""
+        (
+            session,
+            start,
+            end,
+            enc_target,
+            distance_target,
+            class_h_target,
+            dec_target,
+        ) = self.inputs[index]
+        camera_inputs = self.feature_All[session]["rgb"][start:end]
+        camera_inputs = torch.tensor(camera_inputs)
+        motion_inputs = self.feature_All[session]["flow"][start:end]
+        motion_inputs = torch.tensor(motion_inputs)
+        enc_target = torch.tensor(enc_target)
+        distance_target = torch.tensor(distance_target)
+        class_h_target = torch.tensor(class_h_target)
+        dec_target = torch.tensor(dec_target)
+        return (
+            camera_inputs,
+            motion_inputs,
+            enc_target,
+            distance_target,
+            class_h_target,
+            dec_target,
+        )
+
+    def __len__(self):
+        return len(self.inputs)
 
 
 class TRNTHUMOSDataLayer(data.Dataset):
@@ -260,7 +386,6 @@ def __getitem__(self, index):
         camera_inputs = self.feature_All[session]["rgb"][start:end]
         camera_inputs = torch.tensor(camera_inputs)
         motion_inputs = self.feature_All[session]["flow"][start:end]
-
         motion_inputs = torch.tensor(motion_inputs)
         enc_target = torch.tensor(enc_target)
         distance_target = torch.tensor(distance_target)

main.py

@@ -14,7 +14,7 @@
 import utils
 
 import transformer_models
-from dataset import TRNTHUMOSDataLayer
+from dataset import TRNTHUMOSDataLayer, TRNTVSeriesDataLayer
 from train import train_one_epoch, evaluate
 from test import test_one_epoch
 import torch.nn as nn
@@ -123,8 +123,11 @@ def main(args):
     )
     lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
 
-    dataset_train = TRNTHUMOSDataLayer(phase="train", args=args)
-    dataset_val = TRNTHUMOSDataLayer(phase="test", args=args)
+    DataLayer = {"thumos": TRNTHUMOSDataLayer, "tvseries": TRNTVSeriesDataLayer}[
+        args.dataset
+    ]
+    dataset_train = DataLayer(phase="train", args=args)
+    dataset_val = DataLayer(phase="test", args=args)
 
     if args.distributed:
         sampler_train = DistributedSampler(dataset_train)
@@ -259,7 +262,7 @@ def main(args):
     args = parser.parse_args()
     # args.dataset = osp.basename(osp.normpath(args.data_root)).upper()
     with open(args.dataset_file, "r") as f:
-        data_info = json.load(f)["THUMOS"]
+        data_info = json.load(f)[args.dataset.upper()]
     args.train_session_set = data_info["train_session_set"]
     args.test_session_set = data_info["test_session_set"]
     args.class_index = data_info["class_index"]