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()