Spaces:
Runtime error
Runtime error
File size: 1,739 Bytes
72e5fb9 f1fb617 72e5fb9 f1fb617 72e5fb9 |
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 |
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image
class DetectionService:
def __init__(self, model_name="facebook/detr-resnet-50"):
self.processor = DetrImageProcessor.from_pretrained(model_name, revision="no_timm")
self.model = DetrForObjectDetection.from_pretrained(model_name, revision="no_timm")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
self.model.eval()
self.frame_counter = 0
self.frame_skip = 5 # Process every 5th frame for performance
def detect_objects(self, image, confidence_threshold=0.9):
"""Detect objects in an image, skipping frames for performance."""
self.frame_counter += 1
if self.frame_counter % self.frame_skip != 0:
return [] # Skip detection for this frame
inputs = self.processor(images=image, return_tensors="pt").to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
target_sizes = torch.tensor([image.size[::-1]]).to(self.device)
results = self.processor.post_process_object_detection(
outputs, target_sizes=target_sizes, threshold=confidence_threshold
)[0]
detections = []
for score, label, box in zip(
results["scores"], results["labels"], results["boxes"]
):
box = box.cpu().numpy().astype(int)
detections.append({
"score": score.item(),
"label": self.model.config.id2label[label.item()],
"box": {"xmin": box[0], "ymin": box[1], "xmax": box[2], "ymax": box[3]}
})
return detections |