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