Skip to content

Pipelines

ClassificationPipeline

Bases: Pipeline

Pipeline specifically designed for Image Classification tasks. The pipeline includes loading images, preprocessing, prediction, and postprocessing.

Parameters:

Name Type Description Default
model SgModule

The classification model (instance of SgModule) used for making predictions.

required
class_names List[str]

List of class names corresponding to the model's output classes.

required
image_processor Optional[Processing]

Single image processor or a list of image processors for preprocessing and postprocessing the images.

None
device Optional[str]

The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.

None
fuse_model bool

If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.

True
Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
class ClassificationPipeline(Pipeline):
    """Pipeline specifically designed for Image Classification tasks.
    The pipeline includes loading images, preprocessing, prediction, and postprocessing.

    :param model:                       The classification model (instance of SgModule) used for making predictions.
    :param class_names:                 List of class names corresponding to the model's output classes.
    :param image_processor:             Single image processor or a list of image processors for preprocessing and postprocessing the images.
    :param device:                      The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.
    :param fuse_model:                  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
    """

    def __init__(
        self,
        model: SgModule,
        class_names: List[str],
        device: Optional[str] = None,
        image_processor: Optional[Processing] = None,
        fuse_model: bool = True,
    ):
        super().__init__(model=model, device=device, image_processor=image_processor, class_names=class_names, fuse_model=fuse_model)

    def _decode_model_output(self, model_output: Union[List, Tuple, torch.Tensor], model_input: np.ndarray) -> List[ClassificationPrediction]:
        """Decode the model output

        :param model_output:    Direct output of the model, without any post-processing. Tensor of shape [B, C]
        :param model_input:     Model input (i.e. images after preprocessing).
        :return:                Predicted Bboxes.
        """
        pred_scores, pred_labels = torch.max(model_output.softmax(dim=1), 1)

        pred_labels = pred_labels.detach().cpu().numpy()  # [B,1]
        pred_scores = pred_scores.detach().cpu().numpy()  # [B,1]

        predictions = list()
        for prediction, confidence, image_input in zip(pred_labels, pred_scores, model_input):
            predictions.append(ClassificationPrediction(confidence=float(confidence), label=int(prediction), image_shape=image_input.shape))
        return predictions

    def _instantiate_image_prediction(self, image: np.ndarray, prediction: ClassificationPrediction) -> ImagePrediction:
        return ImageClassificationPrediction(image=image, prediction=prediction, class_names=self.class_names)

    def _combine_image_prediction_to_images(
        self, images_predictions: Iterable[ImageClassificationPrediction], n_images: Optional[int] = None
    ) -> ImagesClassificationPrediction:
        if n_images is not None and n_images == 1:
            # Do not show tqdm progress bar if there is only one image
            images_predictions = [next(iter(images_predictions))]
        else:
            images_predictions = [image_predictions for image_predictions in tqdm(images_predictions, total=n_images, desc="Predicting Images")]

        return ImagesClassificationPrediction(_images_prediction_lst=images_predictions)

    def _combine_image_prediction_to_video(
        self, images_predictions: Iterable[ImageDetectionPrediction], fps: float, n_images: Optional[int] = None
    ) -> ImagesClassificationPrediction:
        raise NotImplementedError("This feature is not available for Classification task")

DetectionPipeline

Bases: Pipeline

Pipeline specifically designed for object detection tasks. The pipeline includes loading images, preprocessing, prediction, and postprocessing.

Parameters:

Name Type Description Default
model SgModule

The object detection model (instance of SgModule) used for making predictions.

required
class_names List[str]

List of class names corresponding to the model's output classes.

required
post_prediction_callback DetectionPostPredictionCallback

Callback function to process raw predictions from the model.

required
image_processor Optional[Processing]

Single image processor or a list of image processors for preprocessing and postprocessing the images.

None
device Optional[str]

The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.

None
fuse_model bool

If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.

True
Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
class DetectionPipeline(Pipeline):
    """Pipeline specifically designed for object detection tasks.
    The pipeline includes loading images, preprocessing, prediction, and postprocessing.

    :param model:                       The object detection model (instance of SgModule) used for making predictions.
    :param class_names:                 List of class names corresponding to the model's output classes.
    :param post_prediction_callback:    Callback function to process raw predictions from the model.
    :param image_processor:             Single image processor or a list of image processors for preprocessing and postprocessing the images.
    :param device:                      The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.
    :param fuse_model:                  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
    """

    def __init__(
        self,
        model: SgModule,
        class_names: List[str],
        post_prediction_callback: DetectionPostPredictionCallback,
        device: Optional[str] = None,
        image_processor: Optional[Processing] = None,
        fuse_model: bool = True,
    ):
        super().__init__(model=model, device=device, image_processor=image_processor, class_names=class_names, fuse_model=fuse_model)
        self.post_prediction_callback = post_prediction_callback

    def _decode_model_output(self, model_output: Union[List, Tuple, torch.Tensor], model_input: np.ndarray) -> List[DetectionPrediction]:
        """Decode the model output, by applying post prediction callback. This includes NMS.

        :param model_output:    Direct output of the model, without any post-processing.
        :param model_input:     Model input (i.e. images after preprocessing).
        :return:                Predicted Bboxes.
        """
        post_nms_predictions = self.post_prediction_callback(model_output, device=self.device)

        predictions = []
        for prediction, image in zip(post_nms_predictions, model_input):
            prediction = prediction if prediction is not None else torch.zeros((0, 6), dtype=torch.float32)
            prediction = prediction.detach().cpu().numpy()
            predictions.append(
                DetectionPrediction(
                    bboxes=prediction[:, :4],
                    confidence=prediction[:, 4],
                    labels=prediction[:, 5],
                    bbox_format="xyxy",
                    image_shape=image.shape,
                )
            )

        return predictions

    def _instantiate_image_prediction(self, image: np.ndarray, prediction: DetectionPrediction) -> ImagePrediction:
        return ImageDetectionPrediction(image=image, prediction=prediction, class_names=self.class_names)

    def _combine_image_prediction_to_images(
        self, images_predictions: Iterable[ImageDetectionPrediction], n_images: Optional[int] = None
    ) -> ImagesDetectionPrediction:
        if n_images is not None and n_images == 1:
            # Do not show tqdm progress bar if there is only one image
            images_predictions = [next(iter(images_predictions))]
        else:
            images_predictions = [image_predictions for image_predictions in tqdm(images_predictions, total=n_images, desc="Predicting Images")]

        return ImagesDetectionPrediction(_images_prediction_lst=images_predictions)

    def _combine_image_prediction_to_video(
        self, images_predictions: Iterable[ImageDetectionPrediction], fps: float, n_images: Optional[int] = None
    ) -> VideoDetectionPrediction:
        images_predictions = [image_predictions for image_predictions in tqdm(images_predictions, total=n_images, desc="Predicting Video")]
        return VideoDetectionPrediction(_images_prediction_lst=images_predictions, fps=fps)

Pipeline

Bases: ABC

An abstract base class representing a processing pipeline for a specific task. The pipeline includes loading images, preprocessing, prediction, and postprocessing.

Parameters:

Name Type Description Default
model SgModule

The model used for making predictions.

required
image_processor Union[Processing, List[Processing]]

A single image processor or a list of image processors for preprocessing and postprocessing the images.

required
device Optional[str]

The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.

None
dtype Optional[torch.dtype]

Specify the dtype of the inputs. If None, will use the dtype of the model's parameters.

None
fuse_model bool

If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.

True
Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
 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
243
244
245
246
247
248
249
250
251
class Pipeline(ABC):
    """An abstract base class representing a processing pipeline for a specific task.
    The pipeline includes loading images, preprocessing, prediction, and postprocessing.

    :param model:           The model used for making predictions.
    :param image_processor: A single image processor or a list of image processors for preprocessing and postprocessing the images.
    :param device:          The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.
    :param dtype:           Specify the dtype of the inputs. If None, will use the dtype of the model's parameters.
    :param fuse_model:                  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
    """

    def __init__(
        self,
        model: SgModule,
        image_processor: Union[Processing, List[Processing]],
        class_names: List[str],
        device: Optional[str] = None,
        fuse_model: bool = True,
        dtype: Optional[torch.dtype] = None,
    ):
        model_device: torch.device = infer_model_device(model=model)
        if device:
            device: torch.device = resolve_torch_device(device=device)

        self.device: torch.device = device or model_device
        self.dtype = dtype or next(model.parameters()).dtype
        self.model = model.to(device) if device and device != model_device else model
        self.class_names = class_names

        if isinstance(image_processor, list):
            image_processor = ComposeProcessing(image_processor)
        self.image_processor = image_processor

        self.fuse_model = fuse_model  # If True, the model will be fused in the first forward pass, to make sure it gets the right input_size

    def _fuse_model(self, input_example: torch.Tensor):
        logger.info("Fusing some of the model's layers. If this takes too much memory, you can deactivate it by setting `fuse_model=False`")
        self.model = copy.deepcopy(self.model)
        self.model.eval()
        self.model.prep_model_for_conversion(input_size=input_example.shape[-2:])
        self.fuse_model = False

    def __call__(self, inputs: Union[str, ImageSource, List[ImageSource]], batch_size: Optional[int] = 32) -> ImagesPredictions:
        """Predict an image or a list of images.

        Supported types include:
            - str:              A string representing either a video, an image or an URL.
            - numpy.ndarray:    A numpy array representing the image
            - torch.Tensor:     A PyTorch tensor representing the image
            - PIL.Image.Image:  A PIL Image object
            - List:             A list of images of any of the above image types (list of videos not supported).

        :param inputs:      inputs to the model, which can be any of the above-mentioned types.
        :param batch_size:  Maximum number of images to process at the same time.
        :return:            Results of the prediction.
        """

        if includes_video_extension(inputs):
            return self.predict_video(inputs, batch_size)
        elif check_image_typing(inputs):
            return self.predict_images(inputs, batch_size)
        else:
            raise ValueError(f"Input {inputs} not supported for prediction.")

    def predict_images(self, images: Union[ImageSource, List[ImageSource]], batch_size: Optional[int] = 32) -> ImagesPredictions:
        """Predict an image or a list of images.

        :param images:      Images to predict.
        :param batch_size:  The size of each batch.
        :return:            Results of the prediction.
        """
        from super_gradients.training.utils.media.image import load_images

        images = load_images(images)

        result_generator = self._generate_prediction_result(images=images, batch_size=batch_size)
        return self._combine_image_prediction_to_images(result_generator, n_images=len(images))

    def predict_video(self, video_path: str, batch_size: Optional[int] = 32) -> VideoPredictions:
        """Predict on a video file, by processing the frames in batches.

        :param video_path:  Path to the video file.
        :param batch_size:  The size of each batch.
        :return:            Results of the prediction.
        """
        video_frames, fps = load_video(file_path=video_path)
        result_generator = self._generate_prediction_result(images=video_frames, batch_size=batch_size)
        return self._combine_image_prediction_to_video(result_generator, fps=fps, n_images=len(video_frames))

    def predict_webcam(self) -> None:
        """Predict using webcam"""

        def _draw_predictions(frame: np.ndarray) -> np.ndarray:
            """Draw the predictions on a single frame from the stream."""
            frame_prediction = next(iter(self._generate_prediction_result(images=[frame])))
            return frame_prediction.draw()

        video_streaming = WebcamStreaming(frame_processing_fn=_draw_predictions, fps_update_frequency=1)
        video_streaming.run()

    def _generate_prediction_result(self, images: Iterable[np.ndarray], batch_size: Optional[int] = None) -> Iterable[ImagePrediction]:
        """Run the pipeline on the images as single batch or through multiple batches.

        NOTE: A core motivation to have this function as a generator is that it can be used in a lazy way (if images is generator itself),
              i.e. without having to load all the images into memory.

        :param images:      Iterable of numpy arrays representing images.
        :param batch_size:  The size of each batch.
        :return:            Iterable of Results object, each containing the results of the prediction and the image.
        """
        if batch_size is None:
            yield from self._generate_prediction_result_single_batch(images)
        else:
            for batch_images in generate_batch(images, batch_size):
                yield from self._generate_prediction_result_single_batch(batch_images)

    def _generate_prediction_result_single_batch(self, images: Iterable[np.ndarray]) -> Iterable[ImagePrediction]:
        """Run the pipeline on images. The pipeline is made of 4 steps:
            1. Load images - Loading the images into a list of numpy arrays.
            2. Preprocess - Encode the image in the shape/format expected by the model
            3. Predict - Run the model on the preprocessed image
            4. Postprocess - Decode the output of the model so that the predictions are in the shape/format of original image.

        :param images:  Iterable of numpy arrays representing images.
        :return:        Iterable of Results object, each containing the results of the prediction and the image.
        """
        # Make sure the model is on the correct device, as it might have been moved after init
        model_device: torch.device = infer_model_device(model=self.model)
        if self.device != model_device:
            self.model = self.model.to(self.device)

        images = list(images)  # We need to load all the images into memory, and to reuse it afterwards.

        # Preprocess
        preprocessed_images, processing_metadatas = [], []
        for image in images:
            preprocessed_image, processing_metadata = self.image_processor.preprocess_image(image=image.copy())
            preprocessed_images.append(preprocessed_image)
            processing_metadatas.append(processing_metadata)

        # Predict
        with eval_mode(self.model), torch.no_grad(), torch.cuda.amp.autocast():
            torch_inputs = torch.from_numpy(np.array(preprocessed_images)).to(self.device)
            torch_inputs = torch_inputs.to(self.dtype)
            if self.fuse_model:
                self._fuse_model(torch_inputs)
            model_output = self.model(torch_inputs)
            predictions = self._decode_model_output(model_output, model_input=torch_inputs)

        # Postprocess
        postprocessed_predictions = []
        for image, prediction, processing_metadata in zip(images, predictions, processing_metadatas):
            prediction = self.image_processor.postprocess_predictions(predictions=prediction, metadata=processing_metadata)
            postprocessed_predictions.append(prediction)

        # Yield results one by one
        for image, prediction in zip(images, postprocessed_predictions):
            yield self._instantiate_image_prediction(image=image, prediction=prediction)

    @abstractmethod
    def _decode_model_output(self, model_output: Union[List, Tuple, torch.Tensor], model_input: np.ndarray) -> List[Prediction]:
        """Decode the model outputs, move each prediction to numpy and store it in a Prediction object.

        :param model_output:    Direct output of the model, without any post-processing.
        :param model_input:     Model input (i.e. images after preprocessing).
        :return:                Model predictions, without any post-processing.
        """
        raise NotImplementedError

    @abstractmethod
    def _instantiate_image_prediction(self, image: np.ndarray, prediction: Prediction) -> ImagePrediction:
        """Instantiate an object wrapping an image and the pipeline's prediction.

        :param image:       Image to predict.
        :param prediction:  Model prediction on that image.
        :return:            Object wrapping an image and the pipeline's prediction.
        """
        raise NotImplementedError

    @abstractmethod
    def _combine_image_prediction_to_images(self, images_prediction_lst: Iterable[ImagePrediction], n_images: Optional[int] = None) -> ImagesPredictions:
        """Instantiate an object wrapping the list of images and the pipeline's predictions on them.

        :param images_prediction_lst:   List of image predictions.
        :param n_images:                (Optional) Number of images in the list. This used for tqdm progress bar to work with iterables, but is not required.
        :return:                        Object wrapping the list of image predictions.
        """
        raise NotImplementedError

    @abstractmethod
    def _combine_image_prediction_to_video(
        self, images_prediction_lst: Iterable[ImagePrediction], fps: float, n_images: Optional[int] = None
    ) -> VideoPredictions:
        """Instantiate an object holding the video frames and the pipeline's predictions on it.

        :param images_prediction_lst:   List of image predictions.
        :param fps:                     Frames per second.
        :param n_images:                (Optional) Number of images in the list. This used for tqdm progress bar to work with iterables, but is not required.
        :return:                        Object wrapping the list of image predictions as a Video.
        """
        raise NotImplementedError

__call__(inputs, batch_size=32)

Predict an image or a list of images.

Supported types include: - str: A string representing either a video, an image or an URL. - numpy.ndarray: A numpy array representing the image - torch.Tensor: A PyTorch tensor representing the image - PIL.Image.Image: A PIL Image object - List: A list of images of any of the above image types (list of videos not supported).

Parameters:

Name Type Description Default
inputs Union[str, ImageSource, List[ImageSource]]

inputs to the model, which can be any of the above-mentioned types.

required
batch_size Optional[int]

Maximum number of images to process at the same time.

32

Returns:

Type Description
ImagesPredictions

Results of the prediction.

Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def __call__(self, inputs: Union[str, ImageSource, List[ImageSource]], batch_size: Optional[int] = 32) -> ImagesPredictions:
    """Predict an image or a list of images.

    Supported types include:
        - str:              A string representing either a video, an image or an URL.
        - numpy.ndarray:    A numpy array representing the image
        - torch.Tensor:     A PyTorch tensor representing the image
        - PIL.Image.Image:  A PIL Image object
        - List:             A list of images of any of the above image types (list of videos not supported).

    :param inputs:      inputs to the model, which can be any of the above-mentioned types.
    :param batch_size:  Maximum number of images to process at the same time.
    :return:            Results of the prediction.
    """

    if includes_video_extension(inputs):
        return self.predict_video(inputs, batch_size)
    elif check_image_typing(inputs):
        return self.predict_images(inputs, batch_size)
    else:
        raise ValueError(f"Input {inputs} not supported for prediction.")

predict_images(images, batch_size=32)

Predict an image or a list of images.

Parameters:

Name Type Description Default
images Union[ImageSource, List[ImageSource]]

Images to predict.

required
batch_size Optional[int]

The size of each batch.

32

Returns:

Type Description
ImagesPredictions

Results of the prediction.

Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
115
116
117
118
119
120
121
122
123
124
125
126
127
def predict_images(self, images: Union[ImageSource, List[ImageSource]], batch_size: Optional[int] = 32) -> ImagesPredictions:
    """Predict an image or a list of images.

    :param images:      Images to predict.
    :param batch_size:  The size of each batch.
    :return:            Results of the prediction.
    """
    from super_gradients.training.utils.media.image import load_images

    images = load_images(images)

    result_generator = self._generate_prediction_result(images=images, batch_size=batch_size)
    return self._combine_image_prediction_to_images(result_generator, n_images=len(images))

predict_video(video_path, batch_size=32)

Predict on a video file, by processing the frames in batches.

Parameters:

Name Type Description Default
video_path str

Path to the video file.

required
batch_size Optional[int]

The size of each batch.

32

Returns:

Type Description
VideoPredictions

Results of the prediction.

Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
129
130
131
132
133
134
135
136
137
138
def predict_video(self, video_path: str, batch_size: Optional[int] = 32) -> VideoPredictions:
    """Predict on a video file, by processing the frames in batches.

    :param video_path:  Path to the video file.
    :param batch_size:  The size of each batch.
    :return:            Results of the prediction.
    """
    video_frames, fps = load_video(file_path=video_path)
    result_generator = self._generate_prediction_result(images=video_frames, batch_size=batch_size)
    return self._combine_image_prediction_to_video(result_generator, fps=fps, n_images=len(video_frames))

predict_webcam()

Predict using webcam

Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
140
141
142
143
144
145
146
147
148
149
def predict_webcam(self) -> None:
    """Predict using webcam"""

    def _draw_predictions(frame: np.ndarray) -> np.ndarray:
        """Draw the predictions on a single frame from the stream."""
        frame_prediction = next(iter(self._generate_prediction_result(images=[frame])))
        return frame_prediction.draw()

    video_streaming = WebcamStreaming(frame_processing_fn=_draw_predictions, fps_update_frequency=1)
    video_streaming.run()

PoseEstimationPipeline

Bases: Pipeline

Pipeline specifically designed for pose estimation tasks. The pipeline includes loading images, preprocessing, prediction, and postprocessing.

Parameters:

Name Type Description Default
model SgModule

The object detection model (instance of SgModule) used for making predictions.

required
post_prediction_callback

Callback function to process raw predictions from the model.

required
image_processor Optional[Processing]

Single image processor or a list of image processors for preprocessing and postprocessing the images.

None
device Optional[str]

The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.

None
fuse_model bool

If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.

True
Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
class PoseEstimationPipeline(Pipeline):
    """Pipeline specifically designed for pose estimation tasks.
    The pipeline includes loading images, preprocessing, prediction, and postprocessing.

    :param model:                       The object detection model (instance of SgModule) used for making predictions.
    :param post_prediction_callback:    Callback function to process raw predictions from the model.
    :param image_processor:             Single image processor or a list of image processors for preprocessing and postprocessing the images.
    :param device:                      The device on which the model will be run. If None, will run on current model device. Use "cuda" for GPU support.
    :param fuse_model:                  If True, create a copy of the model, and fuse some of its layers to increase performance. This increases memory usage.
    """

    def __init__(
        self,
        model: SgModule,
        edge_links: Union[np.ndarray, List[Tuple[int, int]]],
        edge_colors: Union[np.ndarray, List[Tuple[int, int, int]]],
        keypoint_colors: Union[np.ndarray, List[Tuple[int, int, int]]],
        post_prediction_callback,
        device: Optional[str] = None,
        image_processor: Optional[Processing] = None,
        fuse_model: bool = True,
    ):
        super().__init__(model=model, device=device, image_processor=image_processor, class_names=None, fuse_model=fuse_model)
        self.post_prediction_callback = post_prediction_callback
        self.edge_links = np.asarray(edge_links, dtype=int)
        self.edge_colors = np.asarray(edge_colors, dtype=int)
        self.keypoint_colors = np.asarray(keypoint_colors, dtype=int)

    def _decode_model_output(self, model_output: Union[List, Tuple, torch.Tensor], model_input: np.ndarray) -> List[PoseEstimationPrediction]:
        """Decode the model output, by applying post prediction callback. This includes NMS.

        :param model_output:    Direct output of the model, without any post-processing.
        :param model_input:     Model input (i.e. images after preprocessing).
        :return:                Predicted Bboxes.
        """
        list_of_predictions = self.post_prediction_callback(model_output)
        decoded_predictions = []
        for image_level_predictions, image in zip(list_of_predictions, model_input):
            decoded_predictions.append(
                PoseEstimationPrediction(
                    poses=image_level_predictions.poses.cpu().numpy() if torch.is_tensor(image_level_predictions.poses) else image_level_predictions.poses,
                    scores=image_level_predictions.scores.cpu().numpy() if torch.is_tensor(image_level_predictions.scores) else image_level_predictions.scores,
                    bboxes_xyxy=image_level_predictions.bboxes_xyxy.cpu().numpy()
                    if torch.is_tensor(image_level_predictions.bboxes_xyxy)
                    else image_level_predictions.bboxes_xyxy,
                    image_shape=image.shape,
                    edge_links=self.edge_links,
                    edge_colors=self.edge_colors,
                    keypoint_colors=self.keypoint_colors,
                )
            )

        return decoded_predictions

    def _instantiate_image_prediction(self, image: np.ndarray, prediction: PoseEstimationPrediction) -> ImagePrediction:
        return ImagePoseEstimationPrediction(image=image, prediction=prediction, class_names=self.class_names)

    def _combine_image_prediction_to_images(
        self, images_predictions: Iterable[PoseEstimationPrediction], n_images: Optional[int] = None
    ) -> ImagesPoseEstimationPrediction:
        if n_images is not None and n_images == 1:
            # Do not show tqdm progress bar if there is only one image
            images_predictions = [next(iter(images_predictions))]
        else:
            images_predictions = [image_predictions for image_predictions in tqdm(images_predictions, total=n_images, desc="Predicting Images")]

        return ImagesPoseEstimationPrediction(_images_prediction_lst=images_predictions)

    def _combine_image_prediction_to_video(
        self, images_predictions: Iterable[ImageDetectionPrediction], fps: float, n_images: Optional[int] = None
    ) -> VideoPoseEstimationPrediction:
        images_predictions = [image_predictions for image_predictions in tqdm(images_predictions, total=n_images, desc="Predicting Video")]
        return VideoPoseEstimationPrediction(_images_prediction_lst=images_predictions, fps=fps)

eval_mode(model)

Set a model in evaluation mode, undo at the end.

Parameters:

Name Type Description Default
model SgModule

The model to set in evaluation mode.

required
Source code in V3_3/src/super_gradients/training/pipelines/pipelines.py
39
40
41
42
43
44
45
46
47
48
@contextmanager
def eval_mode(model: SgModule) -> None:
    """Set a model in evaluation mode, undo at the end.

    :param model: The model to set in evaluation mode.
    """
    _starting_mode = model.training
    model.eval()
    yield
    model.train(mode=_starting_mode)