|
|
import gradio as gr |
|
|
import spaces |
|
|
import supervision as sv |
|
|
from rfdetr import RFDETRBase |
|
|
from rfdetr.util.coco_classes import COCO_CLASSES |
|
|
|
|
|
MARKDOWN = """ |
|
|
# RF-DETR 🔥 |
|
|
|
|
|
<div style="display: flex; align-items: center; gap: 8px;"> |
|
|
<a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-rf-detr-on-detection-dataset.ipynb"> |
|
|
<img src="https://colab.research.google.com/assets/colab-badge.svg" alt="colab" /> |
|
|
</a> |
|
|
<a href="https://blog.roboflow.com/rf-detr"> |
|
|
<img src="https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg" alt="roboflow" /> |
|
|
</a> |
|
|
<a href="https://github.com/roboflow/rf-detr"> |
|
|
<img src="https://badges.aleen42.com/src/github.svg" alt="roboflow" /> |
|
|
</a> |
|
|
</div> |
|
|
|
|
|
RF-DETR is a real-time, transformer-based object detection model architecture developed |
|
|
by [Roboflow](https://roboflow.com/) and released under the Apache 2.0 license. |
|
|
""" |
|
|
|
|
|
COLOR = sv.ColorPalette.from_hex([ |
|
|
"#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff", |
|
|
"#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00" |
|
|
]) |
|
|
|
|
|
MODEL = RFDETRBase() |
|
|
|
|
|
|
|
|
@spaces.GPU() |
|
|
def inference(image, confidence): |
|
|
detections = MODEL.predict(image, threshold=confidence) |
|
|
|
|
|
text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size) |
|
|
thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size) |
|
|
|
|
|
bbox_annotator = sv.BoxAnnotator(color=COLOR, thickness=thickness) |
|
|
label_annotator = sv.LabelAnnotator( |
|
|
color=COLOR, |
|
|
text_color=sv.Color.BLACK, |
|
|
text_scale=text_scale, |
|
|
smart_position=True |
|
|
) |
|
|
|
|
|
labels = [ |
|
|
f"{COCO_CLASSES[class_id]} {confidence:.2f}" |
|
|
for class_id, confidence |
|
|
in zip(detections.class_id, detections.confidence) |
|
|
] |
|
|
|
|
|
annotated_image = image.copy() |
|
|
annotated_image = bbox_annotator.annotate(annotated_image, detections) |
|
|
annotated_image = label_annotator.annotate(annotated_image, detections, labels) |
|
|
return annotated_image |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
gr.Markdown(MARKDOWN) |
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
input_image = gr.Image( |
|
|
label="Input Image", |
|
|
image_mode='RGB', |
|
|
type='pil', |
|
|
height=600 |
|
|
) |
|
|
confidence_slider = gr.Slider( |
|
|
label="Confidence", |
|
|
minimum=0.0, |
|
|
maximum=1.0, |
|
|
step=0.05, |
|
|
value=0.5, |
|
|
) |
|
|
submit_button = gr.Button("Submit") |
|
|
with gr.Column(): |
|
|
output_image = gr.Image( |
|
|
label="Input Image", |
|
|
image_mode='RGB', |
|
|
type='pil', |
|
|
height=600 |
|
|
) |
|
|
|
|
|
submit_button.click( |
|
|
inference, |
|
|
inputs=[input_image, confidence_slider], |
|
|
outputs=output_image |
|
|
) |
|
|
|
|
|
demo.launch(debug=False, show_error=True) |
|
|
|
