Spaces:
Running
Running
import gradio as gr | |
from PIL import Image | |
import numpy as np | |
from ultralytics import YOLO | |
import os | |
def handle_classify(image=None): | |
"""This function performs YOLOv8 object detection on the given image. | |
Args: | |
image (gr.inputs.Image, optional): Input image to detect objects on. Defaults to None. | |
""" | |
model_path = "racist2.0.pt" | |
model = YOLO(model_path) | |
results = model(image) | |
result = results[0] | |
top5 = [[result.names[class_index], str(round(result.probs.top5conf.tolist()[rank], 4)*100)+'%'] | |
for class_index, rank in zip(result.probs.top5, range(5))] | |
print(top5) | |
return "\n".join(["\t".join(row) for row in top5]) | |
inputs = [ | |
gr.Image(label="Input Image"), | |
] | |
outputs = gr.Textbox() | |
title = "Racist model v2" | |
SAMPLE_DIR = 'samples' | |
examples = [os.path.join(SAMPLE_DIR, path) for path in os.listdir(SAMPLE_DIR)] | |
yolo_app = gr.Interface( | |
fn=handle_classify, | |
inputs=inputs, | |
outputs=outputs, | |
title=title, | |
examples=examples, | |
cache_examples=True, | |
) | |
# Launch the Gradio interface in debug mode with queue enabled | |
yolo_app.launch(debug=True) |