Huggingface Model¶

def load_model(self) -> None:
    """Load model from HuggingFace."""
    from transformers import AutoConfig, AutoModelForObjectDetection, AutoProcessor

    hf_token = os.getenv("HF_TOKEN", self._token)
    assert self.model_path is not None, "model_path must be provided for HuggingFace models"
    config = AutoConfig.from_pretrained(self.model_path, token=hf_token)
    if self._is_zero_shot(config):
        ensure_package_minimum_version("transformers", "4.49.0")
        from transformers import AutoModelForZeroShotObjectDetection

        model_class: Any = AutoModelForZeroShotObjectDetection
    else:
        model_class = AutoModelForObjectDetection
    model = model_class.from_pretrained(self.model_path, token=hf_token)
    if self.image_size is not None:
        # RT-DETR family expects explicit height/width; other models use shortest_edge
        if model.__class__.__name__.startswith("RTDetr"):
            size: dict[str, int | None] = {"height": self.image_size, "width": self.image_size}
        else:
            size = {"shortest_edge": self.image_size, "longest_edge": None}
        # use_fast=True raises error: AttributeError: 'SizeDict' object has no attribute 'keys'
        processor = AutoProcessor.from_pretrained(
            self.model_path, size=size, do_resize=True, use_fast=False, token=hf_token
        )
    else:
        processor = AutoProcessor.from_pretrained(self.model_path, use_fast=False, token=hf_token)
    self.set_model(model, processor)

def set_model(self, model: Any, processor: Any | None = None, **kwargs: Any) -> None:
    """Set the detection model and processor."""
    processor = processor or self.processor
    if processor is None:
        raise ValueError(f"'processor' is required to be set, got {processor}.")
    self._is_zero_shot_model = self._is_zero_shot(model)
    valid_processor = "ImageProcessor" in processor.__class__.__name__ or self._is_zero_shot(processor)
    if "ObjectDetection" not in model.__class__.__name__ or not valid_processor:
        raise ValueError(
            "Given 'model' is not an ObjectDetectionModel or 'processor' is not a valid ImageProcessor."
        )
    self.model = model
    self.model.to(self.device)  # type: ignore[attr-defined]
    self._processor = processor
    if self._is_zero_shot_model:
        self.category_mapping = {i: name for i, name in enumerate(self.text_labels or [])}
        self._category_name_to_id = {name: i for i, name in self.category_mapping.items()}
    else:
        self.category_mapping = self.model.config.id2label  # type: ignore[attr-defined]

def perform_inference(self, image: list | np.ndarray) -> None:
    """Prediction is performed using self.model and the prediction result is set to self._original_predictions.

    Args:
        image: np.ndarray
            A numpy array that contains the image to be predicted. 3 channel image should be in RGB order.
    """
    import torch

    # Confirm model is loaded
    if self.model is None or self.processor is None:
        raise RuntimeError("Model is not loaded, load it by calling .load_model()")

    with torch.no_grad():
        if self._is_zero_shot_model:
            text = self._get_zero_shot_text_input(len(image) if isinstance(image, list) else 1)
            inputs = self.processor(images=image, text=text, return_tensors="pt")
        else:
            inputs = self.processor(images=image, return_tensors="pt")
        inputs = {k: v.to(self.device) if hasattr(v, "to") else v for k, v in inputs.items()}
        outputs = self.model(**inputs)
    self._original_input_ids = inputs.get("input_ids")

    images = image if isinstance(image, list) else [image]
    self._original_shapes = [img.shape for img in images]
    self._original_predictions = outputs

def perform_batch_inference(self, images: list[np.ndarray]) -> None:
    """Native batch inference: process all images in a single processor + model call.

    Unlike the base-class default (which runs images sequentially), this
    feeds the entire list to the HuggingFace processor at once and executes
    one batched forward pass.  The processor pads images to a uniform size
    internally, so images of different resolutions are handled correctly.

    This avoids setting ``_batch_images`` so
    ``convert_original_predictions`` uses the standard multi-image path
    rather than the sequential fallback.

    Args:
        images: List of numpy arrays (H, W, C) in RGB order.
    """
    self.perform_inference(images)

def get_valid_predictions(self, logits: Any, pred_boxes: Any) -> tuple:
    """Get predictions above confidence threshold.

    Args:
        logits: torch.Tensor
        pred_boxes: torch.Tensor

    Returns:
        scores: torch.Tensor
        cat_ids: torch.Tensor
        boxes: torch.Tensor
    """
    import torch

    if self._uses_sigmoid_cls:
        # RT-DETR family: per-class sigmoid, logits shape (Q, num_classes) -- no background class
        probs = logits.sigmoid()
        scores, cat_ids = probs.max(-1)
        valid_mask = scores >= self.confidence_threshold
    else:
        # DETR family: softmax over (num_classes + 1), last index is no-object/background
        probs = logits.softmax(-1)
        scores = probs.max(-1).values
        cat_ids = probs.argmax(-1)
        valid_detections = torch.where(cat_ids < self.num_categories, 1, 0)
        valid_confidences = torch.where(scores >= self.confidence_threshold, 1, 0)
        valid_mask = valid_detections.logical_and(valid_confidences).bool()

    scores = scores[valid_mask]
    cat_ids = cat_ids[valid_mask]
    boxes = pred_boxes[valid_mask]
    return scores, cat_ids, boxes