File size: 2,257 Bytes
1ff6035 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
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()
|