|
import json |
|
import gradio as gr |
|
import yolov5 |
|
from PIL import Image |
|
from huggingface_hub import hf_hub_download |
|
|
|
app_title = "License Plate Object Detection" |
|
models_ids = ['keremberke/yolov5n-license-plate', 'keremberke/yolov5s-license-plate', 'keremberke/yolov5m-license-plate'] |
|
article = f"<p style='text-align: center'> <a href='https://huggingface.co/{models_ids[-1]}'>model</a> | <a href='https://huggingface.co/keremberke/license-plate-object-detection'>dataset</a> | <a href='https://github.com/keremberke/awesome-yolov5-models'>awesome-yolov5-models</a> </p>" |
|
|
|
current_model_id = models_ids[-1] |
|
model = yolov5.load(current_model_id) |
|
|
|
|
|
|
|
|
|
def predict(image, threshold=0.25, model_id=None): |
|
|
|
global current_model_id |
|
global model |
|
if model_id != current_model_id: |
|
model = yolov5.load(model_id) |
|
current_model_id = model_id |
|
|
|
|
|
config_path = hf_hub_download(repo_id=model_id, filename="config.json") |
|
with open(config_path, "r") as f: |
|
config = json.load(f) |
|
input_size = config["input_size"] |
|
|
|
|
|
model.conf = threshold |
|
results = model(image, size=input_size) |
|
numpy_image = results.render()[0] |
|
output_image = Image.fromarray(numpy_image) |
|
predictions = results.pred[0] |
|
print(predictions[:, :4]) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return output_image |
|
|
|
|
|
gr.Interface( |
|
title=app_title, |
|
description="Created by 'keremberke'", |
|
article=article, |
|
fn=predict, |
|
inputs=[ |
|
gr.Image(type="pil"), |
|
gr.Slider(maximum=1, step=0.01, value=0.25), |
|
gr.Dropdown(models_ids, value=models_ids[-1]), |
|
], |
|
outputs=gr.Image(type="pil"), |
|
|
|
|
|
).launch(enable_queue=True) |
|
|