Update main.py

main.py

@@ -371,6 +371,11 @@
 
 
 
+
+
+
+
+
 import os
 import json
 import torch
@@ -386,6 +391,7 @@ from typing import List, Dict, Optional
 from huggingface_hub import hf_hub_download, list_repo_files
 import tempfile
 import shutil
+import traceback
 
 # Configuration
 VIS_CONFIG = {
@@ -473,6 +479,7 @@ class HFVideoDataSet(VideoDataSet):
         
         # Create temporary directory for this session
         self.temp_dir = tempfile.mkdtemp(prefix="hf_video_")
+        print(f"Created temp directory: {self.temp_dir}")
         
         # Download the specific video file if video_name is provided
         if video_name:
@@ -483,24 +490,39 @@ class HFVideoDataSet(VideoDataSet):
                 # Ensure the temp directory exists
                 os.makedirs(self.temp_dir, exist_ok=True)
                 
-                # Copy to temp directory with expected structure
+                # Copy to temp directory with expected structure - FIX: Add proper path separator
                 temp_file_path = os.path.join(self.temp_dir, f"{video_name}.npz")
+                print(f"Copying {downloaded_path} to {temp_file_path}")
                 shutil.copy2(downloaded_path, temp_file_path)
-                print(f"Video file ready: {temp_file_path}")
                 
-                # Verify file exists
+                # Verify file exists and print debug info
                 if not os.path.exists(temp_file_path):
                     raise Exception(f"Failed to copy file to {temp_file_path}")
+                else:
+                    print(f"Video file ready: {temp_file_path}")
+                    print(f"File size: {os.path.getsize(temp_file_path)} bytes")
                     
             except Exception as e:
                 print(f"Error downloading video {video_name}: {str(e)}")
+                # Clean up temp directory on error
+                if hasattr(self, 'temp_dir') and os.path.exists(self.temp_dir):
+                    shutil.rmtree(self.temp_dir)
                 raise e
         
         # Set the feature path to our temp directory
         opt['video_feature_all_test'] = self.temp_dir
+        print(f"Set video_feature_all_test to: {opt['video_feature_all_test']}")
         
         # Initialize parent class
-        super().__init__(opt, subset, video_name)
+        try:
+            super().__init__(opt, subset, video_name)
+            print(f"Successfully initialized dataset with {len(self.video_list)} videos")
+        except Exception as e:
+            print(f"Error initializing parent VideoDataSet: {str(e)}")
+            # Clean up temp directory on error
+            if hasattr(self, 'temp_dir') and os.path.exists(self.temp_dir):
+                shutil.rmtree(self.temp_dir)
+            raise e
     
     def __del__(self):
         # Clean up temporary directory
@@ -513,111 +535,121 @@ class HFVideoDataSet(VideoDataSet):
 
 def eval_frame(opt, model, dataset):
     """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}
-    
-    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())
+    try:
+        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}
+        
+        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())
                             
-            except Exception as e:
-                print(f"Error in batch {n_iter}: {str(e)}")
-                continue
-    
-    # Stack arrays
-    for video_name in dataset.video_list:
-        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
+                            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:
+            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
+        
+    except Exception as e:
+        print(f"Error in eval_frame: {str(e)}")
+        raise e
 
 def eval_map_nms(opt, dataset, output_cls, output_reg):
     """Evaluate with Non-Maximum Suppression"""
-    result_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
-        
-        proposal_dict = []
+    try:
+        result_dict = {}
+        anchors = opt['anchors']
         
-        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]
+        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 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
+            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]
                 
-                for cidx in range(len(cls)):
-                    label = cls[cidx]
-                    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
-    
-    return result_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]
+                        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
+        
+        return result_dict
+        
+    except Exception as e:
+        print(f"Error in eval_map_nms: {str(e)}")
+        raise e
 
 def load_ground_truth(opt, video_name):
     """Load ground truth annotations if available"""
@@ -649,6 +681,8 @@ def load_ground_truth(opt, video_name):
 
 def process_video(video_name, split_number, progress=gr.Progress()):
     """Process a single video for action localization"""
+    dataset = None  # Initialize dataset variable
+    
     try:
         if not video_name or video_name in ["Error: Could not load videos from HF dataset", "Error loading videos"]:
             return "Error: Please select a valid video name"
@@ -694,14 +728,18 @@ def process_video(video_name, split_number, progress=gr.Progress()):
             model.load_state_dict(checkpoint)
         
         model.eval()
+        print("Model loaded successfully")
         
         progress(0.4, desc=f"Downloading video features for {video_name}...")
         
         # Create dataset with HF integration
         try:
             dataset = HFVideoDataSet(opt, subset='test', video_name=video_name)
+            print(f"Dataset created successfully with {len(dataset.video_list)} videos")
         except Exception as e:
-            return f"Error downloading or loading video '{video_name}': {str(e)}\n\nPlease check:\n1. Video name is correct\n2. File exists in HF dataset\n3. Network connection is stable"
+            error_msg = f"Error downloading or loading video '{video_name}': {str(e)}\n\nPlease check:\n1. Video name is correct\n2. File exists in HF dataset\n3. Network connection is stable"
+            print(error_msg)
+            return error_msg
         
         if len(dataset.video_list) == 0:
             return f"Error: No video found with name '{video_name}' in dataset after download"
@@ -709,11 +747,23 @@ def process_video(video_name, split_number, progress=gr.Progress()):
         progress(0.6, desc="Running inference...")
         
         # Run inference
-        output_cls, output_reg, labels_cls, labels_reg = eval_frame(opt, model, dataset)
+        try:
+            output_cls, output_reg, labels_cls, labels_reg = eval_frame(opt, model, dataset)
+            print("Inference completed successfully")
+        except Exception as e:
+            error_msg = f"Error during inference: {str(e)}"
+            print(error_msg)
+            return error_msg
         
         progress(0.8, desc="Processing results...")
         
-        result_dict = eval_map_nms(opt, dataset, output_cls, output_reg)
+        try:
+            result_dict = eval_map_nms(opt, dataset, output_cls, output_reg)
+            print("NMS processing completed")
+        except Exception as e:
+            error_msg = f"Error during NMS processing: {str(e)}"
+            print(error_msg)
+            return error_msg
         
         # Load ground truth
         gt_segments, duration = load_ground_truth(opt, video_name)
@@ -796,21 +846,39 @@ def process_video(video_name, split_number, progress=gr.Progress()):
             output_text += f"F1-Score: {f1:.3f}\n"
         
         progress(1.0, desc="Complete!")
+        print("Processing completed successfully")
         return output_text
         
     except Exception as e:
-        import traceback
         error_details = traceback.format_exc()
-        return f"Error processing video: {str(e)}\n\nDetailed error:\n{error_details}\n\nPlease check:\n1. Model checkpoint exists\n2. Video exists in HF dataset\n3. All dependencies are installed"
+        error_msg = f"Error processing video: {str(e)}\n\nDetailed error:\n{error_details}\n\nPlease check:\n1. Model checkpoint exists\n2. Video exists in HF dataset\n3. All dependencies are installed"
+        print(error_msg)
+        return error_msg
+    finally:
+        # Ensure cleanup happens even if there's an error
+        if dataset is not None and hasattr(dataset, '__del__'):
+            try:
+                dataset.__del__()
+            except Exception as e:
+                print(f"Warning: Error during dataset cleanup: {e}")
 
 def refresh_video_list():
     """Refresh the list of available videos"""
-    return gr.Dropdown(choices=get_available_videos_from_hf())
+    try:
+        new_videos = get_available_videos_from_hf()
+        return gr.Dropdown(choices=new_videos)
+    except Exception as e:
+        print(f"Error refreshing video list: {e}")
+        return gr.Dropdown(choices=["Error refreshing videos"])
 
 # Initialize available videos
 print("Loading available videos from Hugging Face dataset...")
-available_videos = get_available_videos_from_hf()
-if not available_videos or available_videos == ["Error loading videos"]:
+try:
+    available_videos = get_available_videos_from_hf()
+    if not available_videos or available_videos == ["Error loading videos"]:
+        available_videos = ["Error: Could not load videos from HF dataset"]
+except Exception as e:
+    print(f"Error loading initial video list: {e}")
     available_videos = ["Error: Could not load videos from HF dataset"]
 
 print(f"Available videos: {len(available_videos)} videos found")
@@ -867,7 +935,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🎬 Temporal Action Localization"
     
     # Event handlers
     refresh_btn.click(
-        fn=lambda: gr.Dropdown(choices=get_available_videos_from_hf()),
+        fn=refresh_video_list,
         outputs=video_dropdown
     )