Untitled

import customtkinter as ctk
import tkinter as tk
from PIL import Image, ImageTk
import numpy as np
import cv2
import math
import os
from datetime import datetime
import tflite_runtime.interpreter as tflite
from picamera2 import Picamera2
import threading
import time

class CircularProgress(ctk.CTkCanvas):
    def __init__(self, parent, name, color, size=100, **kwargs):
        super().__init__(parent, width=size, height=size, **kwargs)
        self.size = size
        self.color = color
        self.name = name
        self.configure(bg=parent.cget('bg'))
        self.value = 0
        
        # Create background circle
        padding = 2
        self.create_oval(padding, padding, size-padding, size-padding, 
                        width=2, outline='gray80')
        
        # Create text
        self.create_text(size/2, size/2, text=name, fill='white', 
                        font=('Arial', int(size/8)))
        
        self.arc = None
        self.warning_glow = None
        self.draw_progress(0)
    
    def draw_progress(self, percentage):
        self.value = percentage
        angle = 360 * (percentage / 100)
        
        # Remove existing arc
        if self.arc is not None:
            self.delete(self.arc)
        if self.warning_glow is not None:
            self.delete(self.warning_glow)
            
        # Calculate arc coordinates
        start_angle = -90
        end_angle = start_angle + angle
        
        # Add warning glow for >90%
        if percentage > 90:
            self.warning_glow = self.create_arc(4, 4, self.size-4, self.size-4,
                                              start=start_angle, extent=angle,
                                              outline='red', width=4, style='arc')
        
        # Draw progress arc
        self.arc = self.create_arc(4, 4, self.size-4, self.size-4,
                                 start=start_angle, extent=angle,
                                 outline=self.color, width=3, style='arc')

class Timer:
    def __init__(self, duration, callback):
        self.duration = duration
        self.callback = callback
        self.remaining = duration
        self.running = False
        self.thread = None
    
    def start(self):
        self.running = True
        self.remaining = self.duration
        self.thread = threading.Thread(target=self._run)
        self.thread.daemon = True
        self.thread.start()
    
    def _run(self):
        while self.running and self.remaining > 0:
            time.sleep(0.1)
            self.remaining -= 0.1
        if self.running:
            self.callback()
    
    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join()

class WasteClassificationGUI:
    def __init__(self):
        # Define categories first
        self.categories = {
            "Biomüll": "brown",
            "Gelber Sack": "yellow",
            "Papier": "blue",
            "Restmüll": "gray"
        }
        
        # Create image storage folders
        for category in self.categories:
            os.makedirs(f"captured_images/{category}", exist_ok=True)
        
        # Setup other components
        self.setup_window()
        self.setup_camera()
        self.setup_model()
        self.create_main_view()
    
    def setup_window(self):
        self.root = ctk.CTk()
        self.root.geometry("1024x600")
        self.root.title("Waste Classification")
        ctk.set_appearance_mode("dark")
        self.root.attributes('-fullscreen', True)
        
        # Create main container
        self.main_container = ctk.CTkFrame(self.root)
        self.main_container.pack(fill='both', expand=True)
    
    def setup_camera(self):
        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 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()
    
    def create_main_view(self):
        # Clear current view
        for widget in self.main_container.winfo_children():
            widget.destroy()
        
        # Create header with progress circles
        header = ctk.CTkFrame(self.main_container, fg_color="transparent")
        header.pack(fill='x', padx=20, pady=20)
        
        # Configure grid columns to be evenly spaced
        for i in range(4):
            header.grid_columnconfigure(i, weight=1)
        
        # Create circular progress indicators
        self.progress_widgets = {}
        for i, (category, color) in enumerate(self.categories.items()):
            progress = CircularProgress(header, category, color, size=120)
            progress.grid(row=0, column=i, padx=20)
            self.progress_widgets[category] = progress
            
            # Create slider for each progress
            slider = ctk.CTkSlider(self.main_container, 
                                 from_=0, to=100,
                                 command=lambda val, cat=category: 
                                     self.update_progress(cat, val))
            slider.pack(pady=5)
        
        # Create spacer to push content to top and bottom
        spacer = ctk.CTkFrame(self.main_container, fg_color="transparent", height=200)
        spacer.pack(expand=True)
        
        # Create scan button container at bottom
        scan_frame = ctk.CTkFrame(self.main_container, fg_color="transparent")
        scan_frame.pack(side='bottom', fill='x', padx=20, pady=20)
        
        # Create scan button on the left
        scan_button = ctk.CTkButton(scan_frame,
                                  text="Scan",
                                  command=self.capture_and_classify,
                                  width=120, height=40)
        scan_button.pack(side='left')
    
    def update_progress(self, category, value):
        self.progress_widgets[category].draw_progress(value)
    
    def create_timer_bar(self, parent, duration):
        timer_frame = ctk.CTkFrame(parent)
        timer_frame.pack(fill='x', pady=5)
        
        self.timer_bar = ctk.CTkProgressBar(timer_frame)
        self.timer_bar.pack(fill='x', padx=20)
        self.timer_bar.set(1)
        
        def update_timer():
            remaining = duration
            while remaining > 0 and hasattr(self, 'timer_bar'):
                self.timer_bar.set(remaining/duration)
                remaining -= 0.1
                time.sleep(0.1)
        
        threading.Thread(target=update_timer, daemon=True).start()
    
    def show_classification_result(self, class_name, confidence):
        result_window = ctk.CTkToplevel(self.root)
        result_window.geometry("1024x600")
        result_window.attributes('-fullscreen', True)
        
        # Create timer bar
        self.create_timer_bar(result_window, 5)
        
        # Display result
        result_frame = ctk.CTkFrame(result_window)
        result_frame.pack(expand=True)
        
        icon_label = ctk.CTkLabel(result_frame,
                                text="🗑️",
                                font=("Arial", 48))
        icon_label.pack(pady=20)
        
        result_label = ctk.CTkLabel(result_frame,
                                  text=f"{class_name}\n{confidence:.1f}%",
                                  font=("Arial", 24))
        result_label.pack(pady=20)
        
        # Feedback buttons
        feedback_frame = ctk.CTkFrame(result_frame)
        feedback_frame.pack(pady=20)
        
        ctk.CTkButton(feedback_frame,
                     text="👍",
                     command=lambda: self.handle_feedback(True, class_name, result_window),
                     width=80).pack(side='left', padx=10)
        
        ctk.CTkButton(feedback_frame,
                     text="👎",
                     command=lambda: self.handle_feedback(False, class_name, result_window),
                     width=80).pack(side='left', padx=10)
        
        # Auto-close timer
        Timer(5, result_window.destroy).start()
    
    def handle_feedback(self, positive, predicted_class, window):
        window.destroy()
        
        if not positive:
            self.show_correction_dialog(predicted_class)
        else:
            self.show_thank_you()
    
    def show_correction_dialog(self, predicted_class):
        dialog = ctk.CTkToplevel(self.root)
        dialog.geometry("1024x600")
        dialog.attributes('-fullscreen', True)
        
        # Create timer bar
        self.create_timer_bar(dialog, 5)
        
        # Correction options
        options_frame = ctk.CTkFrame(dialog)
        options_frame.pack(expand=True)
        
        label = ctk.CTkLabel(options_frame,
                           text="Please select the correct category:",
                           font=("Arial", 20))
        label.pack(pady=20)
        
        for category in self.categories:
            ctk.CTkButton(options_frame,
                         text=category,
                         command=lambda cat=category: self.handle_correction(cat, dialog),
                         width=200).pack(pady=10)
        
        # Auto-close timer
        Timer(5, dialog.destroy).start()
    
    def show_thank_you(self):
        thank_you = ctk.CTkToplevel(self.root)
        thank_you.geometry("1024x600")
        thank_you.attributes('-fullscreen', True)
        
        # Create timer bar
        self.create_timer_bar(thank_you, 3)
        
        # Thank you message
        message = ctk.CTkLabel(thank_you,
                             text="Thank you for your feedback!",
                             font=("Arial", 24))
        message.pack(expand=True)
        
        # Auto-close timer
        Timer(3, thank_you.destroy).start()
    
    def handle_correction(self, correct_category, dialog):
        dialog.destroy()
        self.show_thank_you()
    
    def capture_and_classify(self):
        # Capture image
        image = self.camera.capture_array()
        
        # Preprocess image
        height = self.input_details[0]['shape'][1]
        width = self.input_details[0]['shape'][2]
        processed_image = cv2.resize(image, (width, height))
        processed_image = processed_image / 255.0
        processed_image = np.expand_dims(processed_image, axis=0).astype(np.float32)
        
        # Run inference
        self.interpreter.set_tensor(self.input_details[0]['index'], processed_image)
        self.interpreter.invoke()
        predictions = self.interpreter.get_tensor(self.output_details[0]['index'])[0]
        
        # Get result
        max_index = np.argmax(predictions)
        confidence = predictions[max_index] * 100
        class_name = list(self.categories.keys())[max_index]
        
        # Show result
        self.show_classification_result(class_name, confidence)
    
    def run(self):
        self.root.mainloop()

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