Filter detections classes

[techfriendly-dev]

How can I implement a class filter with SAHI? Like is done in YOLO function like this example

from ultralytics.yolo.engine.model import YOLO

model = YOLO("yolov8n.pt")
model.predict(source="0", show=True, stream=True, **classes=0**)

Thanks in advance

[fcakyon]

This should work:

from ultralytics import YOLO
from sahi import AutoDetectionModel

model = YOLO('custom_yolov8_model.pt")
model.overrides['classes'] = [0, 15, 16]

sahi_model = AutoDetectionModel(
    model_type='yolov8',
    model=model,
)

# use `sahi_model` with `get_prediction` or `get_sliced_prediction`

[techfriendly-dev]

Is working! Thank you for the tip

I didn't realize that AutoDetectionModel.from_pretrained has model: Optional[Any] = None, optional arg input.

class AutoDetectionModel:
    @staticmethod
    def from_pretrained(
        model_type: str,
        model_path: Optional[str] = None,
        model: Optional[Any] = None,
        config_path: Optional[str] = None,
        device: Optional[str] = None,
        mask_threshold: float = 0.5,
        confidence_threshold: float = 0.3,
        category_mapping: Optional[Dict] = None,
        category_remapping: Optional[Dict] = None,
        load_at_init: bool = True,
        image_size: int = None,
        **kwargs,
    ):

Anyway, Is not described at comment, that is the confusion

  Loads a DetectionModel from given path.
        Args:
            model_type: str
                Name of the detection framework (example: "yolov5", "mmdet", "detectron2")
            model_path: str
                Path of the detection model (ex. 'model.pt')
            config_path: str
                Path of the config file (ex. 'mmdet/configs/cascade_rcnn_r50_fpn_1x.py')
            device: str
                Device, "cpu" or "cuda:0"
            mask_threshold: float
                Value to threshold mask pixels, should be between 0 and 1
            confidence_threshold: float
                All predictions with score < confidence_threshold will be discarded
            category_mapping: dict: str to str
                Mapping from category id (str) to category name (str) e.g. {"1": "pedestrian"}
            category_remapping: dict: str to int
                Remap category ids based on category names, after performing inference e.g. {"car": 3}
            load_at_init: bool
                If True, automatically loads the model at initalization
            image_size: int
                Inference input size.
        Returns:
            Returns an instance of a DetectionModel
        Raises:
            ImportError: If given {model_type} framework is not installed