Spaces:
Running
Running
File size: 1,840 Bytes
716ca0c |
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 60 |
from typing import List
import gradio as gr
import numpy as np
import supervision as sv
from inference.models import YOLOWorld
MARKDOWN = """
# YOLO-World 🌎
Powered by Roboflow [Inference](https://github.com/roboflow/inference) and [Supervision](https://github.com/roboflow/supervision).
"""
MODEL = YOLOWorld(model_id="yolo_world/l")
BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator()
LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
def process_categories(categories: str) -> List[str]:
return [category.strip() for category in categories.split(',')]
def process_image(input_image: np.ndarray, categories: str) -> np.ndarray:
categories = process_categories(categories)
MODEL.set_classes(categories)
results = MODEL.infer(input_image, confidence=0.003)
detections = sv.Detections.from_inference(results).with_nms(0.1)
output_image = input_image.copy()
output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
output_image = LABEL_ANNOTATOR.annotate(output_image, detections)
return output_image
with gr.Blocks() as demo:
gr.Markdown(MARKDOWN)
with gr.Row():
input_image_component = gr.Image(
type='numpy',
label='Input Image'
)
output_image_component = gr.Image(
type='numpy',
label='Output Image'
)
with gr.Row():
categories_text_component = gr.Textbox(
label='Categories',
placeholder='comma separated list of categories',
scale=5
)
submit_button_component = gr.Button('Submit', scale=1)
submit_button_component.click(
fn=process_image,
inputs=[input_image_component, categories_text_component],
outputs=output_image_component
)
demo.launch(debug=False, show_error=True)
|