Spaces:
Sleeping
Sleeping
File size: 1,716 Bytes
538d0bc 713f67e 166ff57 538d0bc 6d9d00c 538d0bc 166ff57 6d9d00c 166ff57 6d9d00c 166ff57 6d9d00c 713f67e 538d0bc 166ff57 713f67e 166ff57 713f67e 166ff57 538d0bc 166ff57 6d9d00c 166ff57 538d0bc |
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 |
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()
|