File size: 1,279 Bytes
a3555b5 |
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 |
import numpy as np
from PIL import Image
import gradio as gr
from ultralytics import YOLO
# Load the YOLO model
m_raw_model = YOLO("M-Raw.pt")
n_raw_model = YOLO("N-Raw.pt")
s_raw_model = YOLO("S-Raw.pt")
def snap(image, model, conf):
# Convert the image to a numpy array
image = np.array(image)
# Run the selected model
results = None
if model == "M-Raw":
results = m_raw_model(image, conf=conf)
elif model == "N-Raw":
results = n_raw_model(image, conf=conf)
elif model == "S-Raw":
results = s_raw_model(image, conf=conf)
# Draw the bounding boxes
resulting_image = results.render()
# Convert the resulting image to a PIL image
resulting_image = Image.fromarray(resulting_image)
# Get the labels
labels = results.pandas().xyxy[0]["name"].values
# Sort the labels by their x-value
labels = labels[np.argsort(results.pandas().xyxy[0]["x"].values)]
return [resulting_image, labels]
demo = gr.Interface(
snap,
[gr.Image(source="webcam", tool=None, streaming=True), gr.inputs.Radio(["M-Raw", "N-Raw", "S-Raw"]), gr.inputs.Slider(0.0, 1.0, 0.5, 0.1, "Confidence")],
["image", "labels"],
title="Baybayin Instance Detection"
)
if __name__ == "__main__":
demo.launch() |