Update dataset.py

dataset.py

@@ -97,7 +97,7 @@ class VideoDataSet(data.Dataset):
                 self.feature_rgb_file = {}
                 self.feature_flow_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt["video_feature_all_train"], file + '.npz')
+                    feature_path = opt["video_feature_all_train"] + file + '.npz'
                     if not os.path.exists(feature_path):
                         raise ValueError(f"Feature file {feature_path} not found")
                     feature_All[file] = np.load(feature_path)['feats']
@@ -110,7 +110,7 @@ class VideoDataSet(data.Dataset):
                 self.feature_rgb_file = {}
                 self.feature_flow_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt["video_feature_all_train"], file + '.npz')
+                    feature_path = opt["video_feature_all_train"] + file + '.npz'
                     if not os.path.exists(feature_path):
                         raise ValueError(f"Feature file {feature_path} not found")
                     feature_All[file] = np.load(feature_path)
@@ -123,7 +123,7 @@ class VideoDataSet(data.Dataset):
                 self.feature_rgb_file = {}
                 self.feature_flow_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt["video_feature_all_train"], file + '.pt')
+                    feature_path = opt["video_feature_all_train"] + file + '.pt'
                     if not os.path.exists(feature_path):
                         raise ValueError(f"Feature file {feature_path} not found")
                     feature_All[file] = torch.load(feature_path)
@@ -164,7 +164,7 @@ class VideoDataSet(data.Dataset):
                 self.feature_rgb_file = {}
                 self.feature_flow_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt['video_feature_all_test'], file + '.npz')
+                    feature_path = os.path.join(opt['video_feature_all_test'], video_name + '.npz')
                     if not os.path.exists(feature_path):
                         raise ValueError(f"Feature file {feature_path} not found")
                     feature_All[file] = np.load(feature_path)['feats']
@@ -177,7 +177,7 @@ class VideoDataSet(data.Dataset):
                 self.feature_rgb_file = {}
                 self.feature_flow_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt['video_feature_all_test'], file + '.npz')
+                    feature_path = os.path.join(opt['video_feature_all_test'], video_name + '.npz')
                     if not os.path.exists(feature_path):
                         raise ValueError(f"Feature file {feature_path} not found")
                     feature_All[file] = np.load(feature_path)
@@ -190,7 +190,7 @@ class VideoDataSet(data.Dataset):
                 self.feature_rgb_file = {}
                 self.feature_flow_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt['video_feature_all_test'], file + '.pt')
+                    feature_path = opt["video_feature_all_test"] + file + '.pt'
                     if not os.path.exists(feature_path):
                         raise ValueError(f"Feature file {feature_path} not found")
                     feature_All[file] = torch.load(feature_path)
@@ -213,27 +213,15 @@ class VideoDataSet(data.Dataset):
             elif opt['data_format'] == "npz":
                 feature_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt["video_feature_all_train"], file + '.npz')
-                    if os.path.exists(feature_path):
-                        feature_file[file] = np.load(feature_path)['feats']
-                    else:
-                        print(f"Warning: Feature file {feature_path} not found for length calculation")
+                    feature_file[file] = np.load(opt["video_feature_all_train"] + file + '.npz')['feats']
             elif opt['data_format'] == "npz_i3d":
                 feature_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt["video_feature_all_train"], file + '.npz')
-                    if os.path.exists(feature_path):
-                        feature_file[file] = np.load(feature_path)
-                    else:
-                        print(f"Warning: Feature file {feature_path} not found for length calculation")
+                    feature_file[file] = np.load(opt["video_feature_all_train"] + file + '.npz')
             elif opt['data_format'] == "pt":
                 feature_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt["video_feature_all_train"], file + '.pt')
-                    if os.path.exists(feature_path):
-                        feature_file[file] = torch.load(feature_path)
-                    else:
-                        print(f"Warning: Feature file {feature_path} not found for length calculation")
+                    feature_file[file] = torch.load(opt["video_feature_all_train"] + file + '.pt')
         else:
             if opt['data_format'] == "h5":
                 feature_file = h5py.File(opt["video_feature_rgb_test"], 'r')
@@ -242,54 +230,35 @@ class VideoDataSet(data.Dataset):
             elif opt['data_format'] == "npz":
                 feature_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt['video_feature_all_test'], file + '.npz')
-                    if os.path.exists(feature_path):
-                        feature_file[file] = np.load(feature_path)['feats']
-                    else:
-                        print(f"Warning: Feature file {feature_path} not found for length calculation")
+                    feature_file[file] = np.load(opt["video_feature_all_test"] + file + '.npz')['feats']
             elif opt['data_format'] == "npz_i3d":
                 feature_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt['video_feature_all_test'], file + '.npz')
-                    if os.path.exists(feature_path):
-                        feature_file[file] = np.load(feature_path)
-                    else:
-                        print(f"Warning: Feature file {feature_path} not found for length calculation")
+                    feature_file[file] = np.load(opt["video_feature_all_test"] + file + '.npz')
             elif opt['data_format'] == "pt":
                 feature_file = {}
                 for file in self.video_list:
-                    feature_path = os.path.join(opt['video_feature_all_test'], file + '.pt')
-                    if os.path.exists(feature_path):
-                        feature_file[file] = torch.load(feature_path)
-                    else:
-                        print(f"Warning: Feature file {feature_path} not found for length calculation")
+                    feature_file[file] = torch.load(opt["video_feature_all_test"] + file + '.pt')
                     
         keys = self.video_list
         if opt['data_format'] == "h5":
             for vidx in range(len(keys)):
-                if keys[vidx] in feature_file:
-                    self.video_len[keys[vidx]] = len(feature_file[keys[vidx]])
+                self.video_len[keys[vidx]] = len(feature_file[keys[vidx]])
         elif opt['data_format'] == "pickle":
             for vidx in range(len(keys)):
-                if keys[vidx] in feature_file:
-                    self.video_len[keys[vidx]] = len(feature_file[keys[vidx]]['rgb'])
+                self.video_len[keys[vidx]] = len(feature_file[keys[vidx]]['rgb'])
         elif opt['data_format'] == "npz":
             for vidx in range(len(keys)):
-                if keys[vidx] in feature_file:
-                    self.video_len[keys[vidx]] = len(feature_file[keys[vidx]])
+                self.video_len[keys[vidx]] = len(feature_file[keys[vidx]])
         elif opt['data_format'] == "npz_i3d":
             for vidx in range(len(keys)):
-                if keys[vidx] in feature_file:
-                    self.video_len[keys[vidx]] = len(feature_file[keys[vidx]]['rgb'])
+                self.video_len[keys[vidx]] = len(feature_file[keys[vidx]]['rgb'])
         elif opt['data_format'] == "pt":
             for vidx in range(len(keys)):
-                if keys[vidx] in feature_file:
-                    self.video_len[keys[vidx]] = len(feature_file[keys[vidx]])
-        
-        if self.video_len:  # Only save if we have any lengths
-            outfile = open(self.video_len_path, "w")
-            json.dump(self.video_len, outfile, indent=2)
-            outfile.close()
+                self.video_len[keys[vidx]] = len(feature_file[keys[vidx]])
+        outfile = open(self.video_len_path, "w")
+        json.dump(self.video_len, outfile, indent=2)
+        outfile.close()
 
     def _getDatasetDict(self):
         anno_database = load_json(self.video_anno_path)
@@ -368,29 +337,29 @@ class VideoDataSet(data.Dataset):
             video_name = self.video_list[index]
             duration = self.match_score[video_name].shape[0]
             for i in range(1, duration + 1):
-                st = i - self._segment_size
+                st = i - self.segment_size
                 ed = i
-                self._inputs_all.append([video_name, st, ed, data_idx])
+                self.inputs_all.append([video_name, st, ed, data_idx])
                 data_idx += 1
                 
-        self._inputs = self._inputs_all.copy()
-        print(f"{self._subset} subset seg numbers: {len(self._inputs)}")
+        self.inputs = self.inputs_all.copy()
+        print(f"{self.subset} subset seg numbers: {len(self.inputs)}")
 
     def _makePropLabelUnit(self, i):
-        video_name = self._inputs_all[i][0]
-        st = self._inputs_all[i][1]
-        ed = self._inputs_all[i][2]
+        video_name = self.inputs_all[i][0]
+        st = self.inputs_all[i][1]
+        ed = self.inputs_all[i][2]
         cls_anc = []
         reg_anc = []
 
-        for j in range(0, len(self._anchors)):
-            v1 = np.zeros(self._num_of_class)
+        for j in range(0, len(self.anchors)):
+            v1 = np.zeros(self.num_of_class)
             v1[-1] = 1
             v2 = np.zeros(2)
             v2[-1] = -1e3
-            y_box = [ed - 1, self._anchors[j]]
+            y_box = [ed - 1, self.anchors[j]]
             
-            subset_label = self._get_train_subset_label(video_name, ed - self._anchors[j], ed)
+            subset_label = self._get_train_label_with_class(video_name, ed - self.anchors[j], ed)
             idx_list = []
             for ii in range(0, subset_label.shape[0]):
                 for jj in range(0, subset_label.shape[1]):
@@ -399,23 +368,23 @@ class VideoDataSet(data.Dataset):
                         idx_list.append(idx - 1)
             
             for idx in idx_list:
-                target_box_idx = self._gt_action_list[video_name][idx]
-                cls = int(target_box_idx[2])
-                iou = calc_iou(y_box_idx, target_box)
-                if iou >= self._pos_threshold or (j == len(self._anchors) - 1 and box_include_idx(y_box, target_box)) or (j == 0 and box_include_idx(target_box, y_box)):
+                target_box = self.gt_action[video_name][idx]
+                cls = int(target_box[2])
+                iou = calc_iou(y_box, target_box)
+                if iou >= self.pos_threshold or (j == len(self.anchors) - 1 and box_include(y_box, target_box)) or (j == 0 and box_include(target_box, y_box)):
                     v1[cls] = 1
                     v1[-1] = 0
-                    v2[0] = 1.0 * (target_box[0] - y_box[0]) / self._anchors[j]
+                    v2[0] = 1.0 * (target_box[0] - y_box[0]) / self.anchors[j]
                     v2[1] = np.log(1.0 * max(1, target_box[1]) / y_box[1])
             
             cls_anc.append(v1)
             reg_anc.append(v2)
 
-        v0 = np.zeros(self._num_of_class)
+        v0 = np.zeros(self.num_of_class)
         v0[-1] = 1
         segment_size = ed - st
-        y_box = [ed - 1, self._anchors[-1]]
-        subset_label = self._get_subset_label(video_name, ed - self._anchors[-1], ed)
+        y_box = [ed - 1, self.anchors[-1]]
+        subset_label = self._get_train_label_with_class(video_name, ed - self.anchors[-1], ed)
         idx_list = []
         for ii in range(0, subset_label.shape[0]):
             for jj in range(0, subset_label.shape[1]):
@@ -424,141 +393,141 @@ class VideoDataSet(data.Dataset):
                     idx_list.append(idx - 1)
         
         for idx in idx_list:
-            target_box = self._gt_action[video_name][idx]
+            target_box = self.gt_action[video_name][idx]
             cls = int(target_box[2])
             iou = calc_iou(y_box, target_box)
             if iou >= 0:
                 v0[cls] = 1
                 v0[-1] = 0
 
-        cls_anc = np.stack(cls._anc, idx=0)
-        reg_anc = np.stack(reg._anc, idx=0)
+        cls_anc = np.stack(cls_anc, axis=0)
+        reg_anc = np.stack(reg_anc, axis=0)
         cls_snip = np.array(v0)
         return cls_anc, reg_anc, cls_snip
 
     def _loadPropLabel(self, filename):
         if os.path.exists(filename):
             prop_label_file = h5py.File(filename, 'r')
-            self._cls_label = np.array(prop_label_file['cls_label'][:])
-            self._reg_label = np.array(prop_label_file['reg_label'][:])
-            self._snip_label = np.array(prop_label_file['snip_label'][:])
+            self.cls_label = np.array(prop_label_file['cls_label'][:])
+            self.reg_label = np.array(prop_label_file['reg_label'][:])
+            self.snip_label = np.array(prop_label_file['snip_label'][:])
             prop_label_file.close()
-            self._action_frame_count = np.sum(self._cls_label.reshape((-1, self._cls_label.shape[-1])), idx=0)
-            self._action_frame_count = torch.Tensor(self._action_frame_count)
+            self.action_frame_count = np.sum(self.cls_label.reshape((-1, self.cls_label.shape[-1])), axis=0)
+            self.action_frame_count = torch.Tensor(self.action_frame_count)
             return
     
         pool = Pool(os.cpu_count() // 2)
-        labels = pool.map(self._makePropLabelUnit, range(0, len(self._inputs_all)))
+        labels = pool.map(self._makePropLabelUnit, range(0, len(self.inputs_all)))
         pool.close()
-        pool pool.join()
+        pool.join()
         
         cls_label = []
         reg_label = []
         snip_label = []
         for i in range(0, len(labels)):
-            cls_label[i].append(labels[i][0])
+            cls_label.append(labels[i][0])
             reg_label.append(labels[i][1])
             snip_label.append(labels[i][2])
-        self._cls_label = np.stack(labels_cls, idx=0)
-        self._reg_label = np.stack(labels_reg, idx=0)
-        self._snip_label = np.stack(labels_snip, idx=0)
+        self.cls_label = np.stack(cls_label, axis=0)
+        self.reg_label = np.stack(reg_label, axis=0)
+        self.snip_label = np.stack(snip_label, axis=0)
         
         outfile = h5py.File(filename, 'w')
-        dset_cls = outfile._create_dataset('/cls_label', self._cls_label.shape, shape=self._cls._label_shape, chunks=True, type=np.float32)
-        dset_cls[_._ :] = self._cls._label[_._ :]
-        dset_reg_label = outfile._create_dataset('/label_reg', self._reg._label.shape, shape=self._reg._label.shape, chunks=True, type=np.float32)
-        dset_reg[_._ :] = self._reg._reg_label[_._ :]
-        dset_snip_label = outfile._create_dataset('/snip_label', self._snip._label.shape, shape=self._snip._label.shape, chunks=True, type=np.float32)
-        dset_snip[_._ :] = self._snip._snip_label[_._ :]
-        outfile._close()
+        dset_cls = outfile.create_dataset('/cls_label', self.cls_label.shape, maxshape=self.cls_label.shape, chunks=True, dtype=np.float32)
+        dset_cls[:, :] = self.cls_label[:, :]
+        dset_reg = outfile.create_dataset('/reg_label', self.reg_label.shape, maxshape=self.reg_label.shape, chunks=True, dtype=np.float32)
+        dset_reg[:, :] = self.reg_label[:, :]
+        dset_snip = outfile.create_dataset('/snip_label', self.snip_label.shape, maxshape=self.snip_label.shape, chunks=True, dtype=np.float32)
+        dset_snip[:, :] = self.snip_label[:, :]
+        outfile.close()
         
         return
 
-    def _getitem_item(self, idx):
-        video_name, st, ed, d_idx_data = self._inputs[idx]
+    def __getitem__(self, index):
+        video_name, st, ed, data_idx = self.inputs[index]
         if st >= 0:
-            feature_data = self._get_base_data(video_name, st, ed)
+            feature = self._get_base_data(video_name, st, ed)
         else:
-            feature_data = self._get_base_data(video_name, idx=0, st, ed)
-            pad_func = torch.nn.ConstantPad2d(st, (0, 0, -st, 0), idx=0)
-            data_feature = pad_func(data_feature)
+            feature = self._get_base_data(video_name, 0, ed)
+            padfunc2d = torch.nn.ConstantPad2d((0, 0, -st, 0), 0)
+            feature = padfunc2d(feature)
         
-        cls_label_data = torch.Tensor(self._cls_label[d_idx_data])
-        reg_label_data = torch.Tensor(self._reg_label[d_idx_data])
-        snip_label_data = torch.Tensor(self._snip_label[d_idx_data])
+        cls_label = torch.Tensor(self.cls_label[data_idx])
+        reg_label = torch.Tensor(self.reg_label[data_idx])
+        snip_label = torch.Tensor(self.snip_label[data_idx])
             
-        return data_feature, cls_label_data, reg_label_data, snip_label_data
+        return feature, cls_label, reg_label, snip_label
 
     def _get_base_data(self, video_name, st, ed):
-        feature_rgb_data = self._feature_rgb_file[video_name]
-        feature_rgb_data = feature_rgb_data[st:ed, :]
+        feature_rgb = self.feature_rgb_file[video_name]
+        feature_rgb = feature_rgb[st:ed, :]
         
-        if self._feature_flow_file is not None:
-            feature_flow_data = self._feature_flow_file[video_name]
-            feature_flow_data = feature_flow_data[st:ed, :]
-            data_feature = np.append(feature_data_rgb, feature_flow_data, idx=1)
+        if self.feature_flow_file is not None:
+            feature_flow = self.feature_flow_file[video_name]
+            feature_flow = feature_flow[st:ed, :]
+            feature = np.append(feature_rgb, feature_flow, axis=1)
         else:
-            data_feature = feature_rgb_data
-        data_feature = torch.from_numpy(np.array(data_feature))
+            feature = feature_rgb
+        feature = torch.from_numpy(np.array(feature))
      
-        return data_feature
+        return feature
 
-    def _get_train_label_with_class(self, video_name, st, idx_ed):
-        duration_data = len(self._match_score_data[video_name])
-        st_padding_data = pad_0
-        ed_padding_data = pad_0
+    def _get_train_label_with_class(self, video_name, st, ed):
+        duration = len(self.match_score[video_name])
+        st_padding = 0
+        ed_padding = 0
         if st < 0:
-            st_padding_data = -st
-            st = pad_0
-        if idx_ed > duration_data:
-            ed_padding_data = idx_ed - duration_data
-            idx_ed = duration_data
+            st_padding = -st
+            st = 0
+        if ed > duration:
+            ed_padding = ed - duration
+            ed = duration
     
-        match_score_data = torch.Tensor(self._match_score_data[video_name][st:idx_ed])
-        if st_padding_data > pad_0:
-            pad_func_2d = torch.nn.ConstantPad(data_2d, (pad_0, pad_0, st_padding_data, pad_0), idx=0)
-            data_match_score = pad_func_2d(data_match_score)
-        if ed_padding_data > pad_0:
-            pad_func_2d = torch.nn(data_ConstantPad2d, (pad_0, pad_0, pad_0, ed_padding_data), idx=pad_0)
-            pad_func_2d = pad(data_func_2d(data_match_score))
-        return data_match_score
+        match_score = torch.Tensor(self.match_score[video_name][st:ed])
+        if st_padding > 0:
+            padfunc2d = torch.nn.ConstantPad2d((0, 0, st_padding, 0), 0)
+            match_score = padfunc2d(match_score)
+        if ed_padding > 0:
+            padfunc2d = torch.nn.ConstantPad2d((0, 0, 0, ed_padding), 0)
+            match_score = padfunc2d(match_score)
+        return match_score
 
-    def _len__(self):
-        return len(self._inputs)
+    def __len__(self):
+        return len(self.inputs)
     
-    def _reset_sample(self):
-        self._inputs = self._inputs_all.copy()
+    def reset_sample(self):
+        self.inputs = self.inputs_all.copy()
         
-    def _select_sample(self, idx):
-        inputs_data = [self._inputs_all[i] for i in idx]
-        self._inputs = inputs_data.copy()
+    def select_sample(self, idx):
+        inputs = [self.inputs_all[i] for i in idx]
+        self.inputs = inputs.copy()
         return
 
 class SuppressDataSet(data.Dataset):
     def __init__(self, opt, subset="train"):
-        self._subset = subset
-        self._mode = opt["mode"]
-        self._data_file = h5py.File(opt["suppress_label_file"].format(self._subset + "_" + opt['setup']), 'r')
-        self._video_list = list(self._data_file.keys())
-        self._inputs = []
-        for idx in range(0, len(self._video_list)):
-            video_name = self._video_list[idx]
-            duration_data = self._data_file[video_name + '/input_seq'].shape[0]
-            for i in range(0, duration_data):
-                self._inputs.append([video_name, i])
+        self.subset = subset
+        self.mode = opt["mode"]
+        self.data_file = h5py.File(opt["suppress_label_file"].format(self.subset + "_" + opt['setup']), 'r')
+        self.video_list = list(self.data_file.keys())
+        self.inputs = []
+        for index in range(0, len(self.video_list)):
+            video_name = self.video_list[index]
+            duration = self.data_file[video_name + '/input'].shape[0]
+            for i in range(0, duration):
+                self.inputs.append([video_name, i])
                 
-        print(f"{self._subset} subset seg numbers: {len(self._inputs)}")
+        print(f"{self.subset} subset seg numbers: {len(self.inputs)}")
         
-    def _getitem__(self, idx):
-        video_name, idx = self._inputs[idx]
+    def __getitem__(self, index):
+        video_name, idx = self.inputs[index]
         
-        input_seq_data = self._data_file[video_name + '/input_seq'][idx]
-        label_data = self._data_file[video_name + '/label_data'][idx]
+        input_seq = self.data_file[video_name + '/input'][idx]
+        label = self.data_file[video_name + '/label'][idx]
         
-        input_seq_data = torch.from_numpy(input_seq_data)
-        label_data = torch.from_numpy(label_data)
+        input_seq = torch.from_numpy(input_seq)
+        label = torch.from_numpy(label)
         
-        return input_seq_data, label_data
+        return input_seq, label
             
-    def _len__(self):
-        return len(self._inputs)
+    def __len__(self):
+        return len(self.inputs)