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