StockDetection / app.py
osamaifti's picture
Update app.py
6af426e verified
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()