|
import gradio as gr |
|
from huggingface_hub import hf_hub_download |
|
import onnxruntime |
|
from huggingface_hub import ModelCard |
|
|
|
card = ModelCard.load('mkhug98/Echo-Yolo') |
|
|
|
model_path = hf_hub_download(repo_id="mkhug98/Echo-Yolo", filename="best.onnx") |
|
|
|
|
|
session = onnxruntime.InferenceSession(model_path) |
|
|
|
|
|
def detect_objects(image): |
|
|
|
image = image.resize((640, 640)) |
|
input_data = image.transpose(2, 0, 1).numpy() |
|
|
|
|
|
outputs = session.run(None, {"images": input_data.astype("float32")}) |
|
bboxes, scores, class_ids = outputs |
|
|
|
|
|
detections = [] |
|
for bbox, score, class_id in zip(bboxes[0], scores[0], class_ids[0]): |
|
x1, y1, x2, y2 = bbox |
|
label = session.get_modelmeta().custom_metadata_map["names"][int(class_id)] |
|
detections.append({ |
|
'label': label, |
|
'confidence': float(score), |
|
'x1': float(x1), |
|
'y1': float(y1), |
|
'x2': float(x2), |
|
'y2': float(y2) |
|
}) |
|
|
|
return detections |
|
|
|
|
|
app = gr.Interface(detect_objects, gr.Image(type="pil"), "label", examples=[ |
|
["example_image.jpg"] |
|
]) |
|
|
|
|
|
app.launch() |