Untitled

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

# Set appearance mode and color theme
ctk.set_appearance_mode("dark")
ctk.set_default_color_theme("blue")

class CircularProgress(ctk.CTkCanvas):
    def __init__(self, parent, size=100, progress_color="#2CC985", **kwargs):
        # Use fg_color instead of bg for CustomTkinter
        super().__init__(
            parent,
            width=size,
            height=size,
            fg_color=parent.cget("fg_color"),
            highlightthickness=0,
            **kwargs
        )
        self.size = size
        self.progress_color = progress_color
        self._percentage = 0
        self.ring_width = size // 10
        self.warning_threshold = 90
        self.warning_color = "#FF5555"
        self.draw_ring()
        
    def draw_ring(self):
        self.delete("progress")
        center = self.size // 2
        radius = (self.size - self.ring_width) // 2
        
        if self._percentage > 0:
            start = -90  # Start from top (- 90 degrees)
            extent = self._percentage * 3.6  # Convert percentage to degrees
            
            # Choose color based on percentage
            color = self.warning_color if self._percentage >= self.warning_threshold else self.progress_color
            
            self.create_arc(
                self.ring_width // 2,
                self.ring_width // 2,
                self.size - self.ring_width // 2,
                self.size - self.ring_width // 2,
                start=start,
                extent=extent,
                tags="progress",
                width=self.ring_width,
                style="arc",
                outline=color
            )
    
    def set_percentage(self, percentage):
        self._percentage = min(100, max(0, percentage))
        self.draw_ring()

class TimerBar(ctk.CTkFrame):
    def __init__(self, parent, duration=5, **kwargs):
        super().__init__(parent, **kwargs)
        self.duration = duration
        self.remaining = duration
        self.bar = ctk.CTkProgressBar(self)
        self.bar.pack(fill="x", padx=5, pady=5)
        self.bar.set(1)
        
    def start(self, callback=None):
        self.callback = callback
        self.update_timer()
        
    def update_timer(self):
        if self.remaining > 0:
            self.bar.set(self.remaining / self.duration)
            self.remaining -= 0.1
            self.after(100, self.update_timer)
        else:
            if self.callback:
                self.callback()

class WasteClassificationGUI:
    def __init__(self):
        self.setup_window()
        self.load_model()
        self.setup_camera()
        self.create_main_screen()
        
    def setup_window(self):
        self.window = ctk.CTk()
        self.window.title("Waste Classification")
        self.window.geometry("1024x600")
        self.window.resizable(False, False)
        
        # Create icons directory if it doesn't exist
        os.makedirs("icons", exist_ok=True)
        
        # Default icon paths
        icon_paths = {
            "scan": "icons/scan.png",
            "thumbs_up": "icons/thumbs_up.png",
            "thumbs_down": "icons/thumbs_down.png"
        }
        
        # Load icons if they exist, otherwise use placeholder images
        self.icons = {}
        for icon_name, path in icon_paths.items():
            if os.path.exists(path):
                self.icons[icon_name] = ctk.CTkImage(
                    Image.open(path),
                    size=(48, 48)
                )
            else:
                # Create a blank image as placeholder
                img = Image.new('RGB', (48, 48), color='gray')
                self.icons[icon_name] = ctk.CTkImage(
                    img,
                    size=(48, 48)
                )
        
    def load_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.class_names = ["Biomüll", "Gelber Sack", "Papier", "Restmüll"]
        self.class_colors = {
            "Biomüll": "#8B4513",
            "Gelber Sack": "#FFD700",
            "Papier": "#4169E1",
            "Restmüll": "#808080"
        }
        
    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 create_main_screen(self):
        self.main_frame = ctk.CTkFrame(self.window)
        self.main_frame.pack(fill="both", expand=True)
        
        # Create circles frame
        circles_frame = ctk.CTkFrame(self.main_frame)
        circles_frame.pack(pady=50)
        
        # Create circles with labels
        self.circles = {}
        circle_size = 120
        for i, category in enumerate(self.class_names):
            circle_frame = ctk.CTkFrame(circles_frame)
            circle_frame.pack(side="left", padx=20)
            
            # Create circular progress
            circle = CircularProgress(
                circle_frame,
                size=circle_size,
                progress_color=self.class_colors[category]
            )
            circle.pack(pady=5)
            
            # Create label
            label = ctk.CTkLabel(
                circle_frame,
                text=category,
                font=("Helvetica", 14, "bold")
            )
            label.pack()
            
            # Create slider
            slider = ctk.CTkSlider(
                circle_frame,
                from_=0,
                to=100,
                command=lambda val, c=circle: c.set_percentage(val)
            )
            slider.pack(pady=5)
            
            self.circles[category] = {
                "progress": circle,
                "slider": slider
            }
        
        # Create scan button
        scan_button = ctk.CTkButton(
            self.main_frame,
            text="",
            image=self.icons["scan"],
            width=60,
            height=60,
            command=self.scan_waste
        )
        scan_button.place(relx=0.1, rely=0.9, anchor="center")
        
    def create_prediction_screen(self, prediction, confidence):
        prediction_frame = ctk.CTkFrame(self.window)
        prediction_frame.pack(fill="both", expand=True)
        
        # Create timer
        timer = TimerBar(prediction_frame, duration=5)
        timer.pack(fill="x", padx=20, pady=10)
        
        # Create prediction display
        result_label = ctk.CTkLabel(
            prediction_frame,
            text=f"Detected: {prediction}\nConfidence: {confidence:.1f}%",
            font=("Helvetica", 24, "bold")
        )
        result_label.pack(pady=50)
        
        # Create feedback buttons
        feedback_frame = ctk.CTkFrame(prediction_frame)
        feedback_frame.pack(pady=20)
        
        ctk.CTkButton(
            feedback_frame,
            text="",
            image=self.icons["thumbs_up"],
            width=60,
            height=60,
            command=lambda: self.handle_feedback(True, prediction)
        ).pack(side="left", padx=10)
        
        ctk.CTkButton(
            feedback_frame,
            text="",
            image=self.icons["thumbs_down"],
            width=60,
            height=60,
            command=lambda: self.handle_feedback(False, prediction)
        ).pack(side="left", padx=10)
        
        # Start timer
        timer.start(lambda: self.switch_screen(prediction_frame, self.main_frame))
        
    def create_feedback_screen(self):
        feedback_frame = ctk.CTkFrame(self.window)
        feedback_frame.pack(fill="both", expand=True)
        
        # Create timer
        timer = TimerBar(feedback_frame, duration=5)
        timer.pack(fill="x", padx=20, pady=10)
        
        # Create category selection
        label = ctk.CTkLabel(
            feedback_frame,
            text="Select correct category:",
            font=("Helvetica", 20, "bold")
        )
        label.pack(pady=20)
        
        categories_frame = ctk.CTkFrame(feedback_frame)
        categories_frame.pack(pady=20)
        
        for category in self.class_names:
            ctk.CTkButton(
                categories_frame,
                text=category,
                fg_color=self.class_colors[category],
                command=lambda c=category: self.save_feedback(c)
            ).pack(side="left", padx=10)
        
        # Start timer
        timer.start(lambda: self.switch_screen(feedback_frame, self.main_frame))
        
    def create_thank_you_screen(self):
        thank_you_frame = ctk.CTkFrame(self.window)
        thank_you_frame.pack(fill="both", expand=True)
        
        # Create timer
        timer = TimerBar(thank_you_frame, duration=3)
        timer.pack(fill="x", padx=20, pady=10)
        
        # Create thank you message
        label = ctk.CTkLabel(
            thank_you_frame,
            text="Thank you for your feedback!",
            font=("Helvetica", 24, "bold")
        )
        label.pack(pady=100)
        
        # Start timer
        timer.start(lambda: self.switch_screen(thank_you_frame, self.main_frame))
        
    def switch_screen(self, old_frame, new_frame):
        old_frame.destroy()
        new_frame.pack(fill="both", expand=True)
        
    def preprocess_image(self, image):
        height = self.input_details[0]['shape'][1]
        width = self.input_details[0]['shape'][2]
        resized = cv2.resize(image, (width, height))
        normalized = resized / 255.0
        return np.expand_dims(normalized, axis=0).astype(np.float32)
        
    def scan_waste(self):
        # Capture image
        image = self.camera.capture_array()
        
        # Save image
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.current_image_path = f"captured_images/capture_{timestamp}.jpg"
        os.makedirs("captured_images", exist_ok=True)
        cv2.imwrite(self.current_image_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
        
        # Process image
        processed_image = self.preprocess_image(image)
        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 prediction
        max_index = np.argmax(predictions)
        confidence = predictions[max_index] * 100
        prediction = self.class_names[max_index]
        
        # Show prediction screen
        self.main_frame.pack_forget()
        self.create_prediction_screen(prediction, confidence)
        
    def handle_feedback(self, is_correct, prediction):
        if is_correct:
            # Save image in correct category folder
            category_folder = f"captured_images/{prediction}"
            os.makedirs(category_folder, exist_ok=True)
            new_path = os.path.join(category_folder, os.path.basename(self.current_image_path))
            os.rename(self.current_image_path, new_path)
            self.create_thank_you_screen()
        else:
            self.create_feedback_screen()
            
    def save_feedback(self, category):
        # Save image in selected category folder
        category_folder = f"captured_images/{category}"
        os.makedirs(category_folder, exist_ok=True)
        new_path = os.path.join(category_folder, os.path.basename(self.current_image_path))
        os.rename(self.current_image_path, new_path)
        self.create_thank_you_screen()
        
    def run(self):
        self.window.mainloop()

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