Update main.py

main.py

@@ -45,7 +45,7 @@ VIS_CONFIG = {
     'frame_highlight_gt': 'red',
     'frame_highlight_pred': 'black',
     'iou_threshold': 0.3,
-    'frame_scale_factor': 0.3,  # Reduced for CPU
+    'frame_scale_factor': 0.3,
     'video_text_scale': 0.5,
     'video_gt_text_color': (180, 119, 31),  # BGR
     'video_pred_text_color': (14, 127, 255),  # BGR
@@ -55,8 +55,8 @@ VIS_CONFIG = {
     'video_pred_text_y': 0.45,
     'video_gt_text_y': 0.55,
     'video_footer_height': 150,
-    'video_gt_bar_y': 0.5,
-    'video_pred_bar_y': 0.8,
+    'video_gt_bar_y': 0.2,
+    'video_pred_bar_y': 0.5,
     'video_bar_height': 0.15,
     'video_bar_text_scale': 0.7,
     'min_segment_duration': 1.0,
@@ -108,7 +108,7 @@ def annotate_video_with_actions(
         (185, 218, 255), (255, 204, 204), (193, 182, 255), (201, 252, 189),
         (144, 128, 112), (112, 25, 25), (102, 51, 102), (0, 128, 128), (171, 71, 0)
     ]
-    action_labels = set(seg['label'] for seg in gt_segments).union(seg['label'] for seg in pred_segments)
+    action_labels = set(seg['label'] for seg in gt_segments).union(set(seg['label'] for seg in pred_segments))
     action_color_map = {label: color_palette[i % len(color_palette)] for i, label in enumerate(action_labels)}
     gt_color_rgb = (gt_text_color[2], gt_text_color[1], gt_text_color[0])
     pred_color_rgb = (pred_text_color[2], pred_text_color[1], pred_text_color[0])
@@ -128,10 +128,10 @@ def annotate_video_with_actions(
         except IOError:
             font = None
             bar_font = None
-    window_size = 20.0
+    window_size = VIS_CONFIG['scroll_window_duration']
     num_windows = int(np.ceil(duration / window_size))
     text_bar_gap = 48
-    text_x = 10
+    text_x = VIS_CONFIG['video_bar_label_x']
     frame_idx = 0
     written_frames = 0
     while cap.isOpened():
@@ -148,12 +148,12 @@ def annotate_video_with_actions(
         window_duration = window_end - window_start
         window_timestamp = timestamp - window_start
         gt_labels = [seg['label'] for seg in gt_segments if seg['start'] <= timestamp <= seg['end']]
-        gt_text = "GT: " + ", ".join(gt_labels) if gt_labels else ""
+        gt_text = "GT: " + ", ".join(gt_labels) if gt_labels else "GT: None"
         pred_labels = [seg['label'] for seg in pred_segments if seg['start'] <= timestamp <= seg['end']]
-        pred_text = "Pred: " + ", ".join(pred_labels) if pred_labels else ""
+        pred_text = "Pred: " + ", ".join(pred_labels) if pred_labels else "Pred: None"
         footer_y = frame_height
-        gt_bar_y = footer_y + int(0.2 * footer_height)
-        pred_bar_y = footer_y + int(0.5 * footer_height)
+        gt_bar_y = footer_y + int(VIS_CONFIG['video_gt_bar_y'] * footer_height)
+        pred_bar_y = footer_y + int(VIS_CONFIG['video_pred_bar_y'] * footer_height)
         bar_height = int(VIS_CONFIG['video_bar_height'] * footer_height)
         if font:
             gt_text_bbox = bar_font.getbbox("GT")
@@ -204,20 +204,18 @@ def annotate_video_with_actions(
             frame_text_bbox = draw.textbbox((0, 0), frame_info, font=font)
             frame_text_width = frame_text_bbox[2] - frame_text_bbox[0]
             frame_text_x = (frame_width - frame_text_width) // 2
-            draw.text((frame_text_x, 10), frame_info, font=font, fill=(0, 0, 0))
+            draw.text((frame_text_x, int(frame_height * VIS_CONFIG['video_frame_text_y'])), frame_info, font=font, fill=(0, 0, 0))
             window_info = f"{window_start:.1f}s - {window_end:.1f}s"
             window_text_bbox = draw.textbbox((0, 0), window_info, font=bar_font)
             window_text_width = window_text_bbox[2] - window_text_bbox[0]
             window_text_x = (frame_width - window_text_width) // 2
             draw.text((window_text_x, footer_y + 10), window_info, font=bar_font, fill=(0, 0, 0))
-            if gt_text:
-                gt_y = int(frame_height * VIS_CONFIG['video_gt_text_y'])
-                draw.text((10, gt_y), gt_text, font=font, fill=gt_color_rgb)
-            if pred_text:
-                pred_y = int(frame_height * VIS_CONFIG['video_pred_text_y'])
-                draw.text((10, pred_y), pred_text, font=font, fill=pred_color_rgb)
             draw.text((text_x, gt_bar_y + bar_height // 2), "GT", font=bar_font, fill=gt_color_rgb)
             draw.text((text_x, pred_bar_y + bar_height // 2), "Pred", font=bar_font, fill=pred_color_rgb)
+            gt_y = int(frame_height * VIS_CONFIG['video_gt_text_y'])
+            pred_y = int(frame_height * VIS_CONFIG['video_pred_text_y'])
+            draw.text((10, gt_y), gt_text, font=font, fill=gt_color_rgb)
+            draw.text((10, pred_y), pred_text, font=font, fill=pred_color_rgb)
             extended_frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
         else:
             frame_info = f"Frame: {frame_idx} | FPS: {fps:.2f}"
@@ -226,7 +224,7 @@ def annotate_video_with_actions(
             cv2.putText(
                 extended_frame,
                 frame_info,
-                (frame_text_x, 30),
+                (frame_text_x, int(frame_height * VIS_CONFIG['video_frame_text_y']) + 20),
                 cv2.FONT_HERSHEY_SIMPLEX,
                 text_scale,
                 (0, 0, 0),
@@ -246,28 +244,26 @@ def annotate_video_with_actions(
                 1,
                 cv2.LINE_AA
             )
-            if gt_text:
-                cv2.putText(
-                    extended_frame,
-                    gt_text,
-                    (10, int(frame_height * VIS_CONFIG['video_gt_text_y'])),
-                    cv2.FONT_HERSHEY_SIMPLEX,
-                    text_scale,
-                    gt_text_color,
-                    text_thickness,
-                    cv2.LINE_AA
-                )
-            if pred_text:
-                cv2.putText(
-                    extended_frame,
-                    pred_text,
-                    (10, int(frame_height * VIS_CONFIG['video_pred_text_y'])),
-                    cv2.FONT_HERSHEY_SIMPLEX,
-                    text_scale,
-                    pred_text_color,
-                    text_thickness,
-                    cv2.LINE_AA
-                )
+            cv2.putText(
+                extended_frame,
+                gt_text,
+                (10, int(frame_height * VIS_CONFIG['video_gt_text_y'])),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                text_scale,
+                gt_text_color,
+                text_thickness,
+                cv2.LINE_AA
+            )
+            cv2.putText(
+                extended_frame,
+                pred_text,
+                (10, int(frame_height * VIS_CONFIG['video_pred_text_y'])),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                text_scale,
+                pred_text_color,
+                text_thickness,
+                cv2.LINE_AA
+            )
             cv2.putText(
                 extended_frame,
                 "GT",
@@ -316,7 +312,7 @@ def visualize_action_lengths(
     if num_frames > VIS_CONFIG['max_frames']:
         frame_interval = duration / (VIS_CONFIG['max_frames'] - 1)
         num_frames = VIS_CONFIG['max_frames']
-    frame_times = np.linspace(0, duration, num_frames, endpoint=False)
+    frame_times = np.linspace(0, duration, num_frames, endpoint=True)
     frames = []
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
@@ -328,12 +324,12 @@ def visualize_action_lengths(
             ret, frame = cap.read()
             if ret:
                 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                frame = cv2.resize(frame, (int(frame.shape[1] * 0.3), int(frame.shape[0] * 0.3)))  # Smaller resize
+                frame = cv2.resize(frame, (int(frame.shape[1] * VIS_CONFIG['frame_scale_factor']), int(frame.shape[0] * VIS_CONFIG['frame_scale_factor'])))
                 frames.append(frame)
             else:
                 frames.append(np.ones((100, 100, 3), dtype=np.uint8) * 255)
         cap.release()
-    fig = plt.figure(figsize=(num_frames * VIS_CONFIG['frame_scale_factor'], 6), constrained_layout=True)
+    fig = plt.figure(figsize=(num_frames * 2, 6), constrained_layout=True)
     gs = fig.add_gridspec(3, num_frames, height_ratios=[3, 1, 1])
     for i, (t, frame) in enumerate(zip(frame_times, frames)):
         ax = fig.add_subplot(gs[0, i])
@@ -401,7 +397,7 @@ def train_one_epoch(opt, model, train_dataset, optimizer, warmup=False):
     train_loader = torch.utils.data.DataLoader(
         train_dataset,
         batch_size=opt['batch_size'], shuffle=True,
-        num_workers=1, pin_memory=True, drop_last=True
+        num_workers=0, pin_memory=False, drop_last=True
     )
     epoch_cost = 0
     epoch_cost_cls = 0
@@ -416,7 +412,7 @@ def train_one_epoch(opt, model, train_dataset, optimizer, warmup=False):
                 g['lr'] = n_iter * opt['lr'] / total_iter
         act_cls, act_reg, snip_cls = model(input_data.float().to(device))
         act_cls.register_hook(partial(cls_loss.collect_grad, cls_label))
-        snip_cls.register_hook(partial_labels=snip_label)
+        snip_cls.register_hook(lambda grad: snip_loss.collect_grad(grad, snip_label))
         cost_reg = 0
         cost_cls = 0
         loss = cls_loss_func_(cls_loss, act_cls)
@@ -435,60 +431,12 @@ def train_one_epoch(opt, model, train_dataset, optimizer, warmup=False):
         optimizer.step()
     return n_iter, epoch_cost, epoch_cost_cls, epoch_cost_reg, epoch_cost_snip
 
-def eval_one_epoch(opt, model, test_dataset):
-    device = torch.device("cpu")
-    cls_loss, reg_loss, tot_loss, output_cls, output_reg, labels_cls, labels_reg, _, _ = eval_frame(opt, model, test_dataset)
-    result_dict = eval_map_nms(opt, test_dataset, output_cls, output_reg, labels_cls, labels_reg)
-    output_dict = {"version": "VERSION 1.3, "results": result_dict, "external_data": None}
-    result_path = opt["result_file"].format(opt['exp"])
-    os.makedirs(os.path.dirname(result_path), exist_ok=True)
-    with open(result_path, "w") as outfile:
-        json.dump(output_dict, outfile, indent=2)
-    IoUmAP = evaluation_detection(opt, verbose=False)
-    IoUmAP_5 = sum(IoUmAP[0:]) / len(IoUmAP[0:]) if IoUmAP else 0
-    return cls_loss, reg_loss, tot_loss, IoUmAP_5
-
-def train(opt):
-    writer = SummaryWriter()
-    device = torch.device("cpu")
-    model = MYNET(opt).to(device)
-    rest_of_model_params = [param for name, param in model.named_parameters() if "history_unit" not in name]
-    optimizer = optim.Adam([
-        {'params': model.history_unit.parameters(), 'lr': 1e-6},
-        {'params': rest_of_model_params}
-    ], lr=opt["lr"], weight_decay=opt["weight_decay"])
-    scheduler = optim.StepLR(optimizer, step_size=opt["lr_step"])
-    train_dataset = VideoDataSet(opt, subset="training")
-    test_dataset = VideoDataSet(opt, subset=opt['inference_subset'])
-    warmup = False
-    for n_epoch in range(opt['epoch']):
-        if n_epoch >= 1:
-            warmup = False
-        n_iter, epoch_cost, epoch_cost_cls, epoch_cost_reg, epoch_cost_snip = train_one_epoch(opt, model, train_dataset, optimizer, warmup)
-        writer.add_scalar('data/cost', epoch_cost / (n_iter + 1), n_epoch)
-        print(f"Epoch {n_epoch}: Train Loss: {epoch_cost / (n_iter + 1):.4f}, cls: {epoch_cost_cls / (n_iter + 1):.4f}, reg: {epoch_cost_reg / (n_iter + 1):.4f}, snip: {epoch_cost_snip / (n_iter + 1):.4f}, lr: {optimizer.param_groups[-1]['lr']:.6f}")
-        scheduler.step()
-        model.eval()
-        cls_loss, reg_loss, tot_loss, IoUmAP_5 = eval_one_epoch(opt, model, test_dataset)
-        writer.add_scalar('data/mAP', IoUmAP_5, n_epoch)
-        print(f"Epoch {n_epoch}: Test Loss: {tot_loss:.4f}, cls: {cls_loss:.4f}, reg: {reg_loss:.4f}, mAP: {IoUmAP_5:.4f}")
-        state = {'epoch': n_epoch + 1, 'state_dict': model.state_dict()}
-        checkpoint_path = os.path.join(opt["checkpoint_path"], f"{opt['exp']}_checkpoint_{n_epoch + 1}.pth.tar")
-        os.makedirs(opt["checkpoint_path"], exist_ok=True)
-        torch.save(state, checkpoint_path)
-        if IoUmAP_5 > getattr(model, 'best_map', 0):
-            model.best_map = IoUmAP_5
-            torch.save(state, os.path.join(opt["checkpoint_path"], f"{opt['exp']}_ckp_best.pth.tar"))
-        model.train()
-    writer.close()
-    return model.best_map
-
-def eval_frame_data(opt, model, dataset):
+def eval_frame(opt, model, dataset):
     device = torch.device("cpu")
     test_loader = torch.utils.data.DataLoader(
         dataset,
-        batch_size=opt['batch_size'], shuffle=True,
-        num_workers=1, pin_memory=True, drop_last=True
+        batch_size=opt['batch_size'], shuffle=False,
+        num_workers=0, pin_memory=False, drop_last=False
     )
     labels_cls = {video_name: [] for video_name in dataset.video_names}
     labels_reg = {video_name: [] for video_name in dataset.video_names}
@@ -499,11 +447,12 @@ def eval_frame_data(opt, model, dataset):
     epoch_cost = 0
     epoch_cost_cls = 0
     epoch_cost_reg = 0
-    for n_iter, (input_data, cls_label, reg_label, _) in enumerate(tqdm(test_loader)):
+    cls_loss_fn = MultiCrossEntropyLoss(focal=True)
+    for n_iter, (input_data, cls_label, reg_label, snip_label) in enumerate(tqdm(test_loader)):
         act_cls, act_reg, _ = model(input_data.float().to(device))
         cost_reg = 0
         cost_cls = 0
-        loss = cls_loss_func(cls_label, act_cls)
+        loss = cls_loss_func_(cls_loss_fn, act_cls)
         cost_cls = loss
         epoch_cost_cls += loss.detach().cpu().numpy()
         loss = regress_loss_func(reg_label, act_reg)
@@ -515,10 +464,10 @@ def eval_frame_data(opt, model, dataset):
         total_frames += input_data.size(0)
         for idx in range(input_data.size(0)):
             video_name, st, ed, data_idx = dataset.inputs[n_iter * opt['batch_size'] + idx]
-            output_cls[video_name].append(act_cls[idx, :].detach().cpu().numpy())
-            output_reg[video_name].append(act_reg[idx, :].detach().cpu().numpy())
-            labels_cls[video_name].append(cls_label[idx, :].cpu().numpy())
-            labels_reg[video_name].append(reg_label[idx, :].cpu().numpy())
+            output_cls[video_name].append(act_cls[idx].detach().cpu().numpy())
+            output_reg[video_name].append(act_reg[idx].detach().cpu().numpy())
+            labels_cls[video_name].append(cls_label[idx].cpu().numpy())
+            labels_reg[video_name].append(reg_label[idx].cpu().numpy())
     end_time = time.time()
     working_time = end_time - start_time
     for video_name in dataset.video_names:
@@ -547,7 +496,7 @@ def eval_map_nms(opt, dataset, output_cls, output_reg, labels_cls, labels_reg):
             reg_anc = output_reg[video_name][idx]
             proposal_anc_dict = []
             for anc_idx in range(len(anchors)):
-                cls = np.argwhere(cls_anc[anc_idx][:-1] > opt['threshold']).reshape(-1)
+                cls = np.argwhere(cls_anc[anc_idx][:-1] > threshold).reshape(-1)
                 if len(cls) == 0:
                     continue
                 ed = idx + anchors[anc_idx] * reg_anc[anc_idx][0]
@@ -571,7 +520,7 @@ def eval_map_nms(opt, dataset, output_cls, output_reg, labels_cls, labels_reg):
 def eval_map_suppress(opt, dataset, output_cls, output_reg, labels_cls, labels_reg):
     device = torch.device("cpu")
     model = SuppressNet(opt).to(device)
-    checkpoint_path = os.path.join(opt["checkpoint_path"], "ckp_best_suppress.pth.tar")
+    checkpoint_path = os.path.join(opt["checkpoint_path"], f"ckp_best_suppress.pth.tar")
     if os.path.exists(checkpoint_path):
         checkpoint = torch.load(checkpoint_path, map_location=device)
         model.load_state_dict(checkpoint['state_dict'])
@@ -595,7 +544,7 @@ def eval_map_suppress(opt, dataset, output_cls, output_reg, labels_cls, labels_r
             reg_anc = output_reg[video_name][idx]
             proposal_anc_dict = []
             for anc_idx in range(len(anchors)):
-                cls = np.argwhere(cls_anc[anc_idx][:-1] > opt['threshold']).reshape(-1)
+                cls = np.argwhere(cls_anc[anc_idx][:-1] > threshold).reshape(-1)
                 if len(cls) == 0:
                     continue
                 ed = idx + anchors[anc_idx] * reg_anc[anc_idx][0]
@@ -605,7 +554,7 @@ def eval_map_suppress(opt, dataset, output_cls, output_reg, labels_cls, labels_r
                     label = cls[cidx]
                     tmp_dict = {
                         "segment": [float(st * frame_to_time / 100.0), float(ed * frame_to_time / 100.0)],
-                        "score": float(cls_anc][anc_idx][label]),
+                        "score": float(cls_anc[anc_idx][label]),
                         "label": dataset.label_name[label],
                         "gentime": float(idx * frame_to_time / 100.0)
                     }
@@ -622,60 +571,13 @@ def eval_map_suppress(opt, dataset, output_cls, output_reg, labels_cls, labels_r
             for cls in range(num_class - 1):
                 if suppress_conf[cls] > opt['sup_threshold']:
                     for proposal in proposal_anc_dict:
-                        if proposal['label"] == dataset.label_name[cls]:
+                        if proposal['label'] == dataset.label_name[cls]:
                             if check_overlap_proposal(proposal_dict, proposal, overlapThresh=opt['soft_nms']) is None:
                                 proposal_dict.append(proposal)
         result_dict[video_name] = proposal_dict
         proposal_dict = []
     return result_dict
 
-def test_frame(opt, video_name=None):
-    device = torch.device("cpu")
-    model = MYNET(opt).to(device)
-    checkpoint_path = os.path.join(opt["checkpoint_path"], "ckp_best.pth.tar")
-    if not os.path.exists(checkpoint_path):
-        print(f"[ERROR] Checkpoint {checkpoint_path} not found.")
-        return 0, 0, 0
-    checkpoint = torch.load(checkpoint_path, map_location=device)
-    model.load_state_dict(checkpoint['state_dict'])
-    model.eval()
-    dataset = VideoDataSet(opt, subset=opt['inference_subset'], video_name=video_name)
-    outfile = h5py.File(opt['output_file'].format(opt['exp']), 'w')
-    cls_loss, reg_loss, tot_loss, output_cls, output_reg, labels_cls, labels_reg, working_time, total_frames = eval_frame(opt, model, dataset)
-    print(f"testing loss: {tot_loss:.4f}, cls_loss: {cls_loss:.4f}, reg_loss: {reg_loss:.4f}")
-    for video_name in dataset.video_names:
-        o_cls = output_cls[video_name]
-        o_reg = output_reg[video_name]
-        l_cls = labels_cls[video_name]
-        l_reg = labels_reg[video_name]
-        dset_predcls = outfile.create_dataset(video_name + '/pred_cls', o_cls.shape, maxshape=o_cls.shape, chunks=True, dtype=np.float32)
-        dset_predcls[:, :] = o_cls[:, :]
-        dset_predreg = outfile.create_dataset(video_name + '/pred_reg', o_reg.shape, maxshape=o_reg.shape, chunks=True, dtype=np.float32)
-        dset_predreg[:, :] = o_reg[:, :]
-        dset_labelcls = outfile.create_dataset(video_name + '/label_cls', l_cls.shape, maxshape=l_cls.shape, chunks=True, dtype=np.float32)
-        dset_labelcls[:, :] = l_cls[:, :]
-        dset_labelreg = outfile.create_dataset(video_name + '/label_reg', l_reg.shape, maxshape=l_reg.shape, chunks=True, dtype=np.float32)
-        dset_labelreg[:, :] = l_reg[:, :]
-    outfile.close()
-    print(f"[INFO] Working time: {working_time:.2f}s, {total_frames / working_time:.1f}fps, {total_frames} frames")
-    return cls_loss, reg_loss, tot_loss
-
-def patch_attention(m):
-    forward_orig = m.forward
-    def wrap(*args, **kwargs):
-        kwargs["need_weights"] = True
-        kwargs["average_attn_weights"] = False
-        return forward_orig(*args, **kwargs)
-    m.forward = wrap
-
-class SaveOutput:
-    def __init__(self):
-        self.outputs = []
-    def __call__(self, module, module_in, module_out):
-        self.outputs.append(module_out[1])
-    def clear(self):
-        self.outputs = []
-
 def test(opt, video_name=None):
     device = torch.device("cpu")
     model = MYNET(opt).to(device)
@@ -701,12 +603,13 @@ def test(opt, video_name=None):
     with open(result_path, 'w') as f:
         json.dump(output_dict, f, indent=4)
     mAP = evaluation_detection(opt, verbose=False)
+    mAP_value = sum(mAP) / len(mAP) if mAP else 0
     if video_name:
         print(f"\n[INFO] Comparing Predicted and Ground Truth Actions for Video: {video_name}")
-        anno_path = opt["annotations_path"]
+        anno_path = opt["video_anno"].format(opt["split"])
         if not os.path.exists(anno_path):
             print(f"[ERROR] Annotation file {anno_path} not found. Skipping comparison.")
-            return mAP, "", "", ""
+            return mAP_value, "", "", ""
         with open(anno_path, 'r') as f:
             anno_data = json.load(f)
         gt_annotations = anno_data['database'][video_name]['annotations']
@@ -723,7 +626,7 @@ def test(opt, video_name=None):
             'end': pred['segment'][1],
             'duration': pred['segment'][1] - pred['segment'][0],
             'score': pred['score']
-        } for pred in result_dict[video_name]]
+        } for pred in result_dict.get(video_name, [])]
         matches = []
         iou_threshold = VIS_CONFIG['iou_threshold']
         used_gt_indices = set()
@@ -733,7 +636,7 @@ def test(opt, video_name=None):
             for gt_idx, gt in enumerate(gt_segments):
                 if gt_idx in used_gt_indices:
                     continue
-                iou = calc_iou([pred['start'], pred['end']], [gt['start'], gt['end']])
+                iou = calc_iou([pred['end'], pred['duration']], [gt['end'], gt['duration']])
                 if iou > best_iou and iou >= iou_threshold:
                     best_iou = iou
                     best_gt_idx = gt_idx
@@ -778,7 +681,7 @@ def test(opt, video_name=None):
         comparison_text += f"- Total Predictions: {len(pred_segments)}\n"
         comparison_text += f"- Total Ground Truths: {len(gt_segments)}\n"
         comparison_text += f"- Matched Segments: {matched_count}\n"
-        comparison_text += f"- Average Duration Difference: (s): {avg_duration_diff:.2f}s\n"
+        comparison_text += f"- Average Duration Difference (s): {avg_duration_diff:.2f}\n"
         comparison_text += f"- Average IoU (Matched): {avg_iou:.2f}\n"
         video_path = opt.get('video_path', '')
         viz_path = ""
@@ -791,7 +694,7 @@ def test(opt, video_name=None):
                 video_path=video_path,
                 duration=duration
             )
-            video_out_path = annotate_action_with_video(
+            video_out_path = annotate_video_with_actions(
                 video_id=video_name,
                 pred_segments=pred_segments,
                 gt_segments=gt_segments,
@@ -799,12 +702,12 @@ def test(opt, video_name=None):
             )
         else:
             print(f"[WARNING] Video {video_path} not found. Skipping visualization.")
-        return mAP, comparison_text, viz_path, video_out_path
+        return mAP_value, comparison_text, viz_path, video_out_path
 
 def test_online(opt, video_name=None):
     device = torch.device("cpu")
     model = MYNET(opt).to(device)
-    checkpoint_path = os.path.join(opt["checkpoint_path"], "ckp_best.pth.tar")
+    checkpoint_path = os.path.join(opt["checkpoint_path"], f"{opt['exp']}_ckp_best.pth.tar")
     if not os.path.exists(checkpoint_path):
         print(f"[ERROR] Checkpoint {checkpoint_path} not found.")
         return 0
@@ -812,7 +715,7 @@ def test_online(opt, video_name=None):
     model.load_state_dict(checkpoint['state_dict'])
     model.eval()
     sup_model = SuppressNet(opt).to(device)
-    sup_checkpoint_path = os.path.join(opt["checkpoint_path"], "ckp_best_suppress.pth.tar")
+    sup_checkpoint_path = os.path.join(opt["checkpoint_path"], f"ckp_best_suppress.pth.tar")
     if os.path.exists(sup_checkpoint_path):
         checkpoint = torch.load(sup_checkpoint_path, map_location=device)
         sup_model.load_state_dict(checkpoint['state_dict'])
@@ -824,8 +727,8 @@ def test_online(opt, video_name=None):
         dataset,
         batch_size=1,
         shuffle=False,
-        num_workers=1,
-        pin_memory=True
+        num_workers=0,
+        pin_memory=False
     )
     result_dict = {}
     proposal_dict = []
@@ -844,7 +747,7 @@ def test_online(opt, video_name=None):
         for idx in range(duration):
             total_frames += 1
             input_queue[:-1, :] = input_queue[1:, :].clone()
-            input_queue[-1, :] = dataset._get_base_data(video_name, idx)
+            input_queue[-1, :] = dataset._get_base_data(video_name, idx, idx + 1).squeeze(0)
             minput = input_queue.unsqueeze(0).to(device)
             act_cls, act_reg, _ = model(minput)
             act_cls = torch.softmax(act_cls, dim=-1)
@@ -852,7 +755,7 @@ def test_online(opt, video_name=None):
             reg_anc = act_reg.squeeze(0).detach().cpu().numpy()
             proposal_anc_dict = []
             for anc_idx in range(len(anchors)):
-                cls = np.argwhere(cls_anc[anc_idx][:-1] > opt['threshold']).reshape(-1)
+                cls = np.argwhere(cls_anc[anc_idx][:-1] > threshold).reshape(-1)
                 if len(cls) == 0:
                     continue
                 ed = idx + anchors[anc_idx] * reg_anc[anc_idx][0]
@@ -893,7 +796,8 @@ def test_online(opt, video_name=None):
     with open(result_path, "w") as f:
         json.dump(output_dict, f, indent=4)
     mAP = evaluation_detection(opt, verbose=False)
-    return mAP
+    mAP_value = sum(mAP) / len(mAP) if mAP else 0
+    return mAP_value
 
 def main(opt, video_name=None):
     max_perf = 0
@@ -904,8 +808,6 @@ def main(opt, video_name=None):
     elif opt['mode'] == 'test':
         max_perf, comparison_text, viz_path, video_out_path = test(opt, video_name=video_name)
         return max_perf, comparison_text, viz_path, video_out_path
-    elif opt['mode'] == 'test_frame':
-        max_perf = test_frame(opt, video_name=video_name)
     elif opt['mode'] == 'test_online':
         max_perf = test_online(opt, video_name=video_name)
     elif opt['mode'] == 'eval':
@@ -924,13 +826,23 @@ def gradio_interface(video):
     opt_dict['mode'] = 'test'
     opt_dict['video_name'] = video_name
     opt_dict['video_path'] = video
-    opt_dict['feature_dir'] = os.path.join(os.getcwd(), 'data', 'features')
+    opt_dict['video_anno'] = os.path.join(os.getcwd(), 'data', 'annotations.json')
+    opt_dict['video_feature_all_test'] = os.path.join(os.getcwd(), 'data', 'features') + os.sep
     opt_dict['checkpoint_path'] = os.path.join(os.getcwd(), 'checkpoint')
-    opt_dict['result_file'] = os.path.join(os.getcwd(), 'results', f"result_{{}}.json")
-    opt_dict['annotations_path'] = os.path.join(os.getcwd(), 'data', 'annotations.json')
-    opt_dict['frame_result_file'] = os.path.join(os.getcwd(), 'results', f"frame_result_{{}}.h5")
+    opt_dict['result_file'] = os.path.join(os.getcwd(), 'results', 'result_{}.json')
+    opt_dict['frame_result_file'] = os.path.join(os.getcwd(), 'results', 'frame_result_{}.h5')
+    opt_dict['video_len_file'] = os.path.join(os.getcwd(), 'data', 'video_len_{}.json')
+    opt_dict['proposal_label_file'] = os.path.join(os.getcwd(), 'data', 'proposal_label_{}.h5')
+    opt_dict['suppress_label_file'] = os.path.join(os.getcwd(), 'data', 'suppress_label_{}.h5')
     opt_dict['batch_size'] = 1
+    opt_dict['data_format'] = 'npz_i3d'
+    opt_dict['rgb_only'] = False
     opt_dict['anchors'] = [int(item) for item in opt_dict['anchors'].split(',')]
+    opt_dict['predefined_fps'] = 30  # Adjust if needed
+    opt_dict['split'] = 'test'
+    opt_dict['setup'] = 'default'
+    opt_dict['data_rescale'] = 1.0
+    opt_dict['pos_threshold'] = 0.5
     mAP, comparison_text, viz_path, video_out_path = main(opt_dict, video_name=video_name)
     return viz_path, video_out_path, f"mAP: {mAP:.4f}\n\n{comparison_text}"
 
@@ -938,12 +850,23 @@ if __name__ == "__main__":
     opt = opts.parse_opt()
     opt = vars(opt)
     opt['checkpoint_path'] = os.path.join(os.getcwd(), 'checkpoint')
-    opt['result_file'] = os.path.join(os.getcwd(), 'results', f"result_{{}}.json")
-    opt['frame_result_file'] = os.path.join(os.getcwd(), 'results', f"frame_result_{{}}.h5")
-    opt['feature_dir'] = os.path.join(os.getcwd(), 'data', 'features')
+    opt['result_file'] = os.path.join(os.getcwd(), 'results', 'result_{}.json')
+    opt['frame_result_file'] = os.path.join(os.getcwd(), 'results', 'frame_result_{}.h5')
+    opt['video_anno'] = os.path.join(os.getcwd(), 'data', 'annotations.json')
+    opt['video_feature_all_test'] = os.path.join(os.getcwd(), 'data', 'features') + os.sep
+    opt['video_len_file'] = os.path.join(os.getcwd(), 'data', 'video_len_{}.json')
+    opt['proposal_label_file'] = os.path.join(os.getcwd(), 'data', 'proposal_label_{}.h5')
+    opt['suppress_label_file'] = os.path.join(os.getcwd(), 'data', 'suppress_label_{}.h5')
+    opt['data_format'] = 'npz_i3d'
+    opt['rgb_only'] = False
+    opt['predefined_fps'] = 30
+    opt['split'] = 'test'
+    opt['setup'] = 'default'
+    opt['data_rescale'] = 1.0
+    opt['pos_threshold'] = 0.5
     os.makedirs(opt["checkpoint_path"], exist_ok=True)
     os.makedirs(os.path.dirname(opt["result_file"].format(opt['exp'])), exist_ok=True)
-    os.makedirs(opt["feature_dir"], exist_ok=True)
+    os.makedirs(os.path.dirname(opt["video_anno"]), exist_ok=True)
     with open(os.path.join(opt["checkpoint_path"], f"{opt['exp']}_opts.json"), "w") as f:
         json.dump(opt, f, indent=4)
     if opt['seed'] >= 0:
@@ -961,7 +884,7 @@ if __name__ == "__main__":
                 gr.Textbox(label="Results and mAP")
             ],
             title="Action Detection Model",
-            description="Upload a video to detect actions using pre-extracted I3D features. View visualizations and performance metrics."
+            description="Upload a video to detect actions using pre-extracted I3D features. Ensure a corresponding .npz file exists in data/features/. View visualizations and performance metrics."
         )
         iface.launch()
     else: