import cv2 as cv
import numpy as np
import gradio as gr
from mobilenet import MobileNet
from huggingface_hub import hf_hub_download

# Download ONNX model from Hugging Face
model_path = hf_hub_download(repo_id="opencv/image_classification_mobilenet", filename="image_classification_mobilenetv1_2022apr.onnx")
top_k = 1
backend_id = cv.dnn.DNN_BACKEND_OPENCV
target_id = cv.dnn.DNN_TARGET_CPU

# Load MobileNet model
model = MobileNet(modelPath=model_path, topK=top_k, backendId=backend_id, targetId=target_id)

def classify_image(input_image):
    image = cv.resize(input_image, (256, 256))
    image = image[16:240, 16:240, :]

    result = model.infer(image)

    result_str = "\n".join(f"{label}" for label in result)
    return result_str

# Gradio Interface
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Textbox(label="Top Prediction(s)"),
    title="Image Classification with MobileNet (OpenCV DNN)",
    allow_flagging="never",
    description="Upload an image to classify using a MobileNet model loaded with OpenCV DNN."
)

if __name__ == "__main__":
    demo.launch()