EchoYolo / app.py
mkhug98's picture
Update app.py
5a4893d verified
import gradio as gr
from huggingface_hub import hf_hub_download
import onnxruntime
from huggingface_hub import ModelCard
card = ModelCard.load('mkhug98/Echo-Yolo')
# Download the ONNX model from Hugging Face
model_path = hf_hub_download(repo_id="mkhug98/Echo-Yolo", filename="best.onnx")
# Load the ONNX model
session = onnxruntime.InferenceSession(model_path)
# Function to perform object detection
def detect_objects(image):
# Preprocess the image
image = image.resize((640, 640)) # Resize the image to the expected input size
input_data = image.transpose(2, 0, 1).numpy() # Rearrange the dimensions for ONNX
# Perform inference with the ONNX model
outputs = session.run(None, {"images": input_data.astype("float32")})
bboxes, scores, class_ids = outputs
# Create a list of dictionaries for each detected object
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
# Create the Gradio app
app = gr.Interface(detect_objects, gr.Image(type="pil"), "label", examples=[
["example_image.jpg"] # Replace with your own example image
])
# Run the app
app.launch()