Update realtime_detect.py

realtime_detect.py

@@ -1,74 +1,74 @@
-import tkinter as tk
-from tkinter import Label, Button
-import cv2
-from PIL import Image, ImageTk
-import threading
-from ultralytics import YOLO
-
-class WasteDetectorApp:
-    def __init__(self, window):
-        self.window = window
-        self.window.title("Waste Detection Camera App")
-
-        self.model = YOLO("best.pt")  # Replace with your model path
-        self.cap = None
-        self.running = False
-        self.writer = None
-
-        self.label = Label(window)
-        self.label.pack()
-
-        self.start_button = Button(window, text="Start Detection & Save", command=self.start_detection)
-        self.start_button.pack(pady=5)
-
-        self.stop_button = Button(window, text="Stop", command=self.stop_detection)
-        self.stop_button.pack(pady=5)
-
-    def start_detection(self):
-        if not self.running:
-            self.cap = cv2.VideoCapture(0)
-            self.running = True
-
-            # Set up video writer
-            fourcc = cv2.VideoWriter_fourcc(*'XVID')
-            fps = 20.0
-            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            self.writer = cv2.VideoWriter('output.avi', fourcc, fps, (width, height))
-
-            threading.Thread(target=self.detect_loop, daemon=True).start()
-
-    def stop_detection(self):
-        self.running = False
-        if self.cap:
-            self.cap.release()
-        if self.writer:
-            self.writer.release()
-        self.label.config(image='')
-
-    def detect_loop(self):
-        while self.running:
-            ret, frame = self.cap.read()
-            if not ret:
-                break
-
-            results = self.model.predict(frame, verbose=False)
-            annotated_frame = results[0].plot()
-
-            # Write to video file
-            self.writer.write(annotated_frame)
-
-            # Display to GUI
-            img_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
-            img_pil = Image.fromarray(img_rgb)
-            imgtk = ImageTk.PhotoImage(image=img_pil)
-            self.label.imgtk = imgtk
-            self.label.configure(image=imgtk)
-
-        self.cap.release()
-        self.writer.release()
-
-if __name__ == "__main__":
-    root = tk.Tk()
-    app = WasteDetectorApp(root)
-    root.mainloop()
+import tkinter as tk
+from tkinter import Label, Button
+import cv2
+from PIL import Image, ImageTk
+import threading
+from ultralytics import YOLO
+
+class WasteDetectorApp:
+    def __init__(self, window):
+        self.window = window
+        self.window.title("Waste Detection Camera App")
+
+        self.model = YOLO("best.pt")
+        self.cap = None
+        self.running = False
+        self.writer = None
+
+        self.label = Label(window)
+        self.label.pack()
+
+        self.start_button = Button(window, text="Start Detection & Save", command=self.start_detection)
+        self.start_button.pack(pady=5)
+
+        self.stop_button = Button(window, text="Stop", command=self.stop_detection)
+        self.stop_button.pack(pady=5)
+
+    def start_detection(self):
+        if not self.running:
+            self.cap = cv2.VideoCapture(0)
+            self.running = True
+
+            # Set up video writer
+            fourcc = cv2.VideoWriter_fourcc(*'XVID')
+            fps = 20.0
+            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            self.writer = cv2.VideoWriter('output.avi', fourcc, fps, (width, height))
+
+            threading.Thread(target=self.detect_loop, daemon=True).start()
+
+    def stop_detection(self):
+        self.running = False
+        if self.cap:
+            self.cap.release()
+        if self.writer:
+            self.writer.release()
+        self.label.config(image='')
+
+    def detect_loop(self):
+        while self.running:
+            ret, frame = self.cap.read()
+            if not ret:
+                break
+
+            results = self.model.predict(frame, verbose=False)
+            annotated_frame = results[0].plot()
+
+            # Write to video file
+            self.writer.write(annotated_frame)
+
+            # Display to GUI
+            img_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
+            img_pil = Image.fromarray(img_rgb)
+            imgtk = ImageTk.PhotoImage(image=img_pil)
+            self.label.imgtk = imgtk
+            self.label.configure(image=imgtk)
+
+        self.cap.release()
+        self.writer.release()
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    app = WasteDetectorApp(root)
+    root.mainloop()