Update app.py

app.py

@@ -1360,9 +1360,107 @@ def process_video_file(video_file, analysis_types, processing_mode="Hybrid (Goog
                                 else:
                                     detection_stats["objects"]["person"] = 1
                     else:
-                        # Using YOLO or another DNN-based detector code would go here
-                        pass
+                        # Using YOLO or another DNN-based detector
+                        blob = cv2.dnn.blobFromImage(frame, 1/255.0, (416, 416), swapRB=True, crop=False)
+                        opencv_detector.setInput(blob)
+                        
+                        # Get output layer names
+                        layer_names = opencv_detector.getLayerNames()
+                        output_layers = []
+                        
+                        # Handle different OpenCV versions
+                        try:
+                            if cv2.__version__.startswith('4'):
+                                # OpenCV 4.x
+                                output_layers = [layer_names[i - 1] for i in opencv_detector.getUnconnectedOutLayers()]
+                            else:
+                                # OpenCV 3.x
+                                output_layers = [layer_names[i[0] - 1] for i in opencv_detector.getUnconnectedOutLayers()]
+                        except:
+                            # Fallback method
+                            unconnected_layers = opencv_detector.getUnconnectedOutLayers()
+                            if isinstance(unconnected_layers[0], list) or isinstance(unconnected_layers[0], tuple):
+                                output_layers = [layer_names[i[0] - 1] for i in unconnected_layers]
+                            else:
+                                output_layers = [layer_names[i - 1] for i in unconnected_layers]
                         
+                        outputs = opencv_detector.forward(output_layers)
+                        
+                        # Process detections
+                        class_ids = []
+                        confidences = []
+                        boxes = []
+                        
+                        # Define COCO class names
+                        class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
+                                      "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
+                                      "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
+                                      "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
+                                      "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
+                                      "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
+                                      "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
+                                      "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
+                                      "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
+                                      "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"]
+                        
+                        # Process each detection
+                        for output in outputs:
+                            for detection in output:
+                                scores = detection[5:]
+                                class_id = np.argmax(scores)
+                                confidence = scores[class_id]
+                                
+                                if confidence > confidence_threshold:
+                                    # Object detected
+                                    center_x = int(detection[0] * frame.shape[1])
+                                    center_y = int(detection[1] * frame.shape[0])
+                                    w = int(detection[2] * frame.shape[1])
+                                    h = int(detection[3] * frame.shape[0])
+                                    
+                                    # Rectangle coordinates
+                                    x = int(center_x - w / 2)
+                                    y = int(center_y - h / 2)
+                                    
+                                    boxes.append([x, y, w, h])
+                                    confidences.append(float(confidence))
+                                    class_ids.append(class_id)
+                        
+                        # Apply non-maximum suppression
+                        indices = cv2.dnn.NMSBoxes(boxes, confidences, confidence_threshold, 0.4)
+                        
+                        # Draw the detections
+                        if len(indices) > 0:
+                            for i in indices:
+                                if isinstance(i, (list, tuple)):  # Handle different OpenCV versions
+                                    i = i[0]
+                                
+                                box = boxes[i]
+                                x, y, w, h = box
+                                
+                                # Get class name
+                                class_id = class_ids[i]
+                                label = f"{class_names[class_id]}: {int(confidences[i] * 100)}%"
+                                
+                                # Different colors for different classes
+                                color = (0, 255, 0)  # Default color
+                                
+                                # Draw rectangle and label
+                                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
+                                cv2.putText(frame, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+                                
+                                # Add to object trackers for future frames
+                                object_name = class_names[class_id]
+                                object_trackers[f"{object_name}_{i}"] = {
+                                    "bbox": (x, y, w, h),
+                                    "last_seen": frame_count,
+                                    "score": confidences[i]
+                                }
+                                
+                                # Update detection stats
+                                if object_name in detection_stats["objects"]:
+                                    detection_stats["objects"][object_name] += 1
+                                else:
+                                    detection_stats["objects"][object_name] = 1
                 except Exception as e:
                     cv2.putText(frame, f"OpenCV Error: {str(e)}", (10, 110), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
@@ -2139,6 +2237,49 @@ def main():
             # OpenCV parameters
             st.sidebar.markdown("#### OpenCV Parameters")
             
+            # Add YOLO model download option if models aren't found
+            models_dir = os.path.join(os.path.dirname(__file__), "models")
+            weights_path = os.path.join(models_dir, "yolov3.weights")
+            config_path = os.path.join(models_dir, "yolov3.cfg")
+            
+            if not os.path.exists(models_dir):
+                os.makedirs(models_dir, exist_ok=True)
+            
+            if not (os.path.exists(weights_path) and os.path.exists(config_path)):
+                st.sidebar.warning("⚠️ YOLO models not found. Using basic people detector.")
+                
+                if st.sidebar.button("Download YOLO Models"):
+                    with st.sidebar.spinner("Downloading YOLO models..."):
+                        # Ensure models directory exists
+                        os.makedirs(models_dir, exist_ok=True)
+                        
+                        # Download YOLOv3 config
+                        try:
+                            import urllib.request
+                            
+                            # Download config file
+                            if not os.path.exists(config_path):
+                                urllib.request.urlretrieve(
+                                    "https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg", 
+                                    config_path
+                                )
+                                st.sidebar.success("Downloaded YOLOv3 config file")
+                            
+                            # Download weights file (this is large - about 240MB)
+                            if not os.path.exists(weights_path):
+                                urllib.request.urlretrieve(
+                                    "https://pjreddie.com/media/files/yolov3.weights", 
+                                    weights_path
+                                )
+                                st.sidebar.success("Downloaded YOLOv3 weights file")
+                            
+                            st.sidebar.success("YOLO models downloaded successfully! Please refresh the page.")
+                        except Exception as e:
+                            st.sidebar.error(f"Error downloading YOLO models: {str(e)}")
+                            st.sidebar.info("You can manually download the models from: https://pjreddie.com/darknet/yolo/")
+            else:
+                st.sidebar.success("✅ YOLO models found. Using advanced object detection.")
+            
             track_update_frames = st.sidebar.slider(
                 "Update OpenCV tracking every N frames", 
                 min_value=1, 
@@ -2255,6 +2396,32 @@ def main():
                                     "blend_results": blend_results
                                 })
                             
+                            # Add to the OpenCV parameters section:
+                            with st.sidebar.expander("YOLO Class Filters"):
+                                # Allow users to select which object classes to detect
+                                st.markdown("Select which objects to detect:")
+                                
+                                # Create a multiselect with common categories
+                                selected_categories = st.multiselect(
+                                    "Object Categories",
+                                    ["People", "Vehicles", "Animals", "Indoor Objects", "Sports Equipment", "Food", "All"],
+                                    default=["People", "Vehicles"]
+                                )
+                                
+                                # Map categories to actual YOLO classes
+                                yolo_classes = []
+                                if "People" in selected_categories:
+                                    yolo_classes.extend(["person"])
+                                if "Vehicles" in selected_categories:
+                                    yolo_classes.extend(["bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck"])
+                                if "Animals" in selected_categories:
+                                    yolo_classes.extend(["bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe"])
+                                if "All" in selected_categories:
+                                    yolo_classes = None  # Detect all classes
+                                
+                                # Pass this to your processing function in the processing_params
+                                processing_params["yolo_classes"] = yolo_classes
+                            
                             # Process the video with the parameters
                             processed_video, results = process_video_file(uploaded_file, analysis_types, **processing_params)