File size: 1,954 Bytes
2a01df3 89fca02 903bcf4 2a01df3 903bcf4 4edcf13 89fca02 4edcf13 89fca02 903bcf4 89fca02 4edcf13 89fca02 4edcf13 89fca02 4edcf13 89fca02 903bcf4 92b8946 2a01df3 4edcf13 2a01df3 92b8946 2a01df3 903bcf4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import os
import tempfile
import base64
from io import BytesIO
from PIL import Image
def perform(image, task: str="detect"):
# Predict with the model
results = model(image, save=True) # predict on an image
# get the saved prediction
prediction_image = Image.open(os.path.join(results[0].save_dir, os.listdir(results[0].save_dir)[0]))
# Convert PIL image to base64
buffered = BytesIO()
prediction_image.save(buffered, format="JPEG")
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
# initialize the output list to save the bbox
outputs = []
# Access the results
for result in results[0]:
xywh = result.boxes.xyxy[0].numpy().tolist() # top-left-x, top-left-y, bottom-right-x, bottom-right-y
name = result.names[result.boxes.cls.item()] # classes
confs = float(result.boxes.conf) # confidence score of each box
outputs.append({"xywh": xywh, "conf": confs, "name": name})
# return the PIL image (for display) and the outputs with base64
return prediction_image, outputs, {"image_base64": f"data:image/jpeg;base64,{img_base64}"}
def load_model(model_name: str = "yolo11n.pt"):
"""Model loading"""
try:
print(f"Loading the model {model_name}")
model = YOLO(model_name)
except Exception as e:
print(f"Error loading the model {e}")
return model
# loading the model
model = load_model()
# Create the Gradio interface
demo = gr.Interface(
fn=perform,
inputs=[gr.Image(type="pil")], # PIL Image format
outputs=[gr.Image(type="pil"), gr.JSON(), gr.JSON()],
title="Object Detection using YOLO",
description="Upload an image and get image with bbox in it",
examples=None # You can add example images here if needed
)
# Launch the app
if __name__ == "__main__":
demo.launch() |