Skip to content

RoboflowModel

sahi.models.roboflow

Classes

RoboflowDetectionModel

Bases: DetectionModel

Source code in sahi/models/roboflow.py 
 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
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
class RoboflowDetectionModel(DetectionModel):
    def __init__(
        self,
        model: Optional[Any] = None,
        model_path: Optional[str] = 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: Optional[int] = None,
        api_key: Optional[str] = None,
    ):
        """
        Initialize the RoboflowDetectionModel with the given parameters.

        Args:
            model_path: str
                Path for the instance segmentation model weight
            config_path: str
                Path for the mmdetection instance segmentation model config file
            device: Torch device, "cpu", "mps", "cuda", "cuda:0", "cuda:1", etc.
            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 initialization
            image_size: int
                Inference input size.
        """
        self._use_universe = model and isinstance(model, str)
        self._model = model
        self._device = device
        self._api_key = api_key

        super().__init__(
            model=model,
            model_path=model_path,
            config_path=config_path,
            device=device,
            mask_threshold=mask_threshold,
            confidence_threshold=confidence_threshold,
            category_mapping=category_mapping,
            category_remapping=category_remapping,
            load_at_init=False,
            image_size=image_size,
        )

        if load_at_init:
            self.load_model()

    def check_dependencies(self) -> None:
        """
        This function can be implemented to ensure model dependencies are installed.
        """
        if self._use_universe:
            check_requirements(["inference"])
        else:
            check_requirements(["rfdetr"])

    def set_model(self, model: Any, **kwargs):
        """
        This function should be implemented to instantiate a DetectionModel out of an already loaded model
        Args:
            model: Any
                Loaded model
        """
        self.model = model

    def load_model(self):
        """
        This function should be implemented in a way that detection model
        should be initialized and set to self.model.
        (self.model_path, self.config_path, and self.device should be utilized)
        """
        if self._use_universe:
            from inference import get_model
            from inference.core.env import API_KEY
            from inference.core.exceptions import RoboflowAPINotAuthorizedError

            api_key = self._api_key or API_KEY

            try:
                model = get_model(self._model, api_key=api_key)
            except RoboflowAPINotAuthorizedError as e:
                raise ValueError(
                    "Authorization failed. Please pass a valid API key with the `api_key` parameter or set the `ROBOFLOW_API_KEY` environment variable."
                ) from e

            assert model.task_type == "object-detection", "Roboflow model must be an object detection model."

        else:
            from rfdetr import RFDETRBase, RFDETRLarge

            model, model_path = self._model, self.model_path
            model_names = ("RFDETRBase", "RFDETRLarge")
            if hasattr(model, "__name__") and model.__name__ in model_names:
                model_params = dict(
                    resolution=int(self.image_size) if self.image_size else 560,
                    device=self._device,
                    num_classes=len(self.category_mapping.keys()) if self.category_mapping else None,
                )
                if model_path:
                    model_params["pretrain_weights"] = model_path

                model = model(**model_params)
            elif isinstance(model, (RFDETRBase, RFDETRLarge)):
                model = model
            else:
                raise ValueError(
                    f"Model must be a Roboflow model string or one of {model_names} models, got {self.model}."
                )

        self.set_model(model)

    def perform_inference(
        self,
        image: np.ndarray,
    ):
        """
        This function should be implemented in a way that prediction should be
        performed using self.model and the prediction result should be set to self._original_predictions.
        Args:
            image: np.ndarray
                A numpy array that contains the image to be predicted.
        """
        if self._use_universe:
            self._original_predictions = self.model.infer(image, confidence=self.confidence_threshold)
        else:
            self._original_predictions = [self.model.predict(image, threshold=self.confidence_threshold)]

    def _create_object_prediction_list_from_original_predictions(
        self,
        shift_amount_list: Optional[List[List[int]]] = [[0, 0]],
        full_shape_list: Optional[List[List[int]]] = None,
    ):
        """
        This function should be implemented in a way that self._original_predictions should
        be converted to a list of prediction.ObjectPrediction and set to
        self._object_prediction_list. self.mask_threshold can also be utilized.
        Args:
            shift_amount_list: list of list
                To shift the box and mask predictions from sliced image to full sized image, should
                be in the form of List[[shift_x, shift_y],[shift_x, shift_y],...]
            full_shape_list: list of list
                Size of the full image after shifting, should be in the form of
                List[[height, width],[height, width],...]
        """
        # compatibility for sahi v0.8.15
        shift_amount_list = fix_shift_amount_list(shift_amount_list)
        full_shape_list = fix_full_shape_list(full_shape_list)

        object_prediction_list: List[ObjectPrediction] = []

        if self._use_universe:
            from inference.core.entities.responses.inference import (
                ObjectDetectionInferenceResponse as InferenceObjectDetectionInferenceResponse,
            )
            from inference.core.entities.responses.inference import (
                ObjectDetectionPrediction as InferenceObjectDetectionPrediction,
            )

            original_reponses: List[InferenceObjectDetectionInferenceResponse] = self._original_predictions

            assert len(original_reponses) == len(shift_amount_list) == len(full_shape_list), (
                "Length mismatch between original responses, shift amounts, and full shapes."
            )

            for original_reponse, shift_amount, full_shape in zip(
                original_reponses,
                shift_amount_list,
                full_shape_list,
            ):
                for prediction in original_reponse.predictions:
                    prediction: InferenceObjectDetectionPrediction
                    bbox = [
                        prediction.x - prediction.width / 2,
                        prediction.y - prediction.height / 2,
                        prediction.x + prediction.width / 2,
                        prediction.y + prediction.height / 2,
                    ]
                    object_prediction = ObjectPrediction(
                        bbox=bbox,
                        category_id=prediction.class_id,
                        category_name=prediction.class_name,
                        score=prediction.confidence,
                        shift_amount=shift_amount,
                        full_shape=full_shape,
                    )
                    object_prediction_list.append(object_prediction)

        else:
            from supervision.detection.core import Detections

            original_detections: List[Detections] = self._original_predictions

            assert len(original_detections) == len(shift_amount_list) == len(full_shape_list), (
                "Length mismatch between original responses, shift amounts, and full shapes."
            )

            for original_detection, shift_amount, full_shape in zip(
                original_detections,
                shift_amount_list,
                full_shape_list,
            ):
                for xyxy, confidence, class_id in zip(
                    original_detection.xyxy,
                    original_detection.confidence,
                    original_detection.class_id,
                ):
                    object_prediction = ObjectPrediction(
                        bbox=xyxy,
                        category_id=int(class_id),
                        category_name=self.category_mapping.get(int(class_id), None),
                        score=float(confidence),
                        shift_amount=shift_amount,
                        full_shape=full_shape,
                    )
                    object_prediction_list.append(object_prediction)

        object_prediction_list_per_image = [object_prediction_list]
        self._object_prediction_list_per_image = object_prediction_list_per_image
Functions
__init__(model=None, model_path=None, config_path=None, device=None, mask_threshold=0.5, confidence_threshold=0.3, category_mapping=None, category_remapping=None, load_at_init=True, image_size=None, api_key=None)

Initialize the RoboflowDetectionModel with the given parameters.

Parameters:

Name Type Description Default
model_path Optional[str]

str Path for the instance segmentation model weight

 None
config_path Optional[str]

str Path for the mmdetection instance segmentation model config file

 None
device Optional[str]

Torch device, "cpu", "mps", "cuda", "cuda:0", "cuda:1", etc.

 None
mask_threshold float

float Value to threshold mask pixels, should be between 0 and 1

 0.5
confidence_threshold float

float All predictions with score < confidence_threshold will be discarded

 0.3
category_mapping Optional[Dict]

dict: str to str Mapping from category id (str) to category name (str) e.g. {"1": "pedestrian"}

 None
category_remapping Optional[Dict]

dict: str to int Remap category ids based on category names, after performing inference e.g. {"car": 3}

 None
load_at_init bool

bool If True, automatically loads the model at initialization

 True
image_size Optional[int]

int Inference input size.

 None
Source code in sahi/models/roboflow.py 
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
58
59
60
61
62
63
64
65
66
67
68
69
70
def __init__(
    self,
    model: Optional[Any] = None,
    model_path: Optional[str] = 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: Optional[int] = None,
    api_key: Optional[str] = None,
):
    """
    Initialize the RoboflowDetectionModel with the given parameters.

    Args:
        model_path: str
            Path for the instance segmentation model weight
        config_path: str
            Path for the mmdetection instance segmentation model config file
        device: Torch device, "cpu", "mps", "cuda", "cuda:0", "cuda:1", etc.
        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 initialization
        image_size: int
            Inference input size.
    """
    self._use_universe = model and isinstance(model, str)
    self._model = model
    self._device = device
    self._api_key = api_key

    super().__init__(
        model=model,
        model_path=model_path,
        config_path=config_path,
        device=device,
        mask_threshold=mask_threshold,
        confidence_threshold=confidence_threshold,
        category_mapping=category_mapping,
        category_remapping=category_remapping,
        load_at_init=False,
        image_size=image_size,
    )

    if load_at_init:
        self.load_model()
check_dependencies()

This function can be implemented to ensure model dependencies are installed.

Source code in sahi/models/roboflow.py 
72
73
74
75
76
77
78
79
def check_dependencies(self) -> None:
    """
    This function can be implemented to ensure model dependencies are installed.
    """
    if self._use_universe:
        check_requirements(["inference"])
    else:
        check_requirements(["rfdetr"])
load_model()

This function should be implemented in a way that detection model should be initialized and set to self.model. (self.model_path, self.config_path, and self.device should be utilized)

Source code in sahi/models/roboflow.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def load_model(self):
    """
    This function should be implemented in a way that detection model
    should be initialized and set to self.model.
    (self.model_path, self.config_path, and self.device should be utilized)
    """
    if self._use_universe:
        from inference import get_model
        from inference.core.env import API_KEY
        from inference.core.exceptions import RoboflowAPINotAuthorizedError

        api_key = self._api_key or API_KEY

        try:
            model = get_model(self._model, api_key=api_key)
        except RoboflowAPINotAuthorizedError as e:
            raise ValueError(
                "Authorization failed. Please pass a valid API key with the `api_key` parameter or set the `ROBOFLOW_API_KEY` environment variable."
            ) from e

        assert model.task_type == "object-detection", "Roboflow model must be an object detection model."

    else:
        from rfdetr import RFDETRBase, RFDETRLarge

        model, model_path = self._model, self.model_path
        model_names = ("RFDETRBase", "RFDETRLarge")
        if hasattr(model, "__name__") and model.__name__ in model_names:
            model_params = dict(
                resolution=int(self.image_size) if self.image_size else 560,
                device=self._device,
                num_classes=len(self.category_mapping.keys()) if self.category_mapping else None,
            )
            if model_path:
                model_params["pretrain_weights"] = model_path

            model = model(**model_params)
        elif isinstance(model, (RFDETRBase, RFDETRLarge)):
            model = model
        else:
            raise ValueError(
                f"Model must be a Roboflow model string or one of {model_names} models, got {self.model}."
            )

    self.set_model(model)
perform_inference(image)

This function should be implemented in a way that prediction should be performed using self.model and the prediction result should be set to self._original_predictions. Args: image: np.ndarray A numpy array that contains the image to be predicted.

Source code in sahi/models/roboflow.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def perform_inference(
    self,
    image: np.ndarray,
):
    """
    This function should be implemented in a way that prediction should be
    performed using self.model and the prediction result should be set to self._original_predictions.
    Args:
        image: np.ndarray
            A numpy array that contains the image to be predicted.
    """
    if self._use_universe:
        self._original_predictions = self.model.infer(image, confidence=self.confidence_threshold)
    else:
        self._original_predictions = [self.model.predict(image, threshold=self.confidence_threshold)]
set_model(model, **kwargs)

This function should be implemented to instantiate a DetectionModel out of an already loaded model Args: model: Any Loaded model

Source code in sahi/models/roboflow.py 
81
82
83
84
85
86
87
88
def set_model(self, model: Any, **kwargs):
    """
    This function should be implemented to instantiate a DetectionModel out of an already loaded model
    Args:
        model: Any
            Loaded model
    """
    self.model = model

Functions