Create app.py

app.py

@@ -0,0 +1,139 @@
+import streamlit as st
+import cv2
+import numpy as np
+import tempfile
+import os
+from pathlib import Path
+
+# ---------------------------
+# Helper: Load YOLO Model
+# ---------------------------
+@st.cache_resource
+def load_yolo():
+    # Load class labels
+    labelsPath = os.path.join("yolo", "coco.names")
+    with open(labelsPath, "r") as f:
+        classes = f.read().strip().split("\n")
+    
+    # Load YOLO model configuration and weights
+    net = cv2.dnn.readNet(os.path.join("yolo", "yolov4.weights"), os.path.join("yolo", "yolov4.cfg"))
+    
+    # Uncomment these lines if you have GPU support
+    # net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+    # net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+    
+    # Get output layer names
+    layer_names = net.getLayerNames()
+    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers().flatten()]
+    
+    return net, output_layers, classes
+
+net, output_layers, classes = load_yolo()
+
+# ---------------------------
+# Helper: Process Video
+# ---------------------------
+def process_video(video_path, max_frames=100):
+    cap = cv2.VideoCapture(video_path)
+    processed_frames = []
+    frame_count = 0
+
+    while cap.isOpened() and frame_count < max_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        height, width = frame.shape[:2]
+        
+        # Create blob from image
+        blob = cv2.dnn.blobFromImage(frame, 1/255.0, (416, 416), swapRB=True, crop=False)
+        net.setInput(blob)
+        outputs = net.forward(output_layers)
+        
+        boxes = []
+        confidences = []
+        class_ids = []
+        
+        # Loop over detections
+        for output in outputs:
+            for detection in output:
+                scores = detection[5:]
+                class_id = np.argmax(scores)
+                confidence = scores[class_id]
+                if confidence > 0.5:
+                    center_x = int(detection[0] * width)
+                    center_y = int(detection[1] * height)
+                    w = int(detection[2] * width)
+                    h = int(detection[3] * height)
+                    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)
+        
+        # Non-max suppression to remove duplicates
+        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
+        object_count = {}
+        
+        if len(indexes) > 0:
+            for i in indexes.flatten():
+                x, y, w, h = boxes[i]
+                label = str(classes[class_ids[i]])
+                confidence = confidences[i]
+                color = (0, 255, 0)
+                cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
+                cv2.putText(frame, f'{label} {int(confidence * 100)}%', (x, y - 10),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+                object_count[label] = object_count.get(label, 0) + 1
+        
+        # Display object counts on frame
+        y_offset = 30
+        for label, count in object_count.items():
+            cv2.putText(frame, f'{label}: {count}', (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,255), 2)
+            y_offset += 30
+        
+        processed_frames.append(frame)
+        frame_count += 1
+
+    cap.release()
+    return processed_frames
+
+# ---------------------------
+# Streamlit UI
+# ---------------------------
+st.title("Real-Time Object Detection and Counting")
+st.write("Upload a video file to run object detection using YOLOv4.")
+
+# Video file uploader
+uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "mov", "avi"])
+
+if uploaded_file is not None:
+    # Save uploaded file to a temporary file
+    tfile = tempfile.NamedTemporaryFile(delete=False)
+    tfile.write(uploaded_file.read())
+    
+    st.video(uploaded_file)  # Show original video
+    
+    if st.button("Run Object Detection"):
+        st.write("Processing video...")
+        processed_frames = process_video(tfile.name, max_frames=100)
+        
+        # Create a directory for output frames (optional)
+        output_dir = Path("output_frames")
+        output_dir.mkdir(exist_ok=True)
+        frame_paths = []
+        
+        # Save processed frames as images
+        for i, frame in enumerate(processed_frames):
+            frame_path = output_dir / f"frame_{i:03d}.jpg"
+            cv2.imwrite(str(frame_path), frame)
+            frame_paths.append(str(frame_path))
+        
+        st.success("Processing complete!")
+        
+        # Display processed frames as a gallery (or create a video if needed)
+        st.write("Processed Frames:")
+        for frame_path in frame_paths:
+            st.image(frame_path, channels="BGR")
+        
+        # Optionally, you could create a video file from frames and offer a download link.