Spaces:
Sleeping
Sleeping
import gradio as gr | |
from PIL import Image | |
import torch | |
from ultralyticsplus import YOLO, render_result | |
available_models = ["YOLOv8n", "YOLOv8n-GhostNet-P5", "YOLOv8n-GhostNet-P6"] | |
available_models_path = [ | |
"./models/yolov8n.pt", | |
"./models/yolov8n_ghostnet_p5.pt", | |
"./models/yolov8n_ghostnet_p6.pt", | |
] | |
def launch( | |
image: gr.Image = None, | |
selectedModel: gr.Dropdown = available_models[0], | |
conf_threshold: gr.Slider = 0.4, | |
iou_threshold: gr.Slider = 0.50, | |
): | |
selected_model_index = available_models.index(selectedModel) | |
image_size = (256,) | |
try: | |
model = YOLO(available_models_path[selected_model_index]) | |
# pil_image = Image.fromarray(image) | |
results = model.predict( | |
image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size | |
) | |
box = results[0].boxes | |
# print(box) | |
render = render_result(model=model, image=image, result=results[0]) | |
return render | |
except Exception as e: | |
print("error", e) | |
return "./download.jpeg" | |
inputs = [ | |
gr.Image(type="filepath", label="Input Image"), | |
gr.Dropdown( | |
info="Choose which model should be used in this task", | |
choices=available_models, | |
value=available_models[0], | |
label="Models", | |
), | |
# gr.Slider(minimum=256, maximum=1280, value=640, step=32, label="Image Size"), | |
gr.Slider( | |
minimum=0.0, maximum=1.0, value=0.4, step=0.1, label="Confidence Threshold" | |
), | |
gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.1, label="IOU Threshold"), | |
] | |
outputs = gr.Image(type="filepath", label="Output Result") | |
iface = gr.Interface(fn=launch, inputs=inputs, outputs=outputs) | |
iface.launch() | |