Untitled

 avatar
unknown
plain_text
a year ago
11 kB
12
Indexable
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):
        self.setup_window()
        self.setup_camera()
        self.setup_model()
        self.create_main_view()
        
        # Create image storage folders
        self.categories = {
            "Biomüll": "brown",
            "Gelber Sack": "yellow",
            "Papier": "blue",
            "Restmüll": "gray"
        }
        
        for category in self.categories:
            os.makedirs(f"captured_images/{category}", exist_ok=True)
    
    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)
        header.pack(fill='x', padx=20, pady=20)
        
        # 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 scan button
        scan_frame = ctk.CTkFrame(self.main_container)
        scan_frame.pack(side='bottom', pady=20)
        
        scan_button = ctk.CTkButton(scan_frame,
                                  text="Scan",
                                  command=self.capture_and_classify,
                                  width=120, height=40)
        scan_button.pack(side='left', padx=10)
    
    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()
Editor is loading...
Leave a Comment