Untitled
unknown
plain_text
6 months ago
11 kB
3
Indexable
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()
Editor is loading...
Leave a Comment