import tkinter as tk
from tkinter import filedialog
import cv2
from PIL import Image, ImageTk
import numpy as np
from tensorflow.keras.models import load_model

class ShelfClassifierApp:
    def __init__(self, master):
        self.master = master
        self.master.title("Shelf Classifier")

        self.model = load_model('saved_model.h5')  # Load your model

        self.canvas = tk.Canvas(master, width=300, height=300)
        self.canvas.pack()

        self.load_button = tk.Button(master, text="Load Image", command=self.load_image)
        self.load_button.pack()

        self.classify_button = tk.Button(master, text="Classify", command=self.classify_image)
        self.classify_button.pack()

        self.result_label = tk.Label(master, text="")
        self.result_label.pack()

        self.image = None

    def load_image(self):
        file_path = filedialog.askopenfilename()
        if file_path:
            self.image = cv2.imread(file_path)
            self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
            self.display_image(self.image)

    def display_image(self, image):
        image = Image.fromarray(image)
        image = ImageTk.PhotoImage(image)
        self.canvas.create_image(0, 0, anchor=tk.NW, image=image)
        self.canvas.image = image

    def classify_image(self):
        if self.image is not None:
            # Preprocess the image
            resized_image = cv2.resize(self.image, (224, 224))
            resized_image = resized_image.astype('float32') / 255
            resized_image = np.expand_dims(resized_image, axis=0)

            # Make prediction
            prediction = self.model.predict(resized_image)

            # Postprocess the prediction
            class_index = np.argmax(prediction)
            class_label = "Disorganized or Empty" if class_index == 1 else "Organized"

            # Draw bounding box if shelf is disorganized or empty
            if class_index == 1:
                # Draw red rectangle
                image_with_box = cv2.rectangle(self.image, (0, 0), (self.image.shape[1], self.image.shape[0]), (255, 0, 0), 2)
                self.display_image(image_with_box)
            else:
                self.display_image(self.image)

            self.result_label.config(text=class_label)
        else:
            self.result_label.config(text="Please load an image first")

def main():
    root = tk.Tk()
    app = ShelfClassifierApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()