app.py

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
from tensorflow.keras.models import load_model
import joblib

# Load models
reg_lin = joblib.load("linear_model.joblib")
reg_rf = joblib.load("random_forest_model.joblib")
reg_knn = joblib.load("knn_model.joblib")
mlp_model = load_model("mlp_model.keras")
cnn_model = load_model("cnn_model.keras")

def predict(image, model_choice):
    image = image.convert("L").resize((128, 128))
    img_array = np.array(image)
    img_draw = image.convert("RGB")

    if model_choice in ["Linear Regression", "Random Forest", "KNN", "MLP"]:
        img_flat = img_array.flatten().reshape(1, -1) / 255.0

        if model_choice == "Linear Regression":
            pred = reg_lin.predict(img_flat)[0]
        elif model_choice == "Random Forest":
            pred = reg_rf.predict(img_flat)[0]
        elif model_choice == "KNN":
            pred = reg_knn.predict(img_flat)[0]
        elif model_choice == "MLP":
            pred = mlp_model.predict(img_flat)[0]

    elif model_choice == "CNN":
        img_reshaped = img_array.reshape(1, 128, 128, 1) / 255.0
        pred = cnn_model.predict(img_reshaped)[0]

    draw = ImageDraw.Draw(img_draw)
    x, y, w, h = pred
    draw.rectangle([x, y, x + w, y + h], outline="red", width=2)
    return img_draw

interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload Grayscale Image (128x128)"),
        gr.Dropdown(
            choices=["Linear Regression", "Random Forest", "KNN", "MLP", "CNN"],
            label="Select Model"
        )
    ],
    outputs=gr.Image(label="Predicted Bounding Box"),
    title="Object Localization Demo – Group 32",
    description="Upload a grayscale image and select a model to visualize bounding box predictions."
)

interface.launch()