VRE / src /detection.py
nasim-raj-laskar
initial deploy
aeba2d0
raw
history blame contribute delete
487 Bytes
from transformers import pipeline
# Load once
detector = pipeline(
"object-detection",
model="facebook/detr-resnet-50",
device=-1 # CPU)
)
def detect_objects(image, threshold=0.7):
results = detector(image)
filtered = [r for r in results if r["score"] >= threshold]
labels = list({r["label"] for r in filtered})
boxes = [
{"label": r["label"], "score": round(r["score"], 2), "box": r["box"]}
for r in filtered
]
return labels, boxes