from flask import Flask, render_template, request
import cv2
import numpy as np
from tensorflow.keras.models import load_model

app = Flask(__name__)

class ShelfClassifier:
    def __init__(self, model_path):
        self.model = load_model(model_path)

    def classify_image(self, image):
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        resized_image = cv2.resize(image_rgb, (224, 224))
        resized_image = resized_image.astype('float32') / 255
        resized_image = np.expand_dims(resized_image, axis=0)
        prediction = self.model.predict(resized_image)
        class_index = np.argmax(prediction)
        class_label = "Disorganized or Empty" if class_index == 1 else "Organized"
        return class_label

@app.route('/', methods=['GET', 'POST'])
def upload_image():
    if request.method == 'POST':
        # Check if the post request has the file part
        if 'file' not in request.files:
            return render_template('index.html', message='No file part')
        file = request.files['file']
        # If user does not select file, browser also
        # submit an empty part without filename
        if file.filename == '':
            return render_template('index.html', message='No selected file')
        if file:
            # Read uploaded image
            image = cv2.imdecode(np.fromstring(file.read(), np.uint8), cv2.IMREAD_COLOR)
            # Initialize ShelfClassifier with the model
            classifier = ShelfClassifier('saved_model.h5')
            # Perform classification
            class_label = classifier.classify_image(image)
            # Draw bounding box if shelf is disorganized or empty
            if class_label == "Disorganized or Empty":
                # Draw red rectangle
                cv2.rectangle(image, (0, 0), (image.shape[1], image.shape[0]), (255, 0, 0), 2)
            # Save image with bounding box
            cv2.imwrite('result.jpg', image)
            return render_template('result.html', class_label=class_label)
    return render_template('index.html')

if __name__ == '__main__':
    app.run(debug=True)