Update app.py

app.py

@@ -32,17 +32,61 @@ class DroneDetectionPipeline:
         self._load_model()
         
     def _load_model(self) -> None:
-        """Load YOLO model with error handling"""
+        """Load YOLO model with error handling and PyTorch compatibility"""
         try:
             if os.path.exists(self.model_path):
-                self.model = YOLO(self.model_path)
+                # Handle PyTorch 2.6+ security restrictions for model loading
+                import torch
+                from ultralytics.nn.tasks import DetectionModel
+                
+                # Add safe globals for ultralytics models
+                torch.serialization.add_safe_globals([DetectionModel])
+                
+                # Load model with weights_only=False for older model formats
+                try:
+                    self.model = YOLO(self.model_path)
+                except Exception as pytorch_error:
+                    logger.warning(f"⚠️ Standard loading failed: {str(pytorch_error)}")
+                    logger.info("🔄 Attempting alternative loading method...")
+                    
+                    # Alternative loading for older PyTorch model formats
+                    import tempfile
+                    import shutil
+                    
+                    # Create temporary model with legacy loading
+                    with tempfile.TemporaryDirectory() as temp_dir:
+                        temp_model_path = os.path.join(temp_dir, "temp_model.pt")
+                        shutil.copy2(self.model_path, temp_model_path)
+                        
+                        # Patch torch.load temporarily for this specific model
+                        original_load = torch.load
+                        def patched_load(f, **kwargs):
+                            kwargs['weights_only'] = False
+                            return original_load(f, **kwargs)
+                        
+                        torch.load = patched_load
+                        try:
+                            self.model = YOLO(temp_model_path)
+                        finally:
+                            torch.load = original_load
+                
                 logger.info(f"✅ Model loaded successfully from {self.model_path}")
             else:
                 logger.error(f"❌ Model file not found: {self.model_path}")
-                raise FileNotFoundError(f"Model file not found: {self.model_path}")
+                # Try to load a default YOLOv8 model if custom model is not found
+                logger.info("🔄 Loading default YOLOv8n model as fallback...")
+                self.model = YOLO('yolov8n.pt')
+                logger.warning("⚠️ Using default YOLOv8n model - detection accuracy may be reduced")
+                
         except Exception as e:
             logger.error(f"❌ Error loading model: {str(e)}")
-            raise
+            logger.info("🔄 Attempting to load default YOLOv8n model as final fallback...")
+            try:
+                self.model = YOLO('yolov8n.pt')
+                logger.warning("⚠️ Using default YOLOv8n model - please upload your custom 'best.pt' model for optimal drone detection")
+            except Exception as fallback_error:
+                logger.error(f"❌ Failed to load any model: {str(fallback_error)}")
+                raise RuntimeError("Could not load any YOLO model. Please check your model file and dependencies.")
     
     def _initialize_tracker(self, frame: np.ndarray, bbox: Tuple[int, int, int, int]) -> bool:
         """Initialize CSRT tracker with given bounding box"""
@@ -210,22 +254,39 @@ class DroneDetectionPipeline:
 pipeline = DroneDetectionPipeline()
 
 def process_video_gradio(input_video, confidence_threshold):
-    """Gradio interface function"""
+    """Gradio interface function with enhanced error handling"""
     if input_video is None:
         return None, "❌ Please upload a video file"
     
     try:
+        # Check if model is loaded properly
+        if pipeline.model is None:
+            return None, "❌ Model not loaded. Please check that 'best.pt' is uploaded to the space."
+        
         # Update confidence threshold
         pipeline.confidence_threshold = confidence_threshold
         
+        # Add file size check
+        file_size = os.path.getsize(input_video)
+        max_size = 200 * 1024 * 1024  # 200MB
+        if file_size > max_size:
+            return None, f"❌ File too large ({file_size/1024/1024:.1f}MB). Maximum size is 200MB."
+        
+        logger.info(f"📹 Processing video: {input_video} ({file_size/1024/1024:.1f}MB)")
+        
         # Process video
         output_path = pipeline.process_video(input_video)
         
-        return output_path, "✅ Video processing completed successfully!"
+        success_msg = "✅ Video processing completed successfully!"
+        if "yolov8n" in str(pipeline.model.model):
+            success_msg += "\n⚠️ Note: Using default YOLOv8n model. Upload 'best.pt' for optimal drone detection."
+        
+        return output_path, success_msg
         
     except Exception as e:
         error_msg = f"❌ Error processing video: {str(e)}"
         logger.error(error_msg)
+        logger.error("Full traceback:", exc_info=True)
         return None, error_msg
 
 # Gradio Interface
@@ -304,13 +365,24 @@ def create_interface():
         gr.Markdown(article)
         
         # Example section
-        gr.Markdown("### 📋 System Requirements")
+        gr.Markdown("### 📋 System Requirements & Setup")
         gr.Markdown("""
-        - Input video formats: MP4, AVI, MOV, MKV
-        - Maximum file size: 200MB
-        - Recommended resolution: Up to 1920x1080
-        - Processing time: ~1-2 minutes per minute of video
+        - **Input video formats**: MP4, AVI, MOV, MKV
+        - **Maximum file size**: 200MB
+        - **Recommended resolution**: Up to 1920x1080
+        - **Processing time**: ~1-2 minutes per minute of video
+        
+        **📝 Setup Instructions:**
+        1. Upload your custom trained `best.pt` model to this space for optimal drone detection
+        2. If no custom model is found, the system will use YOLOv8n as fallback
+        3. For best results, use the custom drone detection model trained specifically for your use case
+        
+        **🔧 Model Upload:** Go to "Files and versions" tab → Upload your `best.pt` file → Restart space
         """)
+        
+        # Status indicator
+        model_status = "🟢 Custom Model" if os.path.exists("best.pt") else "🟡 Default YOLOv8n (Upload best.pt for better results)"
+        gr.Markdown(f"**Current Model Status**: {model_status}")
     
     return interface