Update main.py

main.py

@@ -2,8 +2,6 @@ import os
 import json
 import torch
 import numpy as np
-import cv2
-from tqdm import tqdm
 import gradio as gr
 import opts_egtea as opts
 from dataset import VideoDataSet, calc_iou
@@ -11,626 +9,350 @@ from models import MYNET, SuppressNet
 from loss_func import cls_loss_func, regress_loss_func
 from eval import evaluation_detection
 from iou_utils import non_max_suppression, check_overlap_proposal
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
-from PIL import Image, ImageDraw, ImageFont
 from typing import List, Dict, Optional
 
-# Visualization Configuration
+# Configuration
 VIS_CONFIG = {
-    'frame_interval': 1.0,
-    'max_frames': 20,
-    'save_dir': './output/visualizations',
-    'video_save_dir': './output/videos',
-    'gt_color': '#1f77b4',  # Blue for ground truth
-    'pred_color': '#ff7f0e',  # Orange for predictions
-    'fontsize_label': 10,
-    'fontsize_title': 14,
-    'frame_highlight_both': 'green',
-    'frame_highlight_gt': 'red',
-    'frame_highlight_pred': 'black',
     'iou_threshold': 0.3,
-    'frame_scale_factor': 0.8,
-    'video_text_scale': 0.5,
-    'video_gt_text_color': (180, 119, 31),  # BGR for OpenCV
-    'video_pred_text_color': (14, 127, 255),  # BGR for OpenCV
-    'video_text_thickness': 1,
-    'video_font_path': './data/Poppins ExtraBold Italic 800.ttf',
-    'video_font_fallback': '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',
-    '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_bar_height': 0.15,
-    'video_bar_text_scale': 0.7,
     'min_segment_duration': 1.0,
-    'video_frame_text_y': 0.05,
-    'video_bar_label_x': 10,
-    'video_bar_label_scale': 0.5,
-    'scroll_window_duration': 20.0,
-    'scroll_speed': 0.2,
 }
 
 # Determine device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-def annotate_video_with_actions(
-    video_id: str,
-    pred_segments: List[Dict],
-    gt_segments: List[Dict],
-    video_path: str,
-    save_dir: str = VIS_CONFIG['video_save_dir'],
-    text_scale: float = VIS_CONFIG['video_text_scale'] * 1.2,
-    gt_text_color: tuple = VIS_CONFIG['video_gt_text_color'],
-    pred_text_color: tuple = VIS_CONFIG['video_pred_text_color'],
-    text_thickness: int = VIS_CONFIG['video_text_thickness']
-) -> str:
-    os.makedirs(save_dir, exist_ok=True)
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        return f"Error: Could not open video {video_path}."
-    
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    duration = total_frames / fps
-    
-    footer_height = VIS_CONFIG['video_footer_height']
-    output_height = frame_height + footer_height
-    output_path = os.path.join(save_dir, f"annotated_{video_id}.avi")
-    fourcc = cv2.VideoWriter_fourcc(*'XVID')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, output_height))
-    
-    if not out.isOpened():
-        cap.release()
-        return f"Error: Could not initialize video writer for {output_path}."
-    
-    min_duration = VIS_CONFIG['min_segment_duration']
-    gt_segments = [seg for seg in gt_segments if seg['duration'] >= min_duration]
-    pred_segments = [seg for seg in pred_segments if seg['duration'] >= min_duration]
-    
-    color_palette = [
-        (128, 0, 0), (60, 20, 220), (0, 128, 0), (128, 0, 128), (79, 69, 54),
-        (128, 128, 0), (0, 0, 128), (130, 0, 75), (34, 139, 34), (0, 85, 204),
-        (149, 146, 209), (235, 206, 135), (250, 230, 230), (191, 226, 159),
-        (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_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])
-    
-    font_path = VIS_CONFIG['video_font_path']
-    font_fallback = VIS_CONFIG['video_font_fallback']
-    font_size = int(20 * text_scale)
-    bar_font_size = int(20 * VIS_CONFIG['video_bar_text_scale'])
-    font = None
-    bar_font = None
-    try:
-        font = ImageFont.truetype(font_path, font_size)
-        bar_font = ImageFont.truetype(font_path, bar_font_size)
-    except IOError:
-        try:
-            font = ImageFont.truetype(font_fallback, font_size)
-            bar_font = ImageFont.truetype(font_fallback, bar_font_size)
-        except IOError:
-            font = None
-            bar_font = None
-    
-    window_size = 20.0
-    num_windows = int(np.ceil(duration / window_size))
-    text_bar_gap = 48
-    text_x = 10
-    
-    frame_idx = 0
-    written_frames = 0
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        extended_frame = np.zeros((output_height, frame_width, 3), dtype=np.uint8)
-        extended_frame[:frame_height, :, :] = frame
-        extended_frame[frame_height:, :, :] = 255
-        
-        timestamp = frame_idx / fps
-        window_idx = int(timestamp // window_size)
-        window_start = window_idx * window_size
-        window_end = min(window_start + window_size, duration)
-        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 ""
-        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 ""
-        
-        footer_y = frame_height
-        gt_bar_y = footer_y + int(0.2 * footer_height)
-        pred_bar_y = footer_y + int(0.5 * footer_height)
-        bar_height = int(VIS_CONFIG['video_bar_height'] * footer_height)
-        
-        if font:
-            gt_text_bbox = bar_font.getbbox("GT")
-            pred_text_bbox = bar_font.getbbox("Pred")
-            gt_text_width = gt_text_bbox[2] - gt_text_bbox[0]
-            pred_text_width = pred_text_bbox[2] - pred_text_bbox[0]
-        else:
-            gt_text_size, _ = cv2.getTextSize("GT", cv2.FONT_HERSHEY_DUPLEX, VIS_CONFIG['video_bar_text_scale'], 1)
-            pred_text_size, _ = cv2.getTextSize("Pred", cv2.FONT_HERSHEY_DUPLEX, VIS_CONFIG['video_bar_text_scale'], 1)
-            gt_text_width = gt_text_size[0]
-            pred_text_width = pred_text_size[0]
-        max_text_width = max(gt_text_width, pred_text_width)
-        bar_start_x = text_x + max_text_width + text_bar_gap
-        bar_width = frame_width - bar_start_x
-        
-        for seg in gt_segments:
-            if seg['start'] <= window_end and seg['end'] >= window_start:
-                start_t = max(seg['start'], window_start)
-                end_t = min(seg['end'], window_start + window_timestamp)
-                start_x = bar_start_x + int(((start_t - window_start) / window_duration) * bar_width)
-                end_x = bar_start_x + int(((end_t - window_start) / window_duration) * bar_width)
-                if end_x > start_x:
-                    cv2.rectangle(
-                        extended_frame,
-                        (start_x, gt_bar_y),
-                        (end_x, gt_bar_y + bar_height),
-                        action_color_map[seg['label']],
-                        -1
-                    )
-        
-        for seg in pred_segments:
-            if seg['start'] <= window_end and seg['end'] >= window_start:
-                start_t = max(seg['start'], window_start)
-                end_t = min(seg['end'], window_start + window_timestamp)
-                start_x = bar_start_x + int(((start_t - window_start) / window_duration) * bar_width)
-                end_x = bar_start_x + int(((end_t - window_start) / window_duration) * bar_width)
-                if end_x > start_x:
-                    cv2.rectangle(
-                        extended_frame,
-                        (start_x, pred_bar_y),
-                        (end_x, pred_bar_y + bar_height),
-                        action_color_map[seg['label']],
-                        -1
-                    )
-        
-        if font:
-            frame_rgb = cv2.cvtColor(extended_frame, cv2.COLOR_BGR2RGB)
-            pil_image = Image.fromarray(frame_rgb)
-            draw = ImageDraw.Draw(pil_image)
-            
-            frame_info = f"Frame: {frame_idx} | FPS: {fps:.2f}"
-            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))
-            
-            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)
-            
-            extended_frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
-        else:
-            frame_info = f"Frame: {frame_idx} | FPS: {fps:.2f}"
-            text_size, _ = cv2.getTextSize(frame_info, cv2.FONT_HERSHEY_DUPLEX, text_scale, text_thickness)
-            frame_text_x = (frame_width - text_size[0]) // 2
-            cv2.putText(
-                extended_frame,
-                frame_info,
-                (frame_text_x, 30),
-                cv2.FONT_HERSHEY_DUPLEX,
-                text_scale,
-                (0, 0, 0),
-                text_thickness,
-                cv2.LINE_AA
-            )
-            window_info = f"{window_start:.1f}s - {window_end:.1f}s"
-            window_text_size, _ = cv2.getTextSize(window_info, cv2.FONT_HERSHEY_DUPLEX, VIS_CONFIG['video_bar_text_scale'], 1)
-            window_text_x = (frame_width - window_text_size[0]) // 2
-            cv2.putText(
-                extended_frame,
-                window_info,
-                (window_text_x, footer_y + 20),
-                cv2.FONT_HERSHEY_DUPLEX,
-                VIS_CONFIG['video_bar_text_scale'],
-                (0, 0, 0),
-                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_DUPLEX,
-                    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_DUPLEX,
-                    text_scale,
-                    pred_text_color,
-                    text_thickness,
-                    cv2.LINE_AA
-                )
-            cv2.putText(
-                extended_frame,
-                "GT",
-                (text_x, gt_bar_y + bar_height // 2 + 5),
-                cv2.FONT_HERSHEY_DUPLEX,
-                VIS_CONFIG['video_bar_text_scale'],
-                gt_text_color,
-                1,
-                cv2.LINE_AA
-            )
-            cv2.putText(
-                extended_frame,
-                "Pred",
-                (text_x, pred_bar_y + bar_height // 2 + 5),
-                cv2.FONT_HERSHEY_DUPLEX,
-                VIS_CONFIG['video_bar_text_scale'],
-                pred_text_color,
-                1,
-                cv2.LINE_AA
-            )
-        
-        out.write(extended_frame)
-        written_frames += 1
-        frame_idx += 1
-    
-    cap.release()
-    out.release()
-    mp4_path = os.path.splitext(output_path)[0] + '.mp4'
-    os.system(f"ffmpeg -i {output_path} -vcodec libx264 -acodec aac {mp4_path} -y")
-    return mp4_path if os.path.exists(mp4_path) else output_path
-
-def visualize_action_lengths(
-    video_id: str,
-    pred_segments: List[Dict],
-    gt_segments: List[Dict],
-    video_path: str,
-    duration: float,
-    save_dir: str = VIS_CONFIG['save_dir'],
-    frame_interval: float = VIS_CONFIG['frame_interval']
-) -> str:
-    os.makedirs(save_dir, exist_ok=True)
-    num_frames = int(duration / frame_interval) + 1
-    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)
-    frames = []
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        frames = [np.ones((100, 100, 3), dtype=np.uint8) * 255 for _ in frame_times]
-    else:
-        for t in frame_times:
-            cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
-            ret, frame = cap.read()
-            if ret:
-                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                frame = cv2.resize(frame, (int(frame.shape[1] * 0.5), int(frame.shape[0] * 0.5)))
-                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)
-    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])
-        gt_hit = any(seg['start'] <= t <= seg['end'] for seg in gt_segments)
-        pred_hit = any(seg['start'] <= t <= seg['end'] for seg in pred_segments)
-        border_color = None
-        if gt_hit and pred_hit:
-            border_color = VIS_CONFIG['frame_highlight_both']
-        elif gt_hit:
-            border_color = VIS_CONFIG['frame_highlight_gt']
-        elif pred_hit:
-            border_color = VIS_CONFIG['frame_highlight_pred']
-        
-        ax.imshow(frame)
-        ax.axis('off')
-        if border_color:
-            for spine in ax.spines.values():
-                spine.set_edgecolor(border_color)
-                spine.set_linewidth(2)
-        ax.set_title(f"{t:.1f}s", fontsize=VIS_CONFIG['fontsize_label'], color=border_color or 'black')
-    
-    ax_gt = fig.add_subplot(gs[1, :])
-    ax_gt.set_xlim(0, duration)
-    ax_gt.set_ylim(0, 1)
-    ax_gt.axis('off')
-    ax_gt.text(-0.02 * duration, 0.5, "Ground Truth", fontsize=VIS_CONFIG['fontsize_title'], va='center', ha='right', weight='bold')
-    
-    for seg in gt_segments:
-        start, end = seg['start'], seg['end']
-        width = end - start
-        label = seg['label'][:10] + '...' if len(seg['label']) > 10 else seg['label']
-        ax_gt.add_patch(patches.Rectangle(
-            (start, 0.3), width, 0.4, facecolor=VIS_CONFIG['gt_color'], edgecolor='black', alpha=0.8
-        ))
-        ax_gt.text((start + end) / 2, 0.5, label, ha='center', va='center', fontsize=VIS_CONFIG['fontsize_label'], color='white')
-        ax_gt.text(start, 0.2, f"{start:.1f}", ha='center', fontsize=8, color='black')
-        ax_gt.text(end, 0.2, f"{end:.1f}", ha='center', fontsize=8, color='black')
-    
-    ax_pred = fig.add_subplot(gs[2, :])
-    ax_pred.set_xlim(0, duration)
-    ax_pred.set_ylim(0, 1)
-    ax_pred.axis('off')
-    ax_pred.text(-0.02 * duration, 0.5, "Prediction", fontsize=VIS_CONFIG['fontsize_title'], va='center', ha='right', weight='bold')
-    
-    for seg in pred_segments:
-        start, end = seg['start'], seg['end']
-        width = end - start
-        label = seg['label'][:10] + '...' if len(seg['label']) > 10 else seg['label']
-        ax_pred.add_patch(patches.Rectangle(
-            (start, 0.3), width, 0.4, facecolor=VIS_CONFIG['pred_color'], edgecolor='black', alpha=0.8
-        ))
-        ax_pred.text((start + end) / 2, 0.5, label, ha='center', va='center', fontsize=VIS_CONFIG['fontsize_label'], color='white')
-        ax_pred.text(start, 0.8, f"{start:.1f}", ha='center', fontsize=8, color='black')
-        ax_pred.text(end, 0.8, f"{end:.1f}", ha='center', fontsize=8, color='black')
-    
-    jpg_path = os.path.join(save_dir, f"viz_{video_id}.png")
-    plt.savefig(jpg_path, dpi=100, bbox_inches='tight')
-    plt.close()
-    return jpg_path
-
 def eval_frame(opt, model, dataset):
-    test_loader = torch.utils.data.DataLoader(dataset, batch_size=opt['batch_size'], shuffle=False, num_workers=0, pin_memory=False)
+    """Evaluate model frame by frame"""
+    test_loader = torch.utils.data.DataLoader(
+        dataset, 
+        batch_size=opt['batch_size'], 
+        shuffle=False, 
+        num_workers=0, 
+        pin_memory=False
+    )
+    
     labels_cls = {video_name: [] for video_name in dataset.video_list}
     labels_reg = {video_name: [] for video_name in dataset.video_list}
     output_cls = {video_name: [] for video_name in dataset.video_list}
     output_reg = {video_name: [] for video_name in dataset.video_list}
     
-    total_frames = 0
-    for n_iter, (input_data, cls_label, reg_label, _) in enumerate(tqdm(test_loader)):
-        input_data = input_data.to(device)
-        cls_label = cls_label.to(device)
-        reg_label = reg_label.to(device)
-        act_cls, act_reg, _ = model(input_data.float())
-        
-        act_cls = torch.softmax(act_cls, dim=-1)
-        total_frames += input_data.size(0)
-        
-        for b in range(input_data.size(0)):
-            video_name, _, _, _ = dataset.inputs[n_iter * opt['batch_size'] + b]
-            output_cls[video_name].append(act_cls[b, :].detach().cpu().numpy())
-            output_reg[video_name].append(act_reg[b, :].detach().cpu().numpy())
-            labels_cls[video_name].append(cls_label[b, :].cpu().numpy())
-            labels_reg[video_name].append(reg_label[b, :].cpu().numpy())
-    
+    model.eval()
+    with torch.no_grad():
+        for n_iter, batch_data in enumerate(test_loader):
+            try:
+                if len(batch_data) == 4:
+                    input_data, cls_label, reg_label, _ = batch_data
+                else:
+                    input_data, cls_label, reg_label = batch_data
+                
+                input_data = input_data.to(device)
+                cls_label = cls_label.to(device) if cls_label is not None else None
+                reg_label = reg_label.to(device) if reg_label is not None else None
+                
+                act_cls, act_reg, _ = model(input_data.float())
+                act_cls = torch.softmax(act_cls, dim=-1)
+                
+                for b in range(input_data.size(0)):
+                    batch_idx = n_iter * opt['batch_size'] + b
+                    if batch_idx < len(dataset.inputs):
+                        video_name = dataset.inputs[batch_idx][0]
+                        output_cls[video_name].append(act_cls[b, :].detach().cpu().numpy())
+                        output_reg[video_name].append(act_reg[b, :].detach().cpu().numpy())
+                        
+                        if cls_label is not None:
+                            labels_cls[video_name].append(cls_label[b, :].cpu().numpy())
+                        if reg_label is not None:
+                            labels_reg[video_name].append(reg_label[b, :].cpu().numpy())
+                            
+            except Exception as e:
+                print(f"Error in batch {n_iter}: {str(e)}")
+                continue
+    
+    # Stack arrays
     for video_name in dataset.video_list:
-        labels_cls[video_name] = np.stack(labels_cls[video_name], axis=0)
-        labels_reg[video_name] = np.stack(labels_reg[video_name], axis=0)
-        output_cls[video_name] = np.stack(output_cls[video_name], axis=0)
-        output_reg[video_name] = np.stack(output_reg[video_name], axis=0)
+        if output_cls[video_name]:
+            output_cls[video_name] = np.stack(output_cls[video_name], axis=0)
+            output_reg[video_name] = np.stack(output_reg[video_name], axis=0)
+        if labels_cls[video_name]:
+            labels_cls[video_name] = np.stack(labels_cls[video_name], axis=0)
+        if labels_reg[video_name]:
+            labels_reg[video_name] = np.stack(labels_reg[video_name], axis=0)
     
     return output_cls, output_reg, labels_cls, labels_reg
 
 def eval_map_nms(opt, dataset, output_cls, output_reg):
+    """Evaluate with Non-Maximum Suppression"""
     result_dict = {}
-    proposal_dict = []
     anchors = opt['anchors']
     
     for video_name in dataset.video_list:
+        if video_name not in output_cls or len(output_cls[video_name]) == 0:
+            result_dict[video_name] = []
+            continue
+            
         duration = dataset.video_len[video_name]
         video_time = float(dataset.video_dict[video_name]["duration"])
         frame_to_time = 100.0 * video_time / duration
         
-        for idx in range(duration):
+        proposal_dict = []
+        
+        for idx in range(min(duration, len(output_cls[video_name]))):
             cls_anc = output_cls[video_name][idx]
             reg_anc = output_reg[video_name][idx]
-            proposal_anc_dict = []
             
             for anc_idx in range(len(anchors)):
+                if anc_idx >= len(cls_anc):
+                    continue
+                    
                 cls = np.argwhere(cls_anc[anc_idx][:-1] > opt['threshold']).reshape(-1)
                 if len(cls) == 0:
                     continue
+                    
                 ed = idx + anchors[anc_idx] * reg_anc[anc_idx][0]
                 length = anchors[anc_idx] * np.exp(reg_anc[anc_idx][1])
                 st = ed - length
+                
                 for cidx in range(len(cls)):
                     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]),
-                        "label": dataset.label_name[label],
-                        "gentime": float(idx * frame_to_time / 100.0)
-                    }
-                    proposal_anc_dict.append(tmp_dict)
-            
-            proposal_dict += proposal_anc_dict
+                    if label < len(dataset.label_name):
+                        tmp_dict = {
+                            "segment": [float(st * frame_to_time / 100.0), float(ed * frame_to_time / 100.0)],
+                            "score": float(cls_anc[anc_idx][label]),
+                            "label": dataset.label_name[label],
+                            "gentime": float(idx * frame_to_time / 100.0)
+                        }
+                        proposal_dict.append(tmp_dict)
         
+        # Apply NMS
         proposal_dict = non_max_suppression(proposal_dict, overlapThresh=opt['soft_nms'])
         result_dict[video_name] = proposal_dict
-        proposal_dict = []
     
     return result_dict
 
-def process_input(video_file, npz_file, checkpoint_path, split_number):
-    # Parse options
-    opt = opts.parse_opt()
-    opt = vars(opt)
-    opt['mode'] = 'test'
-    opt['split'] = str(split_number)
-    opt['checkpoint_path'] = './checkpoint'
-    opt['video_feature_all_test'] = './data/I3D/'
-    opt['anchors'] = [int(item) for item in opt['anchors'].split(',')]
-    opt['batch_size'] = 1  # Single video processing
-    os.makedirs(opt['checkpoint_path'], exist_ok=True)
-    os.makedirs(opt['video_feature_all_test'], exist_ok=True)
-    
-    # Handle input
-    video_name = "user_upload"
-    video_path = None
-    if video_file:
-        video_path = video_file
-        # Placeholder for I3D feature extraction (to be implemented or assumed precomputed)
-        return "Error: Real-time I3D feature extraction not supported. Please upload .npz file."
-    
-    if npz_file:
-        npz_path = os.path.join(opt['video_feature_all_test'], f"{video_name}.npz")
-        os.makedirs(os.path.dirname(npz_path), exist_ok=True)
-        np.savez(npz_path, rgb=np.load(npz_file)['rgb'], flow=np.load(npz_file)['flow'])
-    
-    # Load model
-    model = MYNET(opt).to(device)
-    checkpoint = torch.load(checkpoint_path, map_location=device)
-    model.load_state_dict(checkpoint['state_dict'])
-    model.eval()
-    
-    # Create dataset
-    dataset = VideoDataSet(opt, subset='test', video_name=video_name)
-    
-    # Run inference
-    output_cls, output_reg, labels_cls, labels_reg = eval_frame(opt, model, dataset)
-    result_dict = eval_map_nms(opt, dataset, output_cls, output_reg)
-    
-    # Load annotations if available
+def load_ground_truth(opt, video_name):
+    """Load ground truth annotations if available"""
     gt_segments = []
     duration = 0
-    video_anno_file = opt["video_anno"].format(opt["split"])
-    if os.path.exists(video_anno_file):
-        with open(video_anno_file, 'r') as f:
-            anno_data = json.load(f)
-        if video_name in anno_data['database']:
-            gt_annotations = anno_data['database'][video_name]['annotations']
-            duration = anno_data['database'][video_name]['duration']
-            for anno in gt_annotations:
-                start, end = anno['segment']
-                gt_segments.append({'label': anno['label'], 'start': start, 'end': end, 'duration': end - start})
     
-    pred_segments = []
-    for pred in result_dict.get(video_name, []):
-        start, end = pred['segment']
-        pred_segments.append({
-            'label': pred['label'],
-            'start': start,
-            'end': end,
-            'duration': end - start,
-            'score': pred['score']
-        })
-    
-    # Generate comparison table
-    output_text = f"Predicted Actions for Video: {video_name}\n\n"
-    if gt_segments:
-        matches = []
-        iou_threshold = VIS_CONFIG['iou_threshold']
-        used_gt_indices = set()
-        for pred in pred_segments:
-            best_iou = 0
-            best_gt_idx = None
-            for gt_idx, gt in enumerate(gt_segments):
-                if gt_idx in used_gt_indices:
-                    continue
-                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
-            if best_gt_idx is not None:
-                matches.append({'pred': pred, 'gt': gt_segments[best_gt_idx], 'iou': best_iou})
-                used_gt_indices.add(best_gt_idx)
-            else:
-                matches.append({'pred': pred, 'gt': None, 'iou': 0})
+    try:
+        video_anno_file = opt["video_anno"].format(opt["split"])
+        if os.path.exists(video_anno_file):
+            with open(video_anno_file, 'r') as f:
+                anno_data = json.load(f)
+            
+            if video_name in anno_data['database']:
+                gt_annotations = anno_data['database'][video_name]['annotations']
+                duration = anno_data['database'][video_name]['duration']
+                
+                for anno in gt_annotations:
+                    start, end = anno['segment']
+                    gt_segments.append({
+                        'label': anno['label'], 
+                        'start': start, 
+                        'end': end, 
+                        'duration': end - start
+                    })
+    except Exception as e:
+        print(f"Could not load ground truth: {str(e)}")
+    
+    return gt_segments, duration
+
+def process_video(video_name, split_number):
+    """Process a single video for action localization"""
+    try:
+        # Parse options
+        opt = opts.parse_opt()
+        opt = vars(opt)
+        opt['mode'] = 'test'
+        opt['split'] = str(split_number)
+        opt['checkpoint_path'] = './checkpoint'
+        opt['video_feature_all_test'] = './data/I3D/'
+        opt['anchors'] = [int(item) for item in opt['anchors'].split(',')]
+        opt['batch_size'] = 1
         
-        for gt_idx, gt in enumerate(gt_segments):
-            if gt_idx not in used_gt_indices:
-                matches.append({'pred': None, 'gt': gt, 'iou': 0})
+        # Check if required files exist
+        checkpoint_path = './checkpoint/01_ckp_best.pth.tar'
+        if not os.path.exists(checkpoint_path):
+            return "Error: Model checkpoint not found at ./checkpoint/01_ckp_best.pth.tar"
         
-        output_text += "{:<20} {:<30} {:<30} {:<15} {:<10}\n".format(
-            "Action Label", "Predicted Segment (s)", "Ground Truth Segment (s)", "Duration Diff (s)", "IoU")
-        output_text += "-" * 105 + "\n"
-        for match in matches:
-            pred = match['pred']
-            gt = match['gt']
-            iou = match['iou']
-            if pred and gt:
-                label = pred['label'] if pred['label'] == gt['label'] else f"{pred['label']} (GT: {gt['label']})"
-                pred_str = f"[{pred['start']:.2f}, {pred['end']:.2f}] ({pred['duration']:.2f}s)"
-                gt_str = f"[{gt['start']:.2f}, {gt['end']:.2f}] ({gt['duration']:.2f}s)"
-                duration_diff = pred['duration'] - gt['duration']
-                output_text += "{:<20} {:<30} {:<30} {:<15.2f} {:<10.2f}\n".format(
-                    label, pred_str, gt_str, duration_diff, iou)
-            elif pred:
-                pred_str = f"[{pred['start']:.2f}, {pred['end']:.2f}] ({pred['duration']:.2f}s)"
-                output_text += "{:<20} {:<30} {:<30} {:<15} {:<10.2f}\n".format(
-                    pred['label'], pred_str, "None", "N/A", iou)
-            elif gt:
-                gt_str = f"[{gt['start']:.2f}, {gt['end']:.2f}] ({gt['duration']:.2f}s)"
-                output_text += "{:<20} {:<30} {:<30} {:<15} {:<10.2f}\n".format(
-                    gt['label'], "None", gt_str, "N/A", iou)
+        npz_path = os.path.join(opt['video_feature_all_test'], f"{video_name}.npz")
+        if not os.path.exists(npz_path):
+            return f"Error: Feature file not found at {npz_path}"
         
-        matched_count = sum(1 for m in matches if m['pred'] and m['gt'])
-        avg_duration_diff = np.mean([m['pred']['duration'] - m['gt']['duration'] for m in matches if m['pred'] and m['gt']]) if matched_count > 0 else 0
-        avg_iou = np.mean([m['iou'] for m in matches if m['iou'] > 0]) if any(m['iou'] > 0 for m in matches) else 0
-        output_text += "\nSummary:\n"
-        output_text += f"- Total Predictions: {len(pred_segments)}\n"
-        output_text += f"- Total Ground Truth: {len(gt_segments)}\n"
-        output_text += f"- Matched Segments: {matched_count}\n"
-        output_text += f"- Average Duration Difference (Matched): {avg_duration_diff:.2f}s\n"
-        output_text += f"- Average IoU (Matched): {avg_iou:.2f}\n"
-    else:
-        output_text += "No ground truth annotations available.\nPredicted Segments:\n"
-        for pred in pred_segments:
-            output_text += f"- {pred['label']}: [{pred['start']:.2f}, {pred['end']:.2f}] ({pred['duration']:.2f}s), Score: {pred['score']:.2f}\n"
-    
-    # Generate visualizations
-    viz_path = None
-    video_out_path = None
-    if video_file and os.path.exists(video_file):
-        duration = max([seg['end'] for seg in pred_segments + gt_segments], default=1.0)
-        viz_path = visualize_action_lengths(video_name, pred_segments, gt_segments, video_file, duration)
-        video_out_path = annotate_video_with_actions(video_name, pred_segments, gt_segments, video_file)
-    
-    return output_text, viz_path, video_out_path
+        # Load model
+        model = MYNET(opt).to(device)
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        
+        # Handle different checkpoint formats
+        if 'state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['state_dict'])
+        else:
+            model.load_state_dict(checkpoint)
+        
+        model.eval()
+        
+        # Create dataset
+        dataset = VideoDataSet(opt, subset='test', video_name=video_name)
+        
+        if len(dataset.video_list) == 0:
+            return f"Error: No video found with name '{video_name}' in dataset"
+        
+        # Run inference
+        output_cls, output_reg, labels_cls, labels_reg = eval_frame(opt, model, dataset)
+        result_dict = eval_map_nms(opt, dataset, output_cls, output_reg)
+        
+        # Load ground truth
+        gt_segments, duration = load_ground_truth(opt, video_name)
+        
+        # Process predictions
+        pred_segments = []
+        for pred in result_dict.get(video_name, []):
+            start, end = pred['segment']
+            pred_segments.append({
+                'label': pred['label'],
+                'start': start,
+                'end': end,
+                'duration': end - start,
+                'score': pred['score']
+            })
+        
+        # Generate output text
+        output_text = f"Predicted Actions for Video: {video_name}\n"
+        output_text += "=" * 50 + "\n\n"
+        
+        if pred_segments:
+            output_text += "PREDICTED ACTIONS:\n"
+            output_text += "-" * 30 + "\n"
+            for i, pred in enumerate(pred_segments, 1):
+                output_text += f"{i}. {pred['label']}\n"
+                output_text += f"   Time: [{pred['start']:.2f}s - {pred['end']:.2f}s]\n"
+                output_text += f"   Duration: {pred['duration']:.2f}s\n"
+                output_text += f"   Confidence: {pred['score']:.3f}\n\n"
+        else:
+            output_text += "No actions detected above threshold.\n\n"
+        
+        # Add ground truth comparison if available
+        if gt_segments:
+            output_text += "\nGROUND TRUTH COMPARISON:\n"
+            output_text += "-" * 30 + "\n"
+            
+            # Calculate basic metrics
+            matched_count = 0
+            total_pred = len(pred_segments)
+            total_gt = len(gt_segments)
+            
+            for gt in gt_segments:
+                output_text += f"GT: {gt['label']} [{gt['start']:.2f}s - {gt['end']:.2f}s]\n"
+                
+                # Find best matching prediction
+                best_match = None
+                best_iou = 0
+                for pred in pred_segments:
+                    # Simple overlap calculation
+                    overlap_start = max(gt['start'], pred['start'])
+                    overlap_end = min(gt['end'], pred['end'])
+                    if overlap_end > overlap_start:
+                        overlap = overlap_end - overlap_start
+                        union = (gt['end'] - gt['start']) + (pred['end'] - pred['start']) - overlap
+                        iou = overlap / union if union > 0 else 0
+                        if iou > best_iou:
+                            best_iou = iou
+                            best_match = pred
+                
+                if best_match and best_iou > VIS_CONFIG['iou_threshold']:
+                    matched_count += 1
+                    output_text += f"    → Matched with: {best_match['label']} (IoU: {best_iou:.3f})\n"
+                else:
+                    output_text += f"    → No match found\n"
+                output_text += "\n"
+            
+            # Summary statistics
+            precision = matched_count / total_pred if total_pred > 0 else 0
+            recall = matched_count / total_gt if total_gt > 0 else 0
+            f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
+            
+            output_text += f"\nSUMMARY STATISTICS:\n"
+            output_text += f"Total Predictions: {total_pred}\n"
+            output_text += f"Total Ground Truth: {total_gt}\n"
+            output_text += f"Matched: {matched_count}\n"
+            output_text += f"Precision: {precision:.3f}\n"
+            output_text += f"Recall: {recall:.3f}\n"
+            output_text += f"F1-Score: {f1:.3f}\n"
+        
+        return output_text
+        
+    except Exception as e:
+        return f"Error processing video: {str(e)}\n\nPlease check:\n1. Model checkpoint exists\n2. Feature file exists\n3. All dependencies are installed"
+
+def get_available_videos():
+    """Get list of available videos from I3D features directory"""
+    feature_dir = './data/I3D/'
+    if not os.path.exists(feature_dir):
+        return []
+    
+    videos = []
+    for file in os.listdir(feature_dir):
+        if file.endswith('.npz'):
+            video_name = file.replace('.npz', '')
+            videos.append(video_name)
+    
+    return sorted(videos)
+
+# Initialize available videos
+available_videos = get_available_videos()
+if not available_videos:
+    available_videos = ["No videos found"]
 
 # Gradio Interface
 iface = gr.Interface(
-    fn=process_input,
+    fn=process_video,
     inputs=[
-    #     gr.Video(label="Upload Video (Optional, requires .npz for processing)"),
-        gr.File(label="Upload I3D .npz File"),
-    #     gr.File(label="Upload Model Checkpoint (.pth.tar)", file_types=[".pth.tar"]),
-        gr.Dropdown(label="Split Number", choices=["1", "2", "3"], value="1")
+        gr.Dropdown(
+            label="Select Video", 
+            choices=available_videos, 
+            value=available_videos[0] if available_videos else None,
+            info="Choose from pre-uploaded videos in data/I3D/ folder"
+        ),
+        gr.Dropdown(
+            label="Split Number", 
+            choices=["1", "2", "3"], 
+            value="1",
+            info="Dataset split for annotations"
+        )
     ],
     outputs=[
-        gr.Textbox(label="Action Predictions"),
-        # gr.Image(label="Action Visualization", type="filepath"),
-        # gr.Video(label="Annotated Video")
+        gr.Textbox(
+            label="Action Predictions", 
+            lines=20,
+            max_lines=50,
+            show_copy_button=True
+        )
     ],
-    title="Temporal Action Localization",
-    description="Upload an I3D .npz file and a trained model checkpoint to predict actions. Optionally upload a video to generate visualizations. Select the annotation split number."
+    title="🎬 Temporal Action Localization",
+    description="""
+    This app performs temporal action localization on pre-uploaded videos using I3D features.
+    
+    **How to use:**
+    1. Select a video from the dropdown (videos must be in data/I3D/ folder as .npz files)
+    2. Choose the annotation split number
+    3. Click Submit to get action predictions
+    
+    **Requirements:**
+    - Model checkpoint: `01_ckp_best.pth.tar` in root directory
+    - Video features: `.npz` files in `data/I3D/` folder
+    """,
+    examples=[
+        [available_videos[0] if available_videos and available_videos[0] != "No videos found" else "example_video", "1"],
+    ] if available_videos and available_videos[0] != "No videos found" else None,
+    cache_examples=False,
+    theme=gr.themes.Soft()
 )
 
 if __name__ == '__main__':
-    opt = opts.parse_opt()
-    opt = vars(opt)
-    opt['anchors'] = [int(item) for item in opt['anchors'].split(',')]
-    iface.launch()
+    print(f"Available videos: {available_videos}")
+    print(f"Using device: {device}")
+    iface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )