import gradio as gr
import numpy as np
import pandas as pd
import cv2
import tensorflow as tf
from tensorflow.keras.models import load_model

# Load the CNN model
model = load_model('cnn_bounding_box_model.h5')

# Function to preprocess an image (expecting a numpy array directly from Gradio)
def preprocess_image(image):
    image = cv2.resize(image, (128, 128))
    image = np.array(image, dtype='float32')
    return image

def predict_bounding_box(image):
    preprocessed_image = preprocess_image(image)
    preprocessed_image = np.expand_dims(preprocessed_image, axis=0)
    prediction = model.predict(preprocessed_image)
    print(prediction)
    return prediction[0]

def draw_bounding_box(image, bbox):
    height, width, _ = image.shape
    print("height", height, "width", width)
    x, y, w, h = bbox
    x = int(x)
    y = int(y)
    w = int(w)
    h = int(h)

    
    print("x", x, "y", y, "w", w, "h", h)
    # Draw rectangle
    image = cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 0), 2)
    
    return image

def process_image(image):
    # Convert PIL Image to numpy array
    image_np = np.array(image)
    # Run prediction
    bbox = predict_bounding_box(image_np)
    print(bbox)
    # Draw bounding box on the image
    image_with_bbox = draw_bounding_box(image_np, bbox)
    return image_with_bbox

# Define inputs and outputs for Gradio interface
inputs = gr.Image(type='numpy', label='Upload an Image')
outputs = gr.Image(type='numpy', label='Image with Bounding Box')

# Create Gradio interface
gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, title="Bounding Box Detection",
             description="Upload an image, and the model will predict and draw a bounding box around the object.").launch()