Guy2's picture
update
6503948
raw
history blame
1.05 kB
import gradio as gr
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
import supervision as sv
def anylize(img):
id2label = {k: v['name'] for k,v in categories.items()}
box_annotator = sv.BoxAnnotator()
image = img
with torch.no_grad():
# load image and predict
inputs = image_processor(images=image, return_tensors='pt')
outputs = model(**inputs)
# post-process
target_sizes = torch.tensor([image.shape[:2]])
results = image_processor.post_process_object_detection(
outputs=outputs,
threshold=0.8,
target_sizes=target_sizes
)[0]
# annotate
detections = sv.Detections.from_transformers(transformers_results=results).with_nms(threshold=0.5)
labels = [f"{id2label[class_id]} {confidence:.2f}" for _, confidence, class_id, _ in detections]
frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)
gr.Interface.load("models/Guy2/AirportSec-100epoch").launch()