Update app.py

app.py

@@ -26,7 +26,7 @@ class DroneDetectionPipeline:
         self.tracker = None
         self.tracking_active = False
         self.last_detection_bbox = None
-        self.confidence_threshold = 0.5
+        self.confidence_threshold = 0.55
         
         # Load model
         self._load_model()
@@ -190,43 +190,58 @@ class DroneDetectionPipeline:
                     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]))
-                    
+                # Get all detections (including low confidence ones for tracking)
+                all_detections = self.model(frame, verbose=False, conf=0.1)  # Low threshold to get all detections
+                high_conf_detections = []
+                low_conf_detections = []
+    
+                for result in all_detections:
+                    if result.boxes is not None:
+                        for box in result.boxes:
+                            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+                            confidence = float(box.conf[0].cpu().numpy())
+                            
+                            if confidence >= 0.55:
+                                high_conf_detections.append((x1, y1, x2, y2, confidence))
+                            elif confidence >= 0.1:  # Lower threshold for tracking
+                                low_conf_detections.append((x1, y1, x2, y2, confidence))
+    
+                if high_conf_detections:
+                    # Use high confidence detection - normal detection display
+                    largest_detection = max(high_conf_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)
+                    detection_bbox = (x1, y1, x2-x1, y2-y1)
                     
-                    # Reinitialize tracker with new detection
+                    frame_processed = self._draw_detection(frame_processed, (x1, y1, x2, y2), None, False)
                     self._initialize_tracker(frame, detection_bbox)
                     detection_count += 1
                     
+                elif low_conf_detections and self.last_detection_bbox is not None:
+                    # Use low confidence detection for tracking only
+                    largest_low_conf = max(low_conf_detections, key=lambda d: (d[2]-d[0]) * (d[3]-d[1]))
+                    x1, y1, x2, y2, conf = largest_low_conf
+                    tracking_bbox = (x1, y1, x2-x1, y2-y1)
+                    
+                    # Update last known position and show as tracking
+                    self.last_detection_bbox = tracking_bbox
+                    frame_processed = self._draw_detection(frame_processed, (x1, y1, x2, y2), None, True)
+                    tracking_count += 1
+                    
                 elif self.tracking_active and self.tracker is not None:
-                    # No detections but tracker is active
+                    # Use CSRT tracker when no detections
                     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)
+                        frame_processed = self._draw_detection(frame_processed, (x, y, x+w, y+h), None, True)
                         tracking_count += 1
                     else:
-                        # Tracking failed
+                        # Tracking failed - use last known position
                         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)
+                            frame_processed = self._draw_detection(frame_processed, (x, y, x+w, y+h), None, True)
                         self.tracking_active = False
                 
                 # Write processed frame
@@ -253,7 +268,7 @@ class DroneDetectionPipeline:
 # Initialize pipeline
 pipeline = DroneDetectionPipeline()
 
-def process_video_gradio(input_video, confidence_threshold):
+def process_video_gradio(input_video):
     """Gradio interface function with enhanced error handling"""
     if input_video is None:
         return None, "❌ Please upload a video file"
@@ -264,7 +279,7 @@ def process_video_gradio(input_video, confidence_threshold):
             return None, "❌ Model not loaded. Please check that 'best.pt' is uploaded to the space."
         
         # Update confidence threshold
-        pipeline.confidence_threshold = confidence_threshold
+        # pipeline.confidence_threshold = confidence_threshold
         
         # Add file size check
         file_size = os.path.getsize(input_video)
@@ -332,13 +347,13 @@ def create_interface():
                     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"
-                )
+                # 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",
@@ -357,7 +372,7 @@ def create_interface():
         
         process_btn.click(
             fn=process_video_gradio,
-            inputs=[input_video, confidence_slider],
+            inputs=[input_video],  # Only input_video
             outputs=[output_video, status_text],
             show_progress=True
         )