import gradio as gr
from PIL import Image, ImageDraw
import numpy as np
import tensorflow as tf

nn = tf.keras.models.load_model('nn.keras')

def detect_object(img: Image.Image) -> Image.Image:

    img_gray = img.convert("L")
    arr_gray = np.array(img_gray)                  # shape = (H, W)
    orig_h, orig_w = arr_gray.shape

    img_resized = img_gray.resize((64, 64))        # PIL resize
    arr_resized = np.array(img_resized).astype("float32") / 255.0
    inp = arr_resized.reshape(1, -1)

    pred = nn.predict(inp)[0]
    x_norm, y_norm, w_norm, h_norm = pred

    xmin = x_norm * orig_w
    ymin = y_norm * orig_h
    xmax = xmin + (w_norm * orig_w)
    ymax = ymin + (h_norm * orig_h)

    img_out = img.convert("RGB")
    draw   = ImageDraw.Draw(img_out)
    draw.rectangle(
        [int(xmin), int(ymin), int(xmax), int(ymax)],
        outline="red",
        width=2
    )

    return img_out


interface = gr.Interface(
    fn=detect_object,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Image(type="pil", label="Detected object"),
    title="Object Detection",
    description="Simple Neural Network Model For Object Detection"
)

if __name__ == "__main__":
  interface.launch(share=True)