Base Model¶

def check_dependencies(self, packages: list[str] | None = None) -> None:
    """Ensures required dependencies are installed.

    If 'packages' is None, uses self.required_packages. Subclasses may still call with a custom list for dynamic
    needs.
    """
    pkgs = packages if packages is not None else getattr(self, "required_packages", [])
    if pkgs:
        check_requirements(pkgs)

def load_model(self) -> None:
    """Load the detection model from disk and assign it to ``self.model``.

    Subclasses must override this method. The implementation should use
    ``self.model_path``, ``self.config_path``, and ``self.device`` to
    construct the underlying model object and store it in ``self.model``.
    """
    raise NotImplementedError()

def set_model(self, model: Any, **kwargs: Any) -> None:
    """Set an already-instantiated model as the underlying detection model.

    Subclasses must override this method to assign ``model`` to
    ``self.model`` and perform any additional setup (e.g. category mapping).

    Args:
        model: Any
            A pre-loaded detection model instance.
        **kwargs: Any
            Additional keyword arguments for subclass-specific setup.
    """
    raise NotImplementedError()

def set_device(self, device: str | None = None) -> None:
    """Sets the device pytorch should use for the model.

    Args:
        device: Torch device, "cpu", "mps", "cuda", "cuda:0", "cuda:1", etc.
    """
    self.device = select_device(device)

def unload_model(self) -> None:
    """Unloads the model from CPU/GPU."""
    self.model = None
    empty_cuda_cache()

def perform_inference(self, image: np.ndarray) -> None:
    """Run inference on a single image and store raw predictions.

    Subclasses must override this method. The implementation should run
    the model on ``image`` and assign the raw results to
    ``self._original_predictions``.

    Args:
        image: np.ndarray
            A numpy array (H, W, C) containing the image to run inference on.
    """
    raise NotImplementedError()

def perform_batch_inference(self, images: list[np.ndarray]) -> None:
    """Performs inference on a batch of images.

    Subclasses can override this for native batch support (e.g.
    ``UltralyticsDetectionModel`` passes the full list to YOLO for
    true GPU batching, ``HuggingfaceDetectionModel`` feeds all images
    to the processor in one call).

    The default does **not** run inference here.  It stores images so
    that ``convert_original_predictions`` can call ``perform_inference``
    per image, preserving each model's ``_original_predictions`` format.
    Subclasses with native batch support override this to run inference
    immediately.

    Args:
        images: list[np.ndarray]
            List of numpy arrays (H, W, C) to run inference on.
    """
    self._batch_images = images
    self._original_shapes = [img.shape for img in images]

def convert_original_predictions(
    self,
    shift_amount: list[list[int | float]] | None = [[0, 0]],
    full_shape: list[list[int | float]] | None = None,
) -> None:
    """Convert raw predictions to ObjectPrediction lists.

    Should be called after ``perform_inference`` or ``perform_batch_inference``.

    When the default (sequential) ``perform_batch_inference`` was used,
    this method runs inference + conversion one image at a time so that
    each model's internal ``_original_predictions`` format is preserved.

    Args:
        shift_amount: Per-image shift amounts ``[[shift_x, shift_y], ...]``
            or a single ``[shift_x, shift_y]`` for one image.
        full_shape: Per-image full image sizes ``[[height, width], ...]``
            or a single ``[height, width]`` for one image.
    """
    batch_images = getattr(self, "_batch_images", None)
    if batch_images is not None:
        from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list

        shift_amount_list = fix_shift_amount_list(shift_amount)
        full_shape_list = fix_full_shape_list(full_shape)

        all_preds: list[list[ObjectPrediction]] = []
        for i, image in enumerate(batch_images):
            self.perform_inference(np.ascontiguousarray(image))
            sa: list[list[int | float]] = [shift_amount_list[i]] if shift_amount_list else [[0, 0]]
            fs: list[list[int | float]] | None = [full_shape_list[i]] if full_shape_list else None
            self._create_object_prediction_list_from_original_predictions(
                shift_amount_list=sa,
                full_shape_list=fs,
            )
            if self.category_remapping:
                self._apply_category_remapping()
            all_preds.extend(self._object_prediction_list_per_image or [])
        self._object_prediction_list_per_image = all_preds
        self._batch_images = None  # clear deferred state
        return

    # Standard single-image path
    self._create_object_prediction_list_from_original_predictions(
        shift_amount_list=shift_amount,
        full_shape_list=full_shape,
    )
    if self.category_remapping:
        self._apply_category_remapping()