|
import cv2 as cv |
|
import numpy as np |
|
import gradio as gr |
|
from mobilenet import MobileNet |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
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 |
|
|
|
|
|
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)") |
|
|
|
|
|
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() |
|
|