Create app.py

app.py

@@ -0,0 +1,126 @@
+import cv2
+import torch
+import gradio as gr
+from ultralytics import YOLO
+import time
+
+# Check if MPS (Metal Performance Shaders) is available, otherwise use CPU
+if torch.backends.mps.is_available() and torch.backends.mps.is_built():
+    device = torch.device('mps')
+    print("MPS available, using MPS.")
+else:
+    device = torch.device('cpu')
+    print("MPS not available, using CPU.")
+
+# Load the YOLOv8 model
+model = YOLO("yolov8n.pt").to(device)
+
+# Classes to count: 0 = person, 2 = car
+classes_to_count = [0, 2]  # person and car classes for counting
+
+# Initialize unique ID storage for each class
+unique_people_ids = set()
+unique_car_ids = set()
+
+def process_video(video_input):
+    global unique_people_ids, unique_car_ids
+    unique_people_ids = set()
+    unique_car_ids = set()
+
+    # Open the input video
+    cap = cv2.VideoCapture(video_input)
+    assert cap.isOpened(), "Error reading video file"
+
+    # Get video properties
+    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
+
+    # Set up video writer to store annotated video as frames are processed
+    output_frames = []
+    
+    frame_counter = 0
+    frame_skip = 5  # Process every 3rd frame
+
+    while cap.isOpened():
+        success, frame = cap.read()
+        if not success:
+            break
+
+        if frame_counter % frame_skip != 0:
+            frame_counter += 1
+            continue
+
+        # Calculate video timestamp based on frame number and FPS
+        video_time_elapsed = frame_counter / fps
+        video_timestamp = time.strftime('%H:%M:%S', time.gmtime(video_time_elapsed))
+
+        # Run object detection and tracking on the frame
+        results = model.track(frame, persist=True, device=device, classes=classes_to_count, verbose=False, conf=0.4)
+
+        # Initialize counters for current frame
+        people_count = 0
+        car_count = 0
+
+        # Process detections to track unique IDs
+        for det in results[0].boxes:
+            try:
+                object_id = int(det.id[0])
+            except:
+                pass
+            if object_id is None:
+                continue  # Skip objects without an ID
+
+            if det.cls == 0:  # person class
+                if object_id not in unique_people_ids:
+                    unique_people_ids.add(object_id)  # Add unique person ID
+                    people_count += 1
+            elif det.cls == 2:  # car class
+                if object_id not in unique_car_ids:
+                    unique_car_ids.add(object_id)  # Add unique car ID
+                    car_count += 1
+
+        # Annotate the frame with the current and total counts of unique objects
+        annotated_frame = results[0].plot()
+
+        # Display unique people and car count on the frame
+        cv2.putText(annotated_frame, f'Unique People: {len(unique_people_ids)} | Unique Cars: {len(unique_car_ids)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+        # Store the annotated frame
+        output_frames.append(annotated_frame)
+
+        # Increment frame counter
+        frame_counter += 1
+
+    cap.release()
+
+    # Return processed video frames
+    return output_frames
+
+def video_pipeline(video_file):
+    # Convert video into individual frames
+    output_frames = process_video(video_file)
+    
+    # Encode the frames back into a video
+    output_video_path = 'output.mp4'
+    h, w, _ = output_frames[0].shape
+    out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), 20, (w, h))
+    for frame in output_frames:
+        out.write(frame)
+    out.release()
+
+    return output_video_path
+
+# Gradio Interface
+title = "YOLOv8 Object Tracking with Unique ID Counting"
+description = "Upload a video to detect and count unique people and cars using YOLOv8."
+
+interface = gr.Interface(
+    fn=video_pipeline, 
+    inputs=gr.Video(label="Input Video"), 
+    outputs=gr.Video(label="Processed Video"),
+    title=title,
+    description=description,
+    live=True
+)
+
+# Launch Gradio interface
+interface.launch()