Create app.py

app.py

@@ -0,0 +1,326 @@
+import gradio as gr
+import cv2
+import numpy as np
+import tempfile
+import os
+from pathlib import Path
+from ultralytics import YOLO
+import torch
+from typing import Optional, Tuple, List
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class DroneDetectionPipeline:
+    """
+    Professional drone detection pipeline with YOLO + CSRT tracking
+    Designed for UK International Lab deployment
+    """
+    
+    def __init__(self, model_path: str = "best.pt"):
+        """Initialize the detection pipeline with YOLO model and CSRT tracker"""
+        self.model_path = model_path
+        self.model = None
+        self.tracker = None
+        self.tracking_active = False
+        self.last_detection_bbox = None
+        self.confidence_threshold = 0.5
+        
+        # Load model
+        self._load_model()
+        
+    def _load_model(self) -> None:
+        """Load YOLO model with error handling"""
+        try:
+            if os.path.exists(self.model_path):
+                self.model = YOLO(self.model_path)
+                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}")
+        except Exception as e:
+            logger.error(f"❌ Error loading model: {str(e)}")
+            raise
+    
+    def _initialize_tracker(self, frame: np.ndarray, bbox: Tuple[int, int, int, int]) -> bool:
+        """Initialize CSRT tracker with given bounding box"""
+        try:
+            self.tracker = cv2.TrackerCSRT_create()
+            success = self.tracker.init(frame, bbox)
+            if success:
+                self.tracking_active = True
+                self.last_detection_bbox = bbox
+                logger.info("✅ Tracker initialized successfully")
+            return success
+        except Exception as e:
+            logger.error(f"❌ Error initializing tracker: {str(e)}")
+            return False
+    
+    def _detect_drones(self, frame: np.ndarray) -> List[Tuple[int, int, int, int, float]]:
+        """Run YOLO inference on frame and return detections"""
+        try:
+            results = self.model(frame, verbose=False, conf=self.confidence_threshold)
+            detections = []
+            
+            for result in results:
+                if result.boxes is not None:
+                    for box in result.boxes:
+                        # Extract coordinates and confidence
+                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+                        confidence = float(box.conf[0].cpu().numpy())
+                        detections.append((x1, y1, x2, y2, confidence))
+            
+            return detections
+        except Exception as e:
+            logger.error(f"❌ Error in detection: {str(e)}")
+            return []
+    
+    def _draw_detection(self, frame: np.ndarray, bbox: Tuple[int, int, int, int], 
+                       confidence: float = None, is_tracking: bool = False) -> np.ndarray:
+        """Draw bounding box and label on frame"""
+        x1, y1, x2, y2 = bbox
+        
+        # Choose color: Red for detection, Blue for tracking
+        color = (0, 0, 255) if not is_tracking else (255, 0, 0)
+        label_text = f"Drone (Det)" if not is_tracking else f"Drone (Track)"
+        
+        if confidence is not None and not is_tracking:
+            label_text = f"Drone {confidence:.2f}"
+        
+        # Draw bounding box
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+        
+        # Draw label background
+        label_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
+        cv2.rectangle(frame, (x1, y1 - label_size[1] - 10), 
+                     (x1 + label_size[0], y1), color, -1)
+        
+        # Draw label text
+        cv2.putText(frame, label_text, (x1, y1 - 5), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+        
+        return frame
+    
+    def process_video(self, input_video_path: str, progress_callback=None) -> str:
+        """
+        Process entire video with drone detection and tracking
+        
+        Args:
+            input_video_path: Path to input video
+            progress_callback: Optional callback for progress updates
+            
+        Returns:
+            Path to processed output video
+        """
+        # Create temporary output file
+        output_dir = tempfile.mkdtemp()
+        output_path = os.path.join(output_dir, "drone_detection_output.mp4")
+        
+        cap = cv2.VideoCapture(input_video_path)
+        
+        if not cap.isOpened():
+            raise ValueError("❌ Could not open input video")
+        
+        # Get video properties
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        
+        # Initialize video writer
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+        
+        frame_count = 0
+        detection_count = 0
+        tracking_count = 0
+        
+        logger.info(f"🎬 Processing video: {total_frames} frames at {fps} FPS")
+        
+        try:
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                frame_processed = frame.copy()
+                current_detections = self._detect_drones(frame)
+                
+                if current_detections:
+                    # We have detections - use the largest detection
+                    largest_detection = max(current_detections, 
+                                          key=lambda d: (d[2]-d[0]) * (d[3]-d[1]))
+                    
+                    x1, y1, x2, y2, conf = largest_detection
+                    detection_bbox = (x1, y1, x2-x1, y2-y1)  # Convert to (x, y, w, h)
+                    
+                    # Draw detection
+                    frame_processed = self._draw_detection(frame_processed, 
+                                                         (x1, y1, x2, y2), conf, False)
+                    
+                    # Reinitialize tracker with new detection
+                    self._initialize_tracker(frame, detection_bbox)
+                    detection_count += 1
+                    
+                elif self.tracking_active and self.tracker is not None:
+                    # No detections but tracker is active
+                    success, tracking_bbox = self.tracker.update(frame)
+                    
+                    if success:
+                        x, y, w, h = [int(v) for v in tracking_bbox]
+                        self.last_detection_bbox = tracking_bbox
+                        
+                        # Draw tracking box
+                        frame_processed = self._draw_detection(frame_processed, 
+                                                             (x, y, x+w, y+h), None, True)
+                        tracking_count += 1
+                    else:
+                        # Tracking failed
+                        if self.last_detection_bbox is not None:
+                            # Use last known position
+                            x, y, w, h = [int(v) for v in self.last_detection_bbox]
+                            frame_processed = self._draw_detection(frame_processed, 
+                                                                 (x, y, x+w, y+h), None, True)
+                        self.tracking_active = False
+                
+                # Write processed frame
+                out.write(frame_processed)
+                frame_count += 1
+                
+                # Update progress
+                if progress_callback and frame_count % 10 == 0:
+                    progress = frame_count / total_frames
+                    progress_callback(progress, f"Processing frame {frame_count}/{total_frames}")
+        
+        except Exception as e:
+            logger.error(f"❌ Error during video processing: {str(e)}")
+            raise
+        finally:
+            cap.release()
+            out.release()
+        
+        logger.info(f"✅ Processing completed!")
+        logger.info(f"📊 Stats: {detection_count} detections, {tracking_count} tracking frames")
+        
+        return output_path
+
+# Initialize pipeline
+pipeline = DroneDetectionPipeline()
+
+def process_video_gradio(input_video, confidence_threshold):
+    """Gradio interface function"""
+    if input_video is None:
+        return None, "❌ Please upload a video file"
+    
+    try:
+        # Update confidence threshold
+        pipeline.confidence_threshold = confidence_threshold
+        
+        # Process video
+        output_path = pipeline.process_video(input_video)
+        
+        return output_path, "✅ Video processing completed successfully!"
+        
+    except Exception as e:
+        error_msg = f"❌ Error processing video: {str(e)}"
+        logger.error(error_msg)
+        return None, error_msg
+
+# Gradio Interface
+def create_interface():
+    """Create professional Gradio interface"""
+    
+    title = "🚁 Professional Drone Detection System"
+    description = """
+    **Advanced Drone Detection and Tracking Pipeline**
+    
+    This system uses state-of-the-art YOLO object detection combined with CSRT tracking for robust drone detection in video footage.
+    
+    **Features:**
+    - Real-time drone detection with confidence scoring
+    - Seamless tracking when detections are lost
+    - Professional-grade output for research and analysis
+    
+    **Developed for UK International Lab**
+    """
+    
+    article = """
+    ### Technical Details
+    - **Detection Model**: Custom trained YOLO model optimized for drone detection
+    - **Tracking Algorithm**: CSRT (Channel and Spatial Reliability Tracker)
+    - **Output Format**: MP4 video with annotated detections and tracking
+    
+    ### Usage Instructions
+    1. Upload your video file (supported formats: MP4, AVI, MOV)
+    2. Adjust confidence threshold if needed (0.1-0.9)
+    3. Click "Process Video" and wait for completion
+    4. Download the processed video with drone annotations
+    
+    **Note**: Processing time depends on video length and resolution.
+    """
+    
+    with gr.Blocks(theme=gr.themes.Soft(), title=title) as interface:
+        gr.Markdown(f"# {title}")
+        gr.Markdown(description)
+        
+        with gr.Row():
+            with gr.Column():
+                input_video = gr.Video(
+                    label="📹 Upload Video",
+                    format="mp4"
+                )
+                confidence_slider = gr.Slider(
+                    minimum=0.1,
+                    maximum=0.9,
+                    value=0.5,
+                    step=0.1,
+                    label="🎯 Detection Confidence Threshold"
+                )
+                process_btn = gr.Button(
+                    "🚀 Process Video",
+                    variant="primary",
+                    size="lg"
+                )
+            
+            with gr.Column():
+                output_video = gr.Video(
+                    label="📽️ Processed Video",
+                    format="mp4"
+                )
+                status_text = gr.Textbox(
+                    label="📊 Processing Status",
+                    interactive=False
+                )
+        
+        process_btn.click(
+            fn=process_video_gradio,
+            inputs=[input_video, confidence_slider],
+            outputs=[output_video, status_text],
+            show_progress=True
+        )
+        
+        gr.Markdown(article)
+        
+        # Example section
+        gr.Markdown("### 📋 System Requirements")
+        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
+        """)
+    
+    return interface
+
+# Launch the application
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        quiet=False
+    )