Update app.py

app.py

@@ -1,98 +1,53 @@
-import cv2
-import tkinter as tk
-from tkinter import ttk, filedialog  # Import ttk for themed widgets
+import gradio as gr
 from ultralytics import YOLO
-import threading
-import time
-
-class PedestrianDetector:
-    def __init__(self, master):
-        self.master = master
-        master.title("Pedestrian Detection")
-
-        # Style for themed widgets (optional, but recommended)
-        style = ttk.Style()
-        style.theme_use("clam")  # Or "alt", "default", "classic" - experiment!
-
-        self.model = None
-        self.video_source = None
-
-        # GUI elements with themed widgets
-        self.load_model_button = ttk.Button(master, text="Load Model", command=self.load_model)
-        self.load_model_button.pack(pady=10, padx=10, fill=tk.X)  # Fill horizontally
-
-        self.select_video_button = ttk.Button(master, text="Select Video/Camera", command=self.select_video)
-        self.select_video_button.pack(pady=10, padx=10, fill=tk.X)
-
-        self.start_button = ttk.Button(master, text="Start Detection", command=self.start_detection, state=tk.DISABLED)
-        self.start_button.pack(pady=10, padx=10, fill=tk.X)
-
-        self.stop_button = ttk.Button(master, text="Stop Detection", command=self.stop_detection, state=tk.DISABLED)
-        self.stop_button.pack(pady=10, padx=10, fill=tk.X)
-
-        self.video_label = ttk.Label(master)  # Themed label
-        self.video_label.pack(pady=(0, 10), padx=10) #padding only at the bottom and left and right
-
-        self.detection_running = False
-        self.thread = None
-
-        # Status Label (for messages)
-        self.status_label = ttk.Label(master, text="", wraplength=300) #wraplength to wrap text
-        self.status_label.pack(pady=(0,10),padx=10)
-
-    # ... (rest of the methods: load_model, select_video, start/stop detection, detect_pedestrians)
-
-    def load_model(self):
-        try:
-            model_path = "yolov8n.pt"
-            self.model = YOLO(model_path)
-            self.load_model_button.config(state=tk.DISABLED)
-            self.status_label.config(text="Model loaded successfully.") #update status
-        except Exception as e:
-            self.status_label.config(text=f"Error loading model: {e}") #update status
-
-    def select_video(self):
-        try:
-            self.video_source = 0  # Webcam
-            if self.video_source is not None:
-                self.start_button.config(state=tk.NORMAL)
-                self.status_label.config(text=f"Video source selected: {self.video_source}") #update status
-        except Exception as e:
-            self.status_label.config(text=f"Error selecting video: {e}") #update status
-
-    def start_detection(self):
-        # ... (same as before)
-        self.status_label.config(text="Detection started.") #update status
-
-    def stop_detection(self):
-        # ... (same as before)
-        self.status_label.config(text="Detection stopped.") #update status
-
-    def detect_pedestrians(self):
-        cap = cv2.VideoCapture(self.video_source)
-
-        while self.detection_running:
+import cv2
+import numpy as np
+
+# Load the model (do this outside the function for efficiency)
+try:
+    model = YOLO("yolov8n.pt")  # Or your model path
+except Exception as e:
+    print(f"Error loading model: {e}")
+    exit()  # Exit if the model fails to load
+
+def detect_pedestrians(video):
+    try:
+        cap = cv2.VideoCapture(video)
+        frames = []
+        while True:
             ret, frame = cap.read()
             if not ret:
-                self.status_label.config(text="Video ended.") #update status
-                self.stop_detection()
                 break
 
-            # ... (detection logic - same as before)
+            results = model(frame)
 
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            frame_tk = cv2.resize(frame_rgb, (640, 480))
-            frame_tk = tk.PhotoImage(image=frame_tk)
-            self.video_label.config(image=frame_tk)
-            self.video_label.image = frame_tk
+            for result in results:
+                boxes = result.boxes
+                for box in boxes:
+                    if result.names[int(box.cls)] == 'person':
+                        x1, y1, x2, y2 = map(int, box.xyxy[0])
+                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                        cv2.putText(frame, 'Person', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+            frames.append(frame)
+        cap.release()
 
-            if cv2.waitKey(1) & 0xFF == ord('q'):
-                self.detection_running = False
+        if not frames:  # Handle cases where no frames are read (e.g., corrupted video)
+            return np.zeros((480, 640, 3), dtype=np.uint8)  # Return a blank image
+
+        return frames[-1]  # Return the last frame with detections
+
+    except Exception as e:
+        print(f"Error in detection: {e}")
+        return f"Error: {e}"  # Return the error message as a string
 
-        cap.release()
-        cv2.destroyAllWindows()
 
+iface = gr.Interface(
+    fn=detect_pedestrians,
+    inputs=gr.Video(source="upload"),  # Use gr.Video for video input
+    outputs=gr.Image(),  # Use gr.Image for image output
+    title="Pedestrian Detection",
+    description="Upload a video to detect pedestrians.",
+    allow_flagging="never",  # Removes flag button
+)
 
-root = tk.Tk()
-detector = PedestrianDetector(root)
-root.mainloop()
+iface.launch()