| 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 |
|
|
| |
| 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() |
|
|
| |
| self.writer.write(annotated_frame) |
|
|
| |
| 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() |
|
|
