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