Spaces:
Running
on
T4
Running
on
T4
File size: 3,863 Bytes
094752c 377e9f4 094752c 377e9f4 094752c 377e9f4 094752c 377e9f4 094752c 377e9f4 094752c 377e9f4 094752c 377e9f4 094752c |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import gradio as gr
import spaces
import supervision as sv
from rfdetr import RFDETRBase, RFDETRLarge
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.
"""
IMAGE_EXAMPLES = [
['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, "large"],
['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, "large"],
['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, "base"],
]
COLOR = sv.ColorPalette.from_hex([
"#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff",
"#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00"
])
MODEL_BASE = RFDETRBase(resolution=728)
MODEL_LARGE = RFDETRLarge(resolution=728)
@spaces.GPU()
def inference(image, confidence: float, checkpoint: str):
detections = MODEL_BASE.predict(image, threshold=confidence) \
if checkpoint == "base" \
else MODEL_LARGE.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,
)
with gr.Row():
checkpoint_dropdown = gr.Dropdown(
label="Checkpoint",
choices=["base", "large"],
value="base"
)
submit_button = gr.Button("Submit")
with gr.Column():
output_image = gr.Image(
label="Input Image",
image_mode='RGB',
type='pil',
height=600
)
gr.Examples(
fn=inference,
examples=IMAGE_EXAMPLES,
inputs=[input_image, confidence_slider, checkpoint_dropdown],
outputs=output_image
)
submit_button.click(
inference,
inputs=[input_image, confidence_slider, checkpoint_dropdown],
outputs=output_image
)
demo.launch(debug=False, show_error=True)
|