Spaces:
Sleeping
Sleeping
File size: 1,985 Bytes
5995290 |
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 |
import logging
from typing import Sequence
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from models.detectors.base import DetectionResult, ObjectDetector
class HuggingFaceYoloV12BotSortDetector(ObjectDetector):
"""YOLOv12 model (BoT-SORT + ReID) hosted on Hugging Face."""
REPO_ID = "wish44165/YOLOv12-BoT-SORT-ReID"
WEIGHT_FILE = "MOT_yolov12n.pt"
def __init__(self, score_threshold: float = 0.3) -> None:
self.name = "hf_yolov12_bot_sort"
self.score_threshold = score_threshold
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
logging.info(
"Loading Hugging Face YOLOv12 BoT-SORT weights %s/%s onto %s",
self.REPO_ID,
self.WEIGHT_FILE,
self.device,
)
weight_path = hf_hub_download(repo_id=self.REPO_ID, filename=self.WEIGHT_FILE)
self.model = YOLO(weight_path)
self.model.to(self.device)
self.class_names = self.model.names
def predict(self, frame: np.ndarray, queries: Sequence[str]) -> DetectionResult:
device_arg = 0 if self.device.startswith("cuda") else "cpu"
results = self.model.predict(
source=frame,
device=device_arg,
conf=self.score_threshold,
verbose=False,
)
result = results[0]
boxes = result.boxes
if boxes is None or boxes.xyxy is None:
empty = np.empty((0, 4), dtype=np.float32)
return DetectionResult(empty, [], [], [])
xyxy = boxes.xyxy.cpu().numpy()
scores = boxes.conf.cpu().numpy().tolist()
label_ids = boxes.cls.cpu().numpy().astype(int).tolist()
label_names = [self.class_names.get(idx, f"class_{idx}") for idx in label_ids]
return DetectionResult(
boxes=xyxy,
scores=scores,
labels=label_ids,
label_names=label_names,
)
|