Spaces:
Sleeping
Sleeping
File size: 2,479 Bytes
777385b 280baf3 777385b |
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 |
from typing import List
import os
import numpy as np
import torch
import gradio as gr
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForObjectDetection
import supervision as sv
import spaces
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = AutoImageProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
model = AutoModelForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365").to(device)
BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator()
MASK_ANNOTATOR = sv.MaskAnnotator()
LABEL_ANNOTATOR = sv.LabelAnnotator()
TRACKER = sv.ByteTrack()
def annotate_image(
input_image,
detections,
labels
) -> np.ndarray:
output_image = MASK_ANNOTATOR.annotate(input_image, detections)
output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels)
return output_image
def process_image(
input_image,
confidence_threshold,
):
results = query(input_image, confidence_threshold)
detections = sv.Detections.from_transformers(results[0])
detections = TRACKER.update_with_detections(detections)
final_labels = [model.config.id2label[label] for label in detections.class_id.tolist()]
output_image = annotate_image(input_image, detections, final_labels)
return output_image
def query(image, confidence_threshold):
inputs = processor(images=image, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model(**inputs)
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs=outputs, threshold=confidence_threshold, target_sizes=target_sizes)
return results
def run_demo():
input_image = gr.inputs.Image(label="Input Image")
conf = gr.inputs.Slider(label="Confidence Threshold", minimum=0.1, maximum=1.0, value=0.6, step=0.05)
output_image = gr.outputs.Image(label="Output Image")
def process_and_display(input_image, conf):
output_img = process_image(input_image, conf)
return output_img
gr.Interface(
fn=process_and_display,
inputs=[input_image, conf],
outputs=output_image,
title="Real Time Object Detection with RT-DETR",
description="This Demo uses RT-DETR for object detection in images. Adjust the confidence threshold to see different results.",
capture_session=True,
).launch()
|