Create app.py

app.py

@@ -0,0 +1,96 @@
+import gradio as gr
+import cv2
+from ultralytics import YOLO, solutions
+
+# Initialize the YOLO model
+model = YOLO("yolov8s.pt")
+
+def process_video(video_path, analytics_type):
+    cap = cv2.VideoCapture(video_path)
+    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))
+    output_filename = f"{analytics_type}_output.avi"
+    out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))
+    
+    # Set up analytics based on the selected type
+    analytics = solutions.Analytics(
+        type=analytics_type,
+        writer=out,
+        im0_shape=(w, h),
+        view_img=False
+    )
+    
+    clswise_count = {}
+    frame_count = 0
+
+    while cap.isOpened():
+        success, frame = cap.read()
+        if success:
+            frame_count += 1
+            results = model.track(frame, persist=True, verbose=False)
+
+            if results[0].boxes.id is not None:
+                boxes = results[0].boxes.xyxy.cpu()
+                clss = results[0].boxes.cls.cpu().tolist()
+
+                for box, cls in zip(boxes, clss):
+                    if model.names[int(cls)] in clswise_count:
+                        clswise_count[model.names[int(cls)]] += 1
+                    else:
+                        clswise_count[model.names[int(cls)]] = 1
+
+            # Update analytics based on type
+            if analytics_type == "line":
+                total_counts = sum(clswise_count.values())
+                analytics.update_line(frame_count, total_counts)
+            elif analytics_type == "multiple_line":
+                analytics.update_multiple_lines(clswise_count, list(clswise_count.keys()), frame_count)
+            elif analytics_type == "pie":
+                analytics.update_pie(clswise_count)
+            elif analytics_type == "area":
+                analytics.update_area(frame_count, clswise_count)
+
+            clswise_count = {}  # Reset for next frame
+
+        else:
+            break
+
+    cap.release()
+    out.release()
+
+    return output_filename  # Return the output video file
+
+
+def gradio_app(video, analytics_type):
+    # Save uploaded video locally
+    video_path = video.name
+    output_video = process_video(video_path, analytics_type)
+    
+    # Return processed video for display
+    return output_video
+
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# YOLO Video Processing App")
+    
+    with gr.Row():
+        video_input = gr.Video(label="Upload Video", type="file")
+        analytics_dropdown = gr.Dropdown(
+            ["line", "multiple_line", "pie", "area"], 
+            label="Select Analytics Type", 
+            value="line"
+        )
+    
+    output_video = gr.Video(label="Processed Output")
+    
+    # Button to start processing
+    submit_btn = gr.Button("Process Video")
+
+    # Define the output when the button is clicked
+    submit_btn.click(gradio_app, inputs=[video_input, analytics_dropdown], outputs=output_video)
+
+# Launch the Gradio app
+demo.launch()