Untitled

import customtkinter as ctk
import tkinter as tk
from PIL import Image
import os
import numpy as np
from datetime import datetime
import cv2
import tflite_runtime.interpreter as tflite
from pathlib import Path
import math

class CircularProgress(ctk.CTkCanvas):
    def __init__(self, parent, size=100, fg_color="#2CC985", bg_color="#2B2B2B", text="", text_color="white", **kwargs):
        super().__init__(parent, width=size, height=size, bg=parent.cget("fg_color"), highlightthickness=0, **kwargs)
        self.size = size
        self.fg_color = fg_color
        self.bg_color = bg_color
        self.text = text
        self.text_color = text_color
        self.percentage = 0
        self.warning_threshold = 90
        self.draw()

    def draw(self):
        self.delete("progress")
        self.delete("text")
        
        # Background circle
        padding = 2
        self.create_oval(padding, padding, self.size-padding, self.size-padding, 
                        fill=self.bg_color, width=0, tags="progress")
        
        if self.percentage > 0:
            start_angle = 90
            extent = -(self.percentage/100 * 360)
            
            # Calculate coordinates for arc
            x0 = padding
            y0 = padding
            x1 = self.size - padding
            y1 = self.size - padding
            
            # Draw progress arc
            arc_color = "#FF5555" if self.percentage >= self.warning_threshold else self.fg_color
            self.create_arc(x0, y0, x1, y1, 
                          start=start_angle, extent=extent,
                          fill="", width=3, 
                          style="arc", outline=arc_color,
                          tags="progress")
        
        # Add text
        self.create_text(self.size/2, self.size/2, 
                        text=self.text, fill=self.text_color,
                        font=("Helvetica", int(self.size/8), "bold"),
                        tags="text")

    def set_percentage(self, percentage):
        self.percentage = min(100, max(0, percentage))
        self.draw()

class WasteClassificationGUI:
    def __init__(self):
        self.setup_window()
        self.setup_model()
        self.setup_camera()
        self.create_folders()
        self.setup_styles()
        self.create_main_screen()
        self.current_screen = "main"
        self.timer = None
        
    def setup_window(self):
        self.root = ctk.CTk()
        self.root.geometry("1024x600")
        self.root.title("Waste Classification")
        self.root.attributes('-fullscreen', True)
        ctk.set_appearance_mode("dark")
        
    def setup_model(self):
        MODEL_PATH = "ei-v2-transfer-learning-tensorflow-lite-float32-model.lite"
        self.interpreter = tflite.Interpreter(model_path=MODEL_PATH)
        self.interpreter.allocate_tensors()
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()
        self.height = self.input_details[0]['shape'][1]
        self.width = self.input_details[0]['shape'][2]
        self.class_names = ["Biomüll", "Gelber Sack", "Papier", "Restmüll"]
        
    def setup_camera(self):
        from picamera2 import Picamera2
        self.camera = Picamera2()
        self.camera.preview_configuration.main.size = (1920, 1440)
        self.camera.preview_configuration.main.format = "RGB888"
        self.camera.configure("preview")
        self.camera.set_controls({"AfMode": 2})
        self.camera.start()
        
    def create_folders(self):
        self.base_dir = Path("captured_images")
        for class_name in self.class_names:
            (self.base_dir / class_name).mkdir(parents=True, exist_ok=True)
            
    def setup_styles(self):
        self.colors = {
            "Biomüll": "#8B4513",
            "Gelber Sack": "#FFD700",
            "Papier": "#4169E1",
            "Restmüll": "#808080",
            "background": "#1E1E1E",
            "timer": "#FFFFFF",
            "text": "#FFFFFF"
        }
        
    def create_timer_bar(self, parent, duration=5000):
        timer_frame = ctk.CTkFrame(parent, fg_color="transparent")
        timer_frame.pack(fill="x", padx=10, pady=5)
        
        self.timer_bar = ctk.CTkProgressBar(timer_frame)
        self.timer_bar.pack(fill="x", padx=5)
        self.timer_bar.set(1.0)
        
        def update_timer(remaining):
            if remaining <= 0:
                self.show_main_screen()
                return
            
            self.timer_bar.set(remaining/duration)
            self.root.after(50, update_timer, remaining - 50)
            
        update_timer(duration)
        
    def create_main_screen(self):
        self.main_frame = ctk.CTkFrame(self.root, fg_color=self.colors["background"])
        self.main_frame.pack(fill="both", expand=True)
        
        # Circles container
        circles_frame = ctk.CTkFrame(self.main_frame, fg_color="transparent")
        circles_frame.pack(expand=True)
        
        self.circles = {}
        circle_size = 150
        
        for i, class_name in enumerate(self.class_names):
            circle_frame = ctk.CTkFrame(circles_frame, fg_color="transparent")
            circle_frame.pack(side="left", padx=20)
            
            circle = CircularProgress(circle_frame, size=circle_size,
                                    fg_color=self.colors[class_name],
                                    bg_color="#2B2B2B",
                                    text=class_name)
            circle.pack(pady=10)
            
            slider = ctk.CTkSlider(circle_frame, from_=0, to=100,
                                 command=lambda val, c=circle: c.set_percentage(val))
            slider.pack(pady=5)
            slider.set(0)
            
            self.circles[class_name] = {"circle": circle, "slider": slider}
        
        # Bottom buttons
        button_frame = ctk.CTkFrame(self.main_frame, fg_color="transparent")
        button_frame.pack(side="bottom", fill="x", padx=20, pady=20)
        
        scan_button = ctk.CTkButton(button_frame, text="Scan",
                                  command=self.start_scan,
                                  width=100)
        scan_button.pack(side="left")
        
    def start_scan(self):
        image = self.camera.capture_array()
        prediction, confidence = self.classify_image(image)
        self.show_result_screen(prediction, confidence, image)
        
    def classify_image(self, image):
        processed_image = self.preprocess_image(image)
        self.interpreter.set_tensor(self.input_details[0]['index'], processed_image)
        self.interpreter.invoke()
        output = self.interpreter.get_tensor(self.output_details[0]['index'])
        
        max_index = np.argmax(output[0])
        confidence = output[0][max_index] * 100
        prediction = self.class_names[max_index]
        
        return prediction, confidence
        
    def preprocess_image(self, image):
        resized = cv2.resize(image, (self.width, self.height))
        normalized = resized / 255.0
        return np.expand_dims(normalized, axis=0).astype(np.float32)
        
    def show_result_screen(self, prediction, confidence, image):
        self.clear_current_screen()
        
        result_frame = ctk.CTkFrame(self.root, fg_color=self.colors["background"])
        result_frame.pack(fill="both", expand=True)
        
        self.create_timer_bar(result_frame)
        
        heading = ctk.CTkLabel(result_frame, 
                             text=f"Detected: {prediction}",
                             font=("Helvetica", 24, "bold"))
        heading.pack(pady=20)
        
        confidence_label = ctk.CTkLabel(result_frame,
                                      text=f"Confidence: {confidence:.1f}%",
                                      font=("Helvetica", 18))
        confidence_label.pack(pady=10)
        
        buttons_frame = ctk.CTkFrame(result_frame, fg_color="transparent")
        buttons_frame.pack(pady=20)
        
        ctk.CTkButton(buttons_frame, 
                     text="✓",
                     command=lambda: self.handle_feedback(True, prediction, image),
                     width=60).pack(side="left", padx=10)
                     
        ctk.CTkButton(buttons_frame,
                     text="✗",
                     command=lambda: self.handle_feedback(False, prediction, image),
                     width=60).pack(side="left", padx=10)
                     
    def handle_feedback(self, is_correct, prediction, image):
        if is_correct:
            self.save_image(image, prediction)
            self.show_thank_you_screen()
        else:
            self.show_correction_screen(image)
            
    def show_correction_screen(self, image):
        self.clear_current_screen()
        
        correction_frame = ctk.CTkFrame(self.root, fg_color=self.colors["background"])
        correction_frame.pack(fill="both", expand=True)
        
        self.create_timer_bar(correction_frame)
        
        heading = ctk.CTkLabel(correction_frame,
                             text="Select correct category:",
                             font=("Helvetica", 24, "bold"))
        heading.pack(pady=20)
        
        buttons_frame = ctk.CTkFrame(correction_frame, fg_color="transparent")
        buttons_frame.pack(expand=True)
        
        for class_name in self.class_names:
            btn = ctk.CTkButton(buttons_frame,
                              text=class_name,
                              fg_color=self.colors[class_name],
                              command=lambda c=class_name: self.handle_correction(c, image))
            btn.pack(pady=10)
            
    def handle_correction(self, correct_class, image):
        self.save_image(image, correct_class)
        self.show_thank_you_screen()
        
    def show_thank_you_screen(self):
        self.clear_current_screen()
        
        thank_you_frame = ctk.CTkFrame(self.root, fg_color=self.colors["background"])
        thank_you_frame.pack(fill="both", expand=True)
        
        self.create_timer_bar(thank_you_frame, duration=3000)
        
        heading = ctk.CTkLabel(thank_you_frame,
                             text="Thank you for your feedback!",
                             font=("Helvetica", 24, "bold"))
        heading.pack(expand=True)
        
    def save_image(self, image, class_name):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        save_path = self.base_dir / class_name / f"capture_{timestamp}.jpg"
        cv2.imwrite(str(save_path), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
        
    def clear_current_screen(self):
        for widget in self.root.winfo_children():
            widget.destroy()
            
    def show_main_screen(self):
        self.clear_current_screen()
        self.create_main_screen()
        
    def run(self):
        self.root.mainloop()

if __name__ == "__main__":
    app = WasteClassificationGUI()
    app.run()