app.py

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

def clear_output_on_change(img):
    return gr.update(value="")

def clear_all():
    return None, None

with gr.Blocks(css='''.example * {
    font-style: italic;
    font-size: 18px !important;
    color: #0ea5e9 !important;
    }''') as demo:

    gr.Markdown("### Image Classification with MobileNet (OpenCV DNN)")
    gr.Markdown("Upload an image to classify using a MobileNet model loaded with OpenCV DNN.")

    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload Image")
        output_box = gr.Textbox(label="Top Prediction(s)")

    # Clear output when new image is uploaded
    image_input.change(fn=clear_output_on_change, inputs=image_input, outputs=output_box)

    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
        clear_btn = gr.Button("Clear")

    submit_btn.click(fn=classify_image, inputs=image_input, outputs=output_box)
    clear_btn.click(fn=clear_all, outputs=[image_input, output_box])

    gr.Markdown("Click on any example to try it.", elem_classes=["example"])

    gr.Examples(
        examples=[
            ["examples/squirrel_cls.jpg"],
            ["examples/baboon.jpg"]
        ],
        inputs=image_input
    )

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