Spaces:
Sleeping
Sleeping
| import logging | |
| from typing import List, 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 HuggingFaceYoloV8Detector(ObjectDetector): | |
| """YOLOv8 detector whose weights are fetched from the Hugging Face Hub.""" | |
| REPO_ID = "Ultralytics/YOLOv8" | |
| WEIGHT_FILE = "yolov8s.pt" | |
| def __init__(self, score_threshold: float = 0.3) -> None: | |
| self.name = "hf_yolov8" | |
| self.score_threshold = score_threshold | |
| self.device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| logging.info( | |
| "Loading Hugging Face YOLOv8 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 _filter_indices(self, label_names: Sequence[str], queries: Sequence[str]) -> List[int]: | |
| if not queries: | |
| return list(range(len(label_names))) | |
| allowed = {query.lower().strip() for query in queries if query} | |
| keep = [idx for idx, name in enumerate(label_names) if name.lower() in allowed] | |
| return keep or list(range(len(label_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] | |
| keep_indices = self._filter_indices(label_names, queries) | |
| xyxy = xyxy[keep_indices] if len(xyxy) else xyxy | |
| scores = [scores[i] for i in keep_indices] | |
| label_ids = [label_ids[i] for i in keep_indices] | |
| label_names = [label_names[i] for i in keep_indices] | |
| return DetectionResult( | |
| boxes=xyxy, | |
| scores=scores, | |
| labels=label_ids, | |
| label_names=label_names, | |
| ) | |