Ifeanyi's picture
Update app.py
565e3f9 verified
raw
history blame
4.42 kB
import supervision as sv
from ultralytics import YOLO
from PIL import Image
import gradio as gr
import numpy as np
import cv2
import urllib
# load pre-trained vision model
model = YOLO("yolo12s.pt")
def image_annotate(image:str, annotator:str) -> Image.Image:
"""
Args:
image: the path to the image file
annotator: the type of annotator to use
Returns:
annotated image
"""
# load the input image
image = cv2.imread(image)
# run object detection on the image
result = model(image)[0]
# convert YOLO output to a Supervision-compatible detections format
detections = sv.Detections.from_ultralytics(result)
# select annotator
if annotator == "Box":
box_annotator = sv.BoxAnnotator()
annotated_image_show = box_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Roundbox":
round_box_annotator = sv.RoundBoxAnnotator()
annotated_image_show = round_box_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Boxcorner":
corner_annotator = sv.BoxCornerAnnotator()
annotated_image_show = corner_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Color":
color_annotator = sv.ColorAnnotator()
annotated_image_show = color_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Circle":
circle_annotator = sv.CircleAnnotator()
annotated_image_show = circle_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Dot":
dot_annotator = sv.DotAnnotator()
annotated_image_show = dot_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Triangle":
triangle_annotator = sv.TriangleAnnotator()
annotated_image_show = triangle_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Ellipse":
ellipse_annotator = sv.EllipseAnnotator()
annotated_image_show = ellipse_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Percentage":
percentage_bar_annotator = sv.PercentageBarAnnotator()
annotated_image_show = percentage_bar_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Heatmap":
heatmap_annotator = sv.HeatMapAnnotator()
annotated_image_show = heatmap_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Label":
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections.class_name, detections.confidence)
]
rich_label_annotator = sv.RichLabelAnnotator(
text_color=sv.Color.BLACK,
text_padding=10,
text_position=sv.Position.CENTER)
annotated_image_show = rich_label_annotator.annotate(
scene=image.copy(),
detections=detections,
labels=labels )
elif annotator == "Blur":
blur_annotator = sv.BlurAnnotator()
annotated_image_show = blur_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Pixelate":
pixelate_annotator = sv.PixelateAnnotator()
annotated_image_show = pixelate_annotator.annotate(
scene=image.copy(),
detections=detections)
elif annotator == "Backgroundcolor":
background_overlay_annotator = sv.BackgroundOverlayAnnotator()
annotated_image_show = background_overlay_annotator.annotate(
scene=image.copy(),
detections=detections)
# return annotated image
return annotated_image_show
app = gr.Interface(
fn = image_annotate,
title="Object Detection",
inputs = [gr.Image(type="filepath",label="Image"),gr.Radio(label="Select Annotator",
choices=["Box","Roundbox","Boxcorner","Color","Circle","Dot","Triangle",
"Ellipse","Percentage","Heatmap","Label","Blur",
"Pixelate","Backgroundcolor"],
value = "Box")],
outputs = gr.Image(label = "Annotated Image"),
examples = [["cars.jpg"],
["colorful-backgrounds-for-laptops.jpg"],
["final_animals-homeschooling_credit-alamy.jpg"]]
)
if __name__ == "__main__":
app.launch(mcp_server = True)