Spaces:
Sleeping
Sleeping
File size: 2,669 Bytes
777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 777385b bb7f5b6 |
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 |
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
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: np.ndarray, detections, labels: List[str]) -> 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: np.ndarray, confidence_threshold: float):
results = query(Image.fromarray(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, ", ".join(final_labels)
def query(image: Image.Image, confidence_threshold: float):
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.Image(label="Input Image", type="numpy")
conf = gr.Slider(label="Confidence Threshold", minimum=0.1, maximum=1.0, value=0.6, step=0.05)
output_image = gr.Image(label="Output Image", type="numpy")
output_text = gr.Textbox(label="Detected Classes")
def process_and_display(input_image, conf):
output_img, detected_classes = process_image(input_image, conf)
return output_img, detected_classes
gr.Interface(
fn=process_and_display,
inputs=[input_image, conf],
outputs=[output_image, output_text],
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.",
).launch()
if __name__ == "__main__":
run_demo()
|