Skip to content

Plugins

DeciClient

A client to deci platform and model zoo. requires credentials for connection

Source code in src/super_gradients/common/plugins/deci_client.py
 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
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
class DeciClient:
    """
    A client to deci platform and model zoo.
    requires credentials for connection
    """

    def __init__(self):
        if not client_enabled:
            logger.error("deci-platform-client is not installed. Model cannot be loaded from deci lab." "Please install deci-platform-client>=5.0.0")
            return

        self.lab_client = DeciPlatformClient()

    def _get_file(self, model_name: str, file_name: "AutoNACFileName") -> Optional[str]:
        """Get a file from the DeciPlatform if it exists, otherwise returns None
        :param model_name:      Name of the model to download from, as saved in the platform.
        :param file_name:       Name of the file to download
        :return:            Path were the downloaded file was saved to. None if not found.
        """
        try:
            download_link, etag = self.lab_client.get_autonac_model_file_link(
                model_name=model_name,
                file_name=file_name,
                super_gradients_version=super_gradients.__version__,
            )
        except ApiException as e:
            if e.status == 401:
                logger.error(
                    "Unauthorized. wrong credentials or credentials not defined. "
                    "Please provide credentials via environment variables (DECI_CLIENT_ID, DECI_CLIENT_SECRET)"
                )
            elif e.status == 400 and e.body is not None and "message" in e.body:
                logger.error(f"Deci client: {json.loads(e.body)['message']}")
            else:
                logger.debug(e.body)
            return None
        cache_dir = os.path.join(torch.hub.get_dir(), "deci")
        file_path = os.path.join(cache_dir, etag or "", os.path.basename(file_name))
        file_path = normalize_path(file_path)

        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        if os.path.isfile(file_path):
            return file_path

        file_path = self._download_file_to_cache_dir(
            file_url=download_link,
            file_path=file_path,
        )
        return file_path

    def _download_file_to_cache_dir(self, file_url: str, file_path: str, timeout_seconds: Optional[int] = DOWNLOAD_MODEL_TIMEOUT_SECONDS):
        """
        Download a file from a url to a cache dir. The file will be saved in a subfolder named by the etag.
        This allow us to save multiple versions of the same file and cache them, so when a file with the same etag is
        requested, we can return the cached file.

        :param file_url:  Url to download the file from.
        :param file_path: Path to save the file to.
        :return:        Path were the downloaded file was saved to. (same as file_path)
        """
        # TODO: Use requests with stream and limit the file size and timeouts.
        socket.setdefaulttimeout(timeout_seconds)
        try:
            urlretrieve(file_url, file_path)
        except urllib.error.ContentTooShortError as ex:
            raise RuntimeError("File download did not finish correctly " + str(ex))
        return file_path

    def get_model_arch_params(self, model_name: str) -> Optional[DictConfig]:
        """Get the model arch_params from DeciPlatform.
        :param model_name:  Name of the model as saved in the platform.
        :return:            arch_params. None if arch_params were not found for this specific model on this SG version."""
        arch_params_file = self._get_file(model_name, AutoNACFileName.STRUCTURE_YAML)
        if arch_params_file is None:
            return None

        config_name = os.path.basename(arch_params_file)
        download_dir = os.path.dirname(arch_params_file)

        # The arch_params config files need to be saved inside an "arch_params" folder
        _move_file_to_folder(src_file_path=arch_params_file, dest_dir_name="arch_params")

        return load_arch_params(config_name=config_name, recipes_dir_path=download_dir)

    def get_model_recipe(self, model_name: str) -> Optional[DictConfig]:
        """Get the model recipe from DeciPlatform.
        :param model_name:  Name of the model as saved in the platform.
        :return:            recipe. None if recipe were not found for this specific model on this SG version."""
        recipe_file = self._get_file(model_name, AutoNACFileName.RECIPE_YAML)
        if recipe_file is None:
            return None

        config_name = os.path.basename(recipe_file)
        download_dir = os.path.dirname(recipe_file)

        return load_recipe(config_name=config_name, recipes_dir_path=download_dir)

    def get_model_weights(self, model_name: str) -> Optional[str]:
        """Get the path to model weights (downloaded locally).
        :param model_name:  Name of the model as saved in the platform.
        :return:            model_weights path. None if weights were not found for this specific model on this SG version."""
        return self._get_file(model_name=model_name, file_name=AutoNACFileName.WEIGHTS_PTH)

    @staticmethod
    def load_code_from_zipfile(*, file: str, target_path: str, package_name: str = "deci_model_code") -> None:
        """Load additional code files.
        The zip file is extracted, and code files will be placed in the target_path/package_name and imported dynamically,
        :param file:            path to zip file to extract code files from.
        :param target_path:     path to place code files.
        :param package_name:    name of the package to place code files in."""
        package_path = os.path.join(target_path, package_name)
        # create the directory
        os.makedirs(package_path, exist_ok=True)

        # extract code files
        with ZipFile(file) as zipfile:
            zipfile.extractall(package_path)

        # add an init file that imports all code files
        with open(os.path.join(package_path, "__init__.py"), "w") as init_file:
            all_str = "\n\n__all__ = ["
            for code_file in os.listdir(path=package_path):
                if code_file.endswith(".py") and not code_file.startswith("__init__"):
                    init_file.write(f'import {code_file.replace(".py", "")}\n')
                    all_str += f'"{code_file.replace(".py", "")}", '

            all_str += "]\n\n"
            init_file.write(all_str)

        # include in path and import
        sys.path.insert(1, package_path)
        importlib.import_module(package_name)

    def download_and_load_model_additional_code(self, model_name: str, target_path: str, package_name: str = "deci_model_code") -> None:
        """
        try to download code files for this model.
        if found, code files will be placed in the target_path/package_name and imported dynamically
        """

        file = self._get_file(model_name=model_name, file_name=AutoNACFileName.CODE_ZIP)
        package_path = os.path.join(target_path, package_name)
        if file is not None:
            self.load_code_from_zipfile(file=file, target_path=target_path, package_name=package_name)
            logger.info(
                f"*** IMPORTANT ***: files required for the model {model_name} were downloaded and added to your code in:\n{package_path}\n"
                f"These files will be downloaded to the same location each time the model is fetched from the deci-client.\n"
                f"you can override this by passing models.get(... download_required_code=False) and importing the files yourself"
            )

    def upload_model(
        self,
        model: nn.Module,
        name: str,
        input_dimensions: "Sequence[int]",
        target_hardware_types: "Optional[List[HardwareType]]" = None,
        target_quantization_level: "Optional[QuantizationLevel]" = None,
        target_batch_size: "Optional[int]" = None,
    ):
        """
        This function will upload the trained model to the Deci Lab

        :param model:                            The resulting model from the training process
        :param name:                             The model's name
        :param input_dimensions:                 The model's input dimensions
        :param target_hardware_types:            List of hardware types to optimize the model for
        :param target_quantization_level:        The quantization level to optimize the model for
        :param target_batch_size:                The batch size to optimize the model for
        """
        model_id = self.lab_client.register_model(
            model=model,
            name=name,
            framework=FrameworkType.PYTORCH,
            input_dimensions=input_dimensions,
        )
        if target_hardware_types:
            kwargs = {}
            if target_quantization_level:
                kwargs["quantization_level"] = target_quantization_level
            if target_batch_size:
                kwargs["batch_size"] = target_batch_size
            self.lab_client.optimize_model(model_id=model_id, hardware_types=target_hardware_types, **kwargs)

    def is_model_benchmarking(self, name: str) -> bool:
        """Check if a given model is still benchmarking or not.
        :param name: The mode name.
        """
        benchmark_state = self.lab_client.get_model(name=name)[0]["benchmarkState"]
        return benchmark_state in [ModelBenchmarkState.IN_PROGRESS, ModelBenchmarkState.PENDING]

    def register_experiment(self, name: str, model_name: str, resume: bool):
        """Registers a training experiment in Deci's backend.
        :param name:        Name of the experiment to register
        :param model_name:  Name of the model architecture to connect the experiment to
        """
        try:
            self.lab_client.register_user_architecture(name=model_name)
        except (ApiException, ApiTypeError) as e:
            logger.debug(f"The model was already registered, or validation error: {e}")

        self.lab_client.register_experiment(name=name, model_name=model_name, resume=resume)

    def save_experiment_file(self, file_path: str):
        """
        Uploads a training related file to Deci's location in S3. This can be a TensorBoard file or a log
        :params file_path: The local path of the file to be uploaded
        """
        self.lab_client.save_experiment_file(file_path=file_path)

    def upload_file_to_s3(self, tag: str, level: "SentryLevel", from_path: str):
        """Upload a file to the platform S3 bucket.

        :param tag:         Tag that will be associated to the file.
        :param level:       Logging level that will be used to notify the monitoring system that the file was uploaded.
        :param from_path:   Path of the file to upload.
        """
        data = self.lab_client.upload_log_url(tag=tag, level=level)
        signed_url = S3SignedUrl(**data)
        self.lab_client.upload_file_to_s3(from_path=from_path, s3_signed_url=signed_url)

download_and_load_model_additional_code(model_name, target_path, package_name='deci_model_code')

try to download code files for this model. if found, code files will be placed in the target_path/package_name and imported dynamically

Source code in src/super_gradients/common/plugins/deci_client.py
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def download_and_load_model_additional_code(self, model_name: str, target_path: str, package_name: str = "deci_model_code") -> None:
    """
    try to download code files for this model.
    if found, code files will be placed in the target_path/package_name and imported dynamically
    """

    file = self._get_file(model_name=model_name, file_name=AutoNACFileName.CODE_ZIP)
    package_path = os.path.join(target_path, package_name)
    if file is not None:
        self.load_code_from_zipfile(file=file, target_path=target_path, package_name=package_name)
        logger.info(
            f"*** IMPORTANT ***: files required for the model {model_name} were downloaded and added to your code in:\n{package_path}\n"
            f"These files will be downloaded to the same location each time the model is fetched from the deci-client.\n"
            f"you can override this by passing models.get(... download_required_code=False) and importing the files yourself"
        )

get_model_arch_params(model_name)

Get the model arch_params from DeciPlatform.

Parameters:

Name Type Description Default
model_name str

Name of the model as saved in the platform.

required

Returns:

Type Description
Optional[DictConfig]

arch_params. None if arch_params were not found for this specific model on this SG version.

Source code in src/super_gradients/common/plugins/deci_client.py
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def get_model_arch_params(self, model_name: str) -> Optional[DictConfig]:
    """Get the model arch_params from DeciPlatform.
    :param model_name:  Name of the model as saved in the platform.
    :return:            arch_params. None if arch_params were not found for this specific model on this SG version."""
    arch_params_file = self._get_file(model_name, AutoNACFileName.STRUCTURE_YAML)
    if arch_params_file is None:
        return None

    config_name = os.path.basename(arch_params_file)
    download_dir = os.path.dirname(arch_params_file)

    # The arch_params config files need to be saved inside an "arch_params" folder
    _move_file_to_folder(src_file_path=arch_params_file, dest_dir_name="arch_params")

    return load_arch_params(config_name=config_name, recipes_dir_path=download_dir)

get_model_recipe(model_name)

Get the model recipe from DeciPlatform.

Parameters:

Name Type Description Default
model_name str

Name of the model as saved in the platform.

required

Returns:

Type Description
Optional[DictConfig]

recipe. None if recipe were not found for this specific model on this SG version.

Source code in src/super_gradients/common/plugins/deci_client.py
129
130
131
132
133
134
135
136
137
138
139
140
def get_model_recipe(self, model_name: str) -> Optional[DictConfig]:
    """Get the model recipe from DeciPlatform.
    :param model_name:  Name of the model as saved in the platform.
    :return:            recipe. None if recipe were not found for this specific model on this SG version."""
    recipe_file = self._get_file(model_name, AutoNACFileName.RECIPE_YAML)
    if recipe_file is None:
        return None

    config_name = os.path.basename(recipe_file)
    download_dir = os.path.dirname(recipe_file)

    return load_recipe(config_name=config_name, recipes_dir_path=download_dir)

get_model_weights(model_name)

Get the path to model weights (downloaded locally).

Parameters:

Name Type Description Default
model_name str

Name of the model as saved in the platform.

required

Returns:

Type Description
Optional[str]

model_weights path. None if weights were not found for this specific model on this SG version.

Source code in src/super_gradients/common/plugins/deci_client.py
142
143
144
145
146
def get_model_weights(self, model_name: str) -> Optional[str]:
    """Get the path to model weights (downloaded locally).
    :param model_name:  Name of the model as saved in the platform.
    :return:            model_weights path. None if weights were not found for this specific model on this SG version."""
    return self._get_file(model_name=model_name, file_name=AutoNACFileName.WEIGHTS_PTH)

is_model_benchmarking(name)

Check if a given model is still benchmarking or not.

Parameters:

Name Type Description Default
name str

The mode name.

required
Source code in src/super_gradients/common/plugins/deci_client.py
227
228
229
230
231
232
def is_model_benchmarking(self, name: str) -> bool:
    """Check if a given model is still benchmarking or not.
    :param name: The mode name.
    """
    benchmark_state = self.lab_client.get_model(name=name)[0]["benchmarkState"]
    return benchmark_state in [ModelBenchmarkState.IN_PROGRESS, ModelBenchmarkState.PENDING]

load_code_from_zipfile(*, file, target_path, package_name='deci_model_code') staticmethod

Load additional code files. The zip file is extracted, and code files will be placed in the target_path/package_name and imported dynamically,

Parameters:

Name Type Description Default
file str

path to zip file to extract code files from.

required
target_path str

path to place code files.

required
package_name str

name of the package to place code files in.

'deci_model_code'
Source code in src/super_gradients/common/plugins/deci_client.py
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
@staticmethod
def load_code_from_zipfile(*, file: str, target_path: str, package_name: str = "deci_model_code") -> None:
    """Load additional code files.
    The zip file is extracted, and code files will be placed in the target_path/package_name and imported dynamically,
    :param file:            path to zip file to extract code files from.
    :param target_path:     path to place code files.
    :param package_name:    name of the package to place code files in."""
    package_path = os.path.join(target_path, package_name)
    # create the directory
    os.makedirs(package_path, exist_ok=True)

    # extract code files
    with ZipFile(file) as zipfile:
        zipfile.extractall(package_path)

    # add an init file that imports all code files
    with open(os.path.join(package_path, "__init__.py"), "w") as init_file:
        all_str = "\n\n__all__ = ["
        for code_file in os.listdir(path=package_path):
            if code_file.endswith(".py") and not code_file.startswith("__init__"):
                init_file.write(f'import {code_file.replace(".py", "")}\n')
                all_str += f'"{code_file.replace(".py", "")}", '

        all_str += "]\n\n"
        init_file.write(all_str)

    # include in path and import
    sys.path.insert(1, package_path)
    importlib.import_module(package_name)

register_experiment(name, model_name, resume)

Registers a training experiment in Deci's backend.

Parameters:

Name Type Description Default
name str

Name of the experiment to register

required
model_name str

Name of the model architecture to connect the experiment to

required
Source code in src/super_gradients/common/plugins/deci_client.py
234
235
236
237
238
239
240
241
242
243
244
def register_experiment(self, name: str, model_name: str, resume: bool):
    """Registers a training experiment in Deci's backend.
    :param name:        Name of the experiment to register
    :param model_name:  Name of the model architecture to connect the experiment to
    """
    try:
        self.lab_client.register_user_architecture(name=model_name)
    except (ApiException, ApiTypeError) as e:
        logger.debug(f"The model was already registered, or validation error: {e}")

    self.lab_client.register_experiment(name=name, model_name=model_name, resume=resume)

save_experiment_file(file_path)

Uploads a training related file to Deci's location in S3. This can be a TensorBoard file or a log

Parameters:

Name Type Description Default
file_path str

The local path of the file to be uploaded

required
Source code in src/super_gradients/common/plugins/deci_client.py
246
247
248
249
250
251
def save_experiment_file(self, file_path: str):
    """
    Uploads a training related file to Deci's location in S3. This can be a TensorBoard file or a log
    :params file_path: The local path of the file to be uploaded
    """
    self.lab_client.save_experiment_file(file_path=file_path)

upload_file_to_s3(tag, level, from_path)

Upload a file to the platform S3 bucket.

Parameters:

Name Type Description Default
tag str

Tag that will be associated to the file.

required
level SentryLevel

Logging level that will be used to notify the monitoring system that the file was uploaded.

required
from_path str

Path of the file to upload.

required
Source code in src/super_gradients/common/plugins/deci_client.py
253
254
255
256
257
258
259
260
261
262
def upload_file_to_s3(self, tag: str, level: "SentryLevel", from_path: str):
    """Upload a file to the platform S3 bucket.

    :param tag:         Tag that will be associated to the file.
    :param level:       Logging level that will be used to notify the monitoring system that the file was uploaded.
    :param from_path:   Path of the file to upload.
    """
    data = self.lab_client.upload_log_url(tag=tag, level=level)
    signed_url = S3SignedUrl(**data)
    self.lab_client.upload_file_to_s3(from_path=from_path, s3_signed_url=signed_url)

upload_model(model, name, input_dimensions, target_hardware_types=None, target_quantization_level=None, target_batch_size=None)

This function will upload the trained model to the Deci Lab

Parameters:

Name Type Description Default
model nn.Module

The resulting model from the training process

required
name str

The model's name

required
input_dimensions Sequence[int]

The model's input dimensions

required
target_hardware_types Optional[List[HardwareType]]

List of hardware types to optimize the model for

None
target_quantization_level Optional[QuantizationLevel]

The quantization level to optimize the model for

None
target_batch_size Optional[int]

The batch size to optimize the model for

None
Source code in src/super_gradients/common/plugins/deci_client.py
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
def upload_model(
    self,
    model: nn.Module,
    name: str,
    input_dimensions: "Sequence[int]",
    target_hardware_types: "Optional[List[HardwareType]]" = None,
    target_quantization_level: "Optional[QuantizationLevel]" = None,
    target_batch_size: "Optional[int]" = None,
):
    """
    This function will upload the trained model to the Deci Lab

    :param model:                            The resulting model from the training process
    :param name:                             The model's name
    :param input_dimensions:                 The model's input dimensions
    :param target_hardware_types:            List of hardware types to optimize the model for
    :param target_quantization_level:        The quantization level to optimize the model for
    :param target_batch_size:                The batch size to optimize the model for
    """
    model_id = self.lab_client.register_model(
        model=model,
        name=name,
        framework=FrameworkType.PYTORCH,
        input_dimensions=input_dimensions,
    )
    if target_hardware_types:
        kwargs = {}
        if target_quantization_level:
            kwargs["quantization_level"] = target_quantization_level
        if target_batch_size:
            kwargs["batch_size"] = target_batch_size
        self.lab_client.optimize_model(model_id=model_id, hardware_types=target_hardware_types, **kwargs)

log_detection_results_to_wandb(prediction, show_confidence=True)

Log predictions for object detection to Weights & Biases using interactive bounding box overlays.

Parameters:

Name Type Description Default
prediction ImagesDetectionPrediction

The model predictions (a super_gradients.training.models.prediction_results.ImagesDetectionPrediction object)

required
show_confidence bool

Whether to log confidence scores to Weights & Biases or not.

True
Source code in src/super_gradients/common/plugins/wandb/log_predictions.py
49
50
51
52
53
54
55
56
57
58
59
def log_detection_results_to_wandb(prediction: ImagesDetectionPrediction, show_confidence: bool = True):
    """Log predictions for object detection to Weights & Biases using interactive bounding box overlays.

    :param prediction:        The model predictions (a `super_gradients.training.models.prediction_results.ImagesDetectionPrediction` object)
    :param show_confidence:   Whether to log confidence scores to Weights & Biases or not.
    """
    if wandb.run is None:
        raise wandb.Error("Images and bounding boxes cannot be visualized on Weights & Biases without initializing a run using `wandb.init()`")
    for prediction in prediction._images_prediction_lst:
        wandb_image = visualize_image_detection_prediction_on_wandb(prediction=prediction, show_confidence=show_confidence)
        wandb.log({"Predictions": wandb_image})

plot_detection_dataset_on_wandb(detection_dataset, max_examples=None, dataset_name=None, reverse_channels=True)

Log a detection dataset to Weights & Biases Table.

Parameters:

Name Type Description Default
detection_dataset DetectionDataset

The Detection Dataset (a super_gradients.training.datasets.detection_datasets.DetectionDataset object)

required
max_examples int

Maximum number of examples from the detection dataset to plot (an int).

None
dataset_name str

Name of the dataset (a str).

None
reverse_channels bool

Reverse the order of channels on the images while plotting.

True
Source code in src/super_gradients/common/plugins/wandb/log_predictions.py
 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
def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_examples: int = None, dataset_name: str = None, reverse_channels: bool = True):
    """Log a detection dataset to Weights & Biases Table.

    :param detection_dataset:       The Detection Dataset (a `super_gradients.training.datasets.detection_datasets.DetectionDataset` object)
    :param max_examples:            Maximum number of examples from the detection dataset to plot (an `int`).
    :param dataset_name:            Name of the dataset (a `str`).
    :param reverse_channels:        Reverse the order of channels on the images while plotting.
    """
    max_examples = len(detection_dataset) if max_examples is None else max_examples
    wandb_table = wandb.Table(columns=["Images", "Class-Frequencies"])
    input_format = detection_dataset.output_target_format
    target_format_transform = DetectionTargetsFormatTransform(input_format=input_format, output_format=XYXY_LABEL)
    class_id_to_labels = {int(_id): str(_class_name) for _id, _class_name in enumerate(detection_dataset.classes)}
    for data_idx in tqdm(range(max_examples), desc="Plotting Examples on Weights & Biases"):
        image, targets, *_ = detection_dataset[data_idx]
        image = image.transpose(1, 2, 0).astype(np.int32)
        sample = target_format_transform({"image": image, "target": targets})
        boxes = sample["target"][:, 0:4]
        boxes = boxes[(boxes != 0).any(axis=1)]
        classes = targets[:, 0].tolist()
        wandb_boxes = []
        class_frequencies = {str(_class_name): 0 for _id, _class_name in enumerate(detection_dataset.classes)}
        for idx in range(boxes.shape[0]):
            wandb_boxes.append(
                {
                    "position": {
                        "minX": float(boxes[idx][0] / image.shape[1]),
                        "maxX": float(boxes[idx][2] / image.shape[1]),
                        "minY": float(boxes[idx][1] / image.shape[0]),
                        "maxY": float(boxes[idx][3] / image.shape[0]),
                    },
                    "class_id": int(classes[idx]),
                    "box_caption": str(class_id_to_labels[int(classes[idx])]),
                }
            )
            class_frequencies[str(class_id_to_labels[int(classes[idx])])] += 1
        image = image[:, :, ::-1] if reverse_channels else image
        wandb_table.add_data(wandb.Image(image, boxes={"ground_truth": {"box_data": wandb_boxes, "class_labels": class_id_to_labels}}), class_frequencies)
    dataset_name = "Dataset" if dataset_name is None else dataset_name
    wandb.log({dataset_name: wandb_table}, commit=False)

visualize_image_detection_prediction_on_wandb(prediction, show_confidence, reverse_channels=False)

Visualize detection results on a single image.

Parameters:

Name Type Description Default
prediction ImageDetectionPrediction

Prediction results of a single image (a super_gradients.training.models.prediction_results.ImageDetectionPrediction object)

required
show_confidence bool

Whether to log confidence scores to Weights & Biases or not.

required
reverse_channels bool

Reverse the order of channels on the images while plotting.

False
Source code in src/super_gradients/common/plugins/wandb/log_predictions.py
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
def visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool, reverse_channels: bool = False):
    """Visualize detection results on a single image.

    :param prediction:          Prediction results of a single image
                                (a `super_gradients.training.models.prediction_results.ImageDetectionPrediction` object)
    :param show_confidence:     Whether to log confidence scores to Weights & Biases or not.
    :param reverse_channels:    Reverse the order of channels on the images while plotting.
    """
    boxes = []
    image = prediction.image.copy()
    image = image[:, :, ::-1] if reverse_channels else image
    height, width, _ = image.shape
    class_id_to_labels = {int(_id): str(_class_name) for _id, _class_name in enumerate(prediction.class_names)}

    for pred_i in range(len(prediction.prediction)):
        class_id = int(prediction.prediction.labels[pred_i])
        box = {
            "position": {
                "minX": float(int(prediction.prediction.bboxes_xyxy[pred_i, 0]) / width),
                "maxX": float(int(prediction.prediction.bboxes_xyxy[pred_i, 2]) / width),
                "minY": float(int(prediction.prediction.bboxes_xyxy[pred_i, 1]) / height),
                "maxY": float(int(prediction.prediction.bboxes_xyxy[pred_i, 3]) / height),
            },
            "class_id": int(class_id),
            "box_caption": str(prediction.class_names[class_id]),
        }
        if show_confidence:
            box["scores"] = {"confidence": float(round(prediction.prediction.confidence[pred_i], 2))}
        boxes.append(box)

    return wandb.Image(image, boxes={"predictions": {"box_data": boxes, "class_labels": class_id_to_labels}})

WandBDetectionValidationPredictionLoggerCallback

Bases: Callback

Source code in src/super_gradients/common/plugins/wandb/validation_logger.py
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
class WandBDetectionValidationPredictionLoggerCallback(Callback):
    def __init__(
        self,
        class_names,
        max_predictions_plotted: Optional[int] = None,
        post_prediction_callback: Optional[DetectionPostPredictionCallback] = None,
        reverse_channels: bool = True,
    ) -> None:
        """A callback for logging object detection predictions to Weights & Biases during training. This callback is logging images on each batch in validation
        and accumulating generated images in a `wandb.Table` in the RAM. This could potentially cause OOM errors for very large datasets like COCO. In order to
        avoid this, it is recommended to explicitly set the parameter `max_predictions_plotted` to a small value, thus limiting the number of images logged in
        the table.

        :param class_names:                 A list of class names.
        :param max_predictions_plotted:     Maximum number of predictions to be plotted per epoch. This is set to `None` by default which means that the
                                            predictions corresponding to all images from `context.inputs` is logged, otherwise only `max_predictions_plotted`
                                            number of images is logged. Since `WandBDetectionValidationPredictionLoggerCallback` accumulates the generated
                                            images in the RAM, it is advisable that the value of this parameter be explicitly specified for larger datasets in
                                            order to avoid out-of-memory errors.
        :param post_prediction_callback:    `DetectionPostPredictionCallback` for post-processing outputs of the model.
        :param reverse_channels:            Reverse the order of channels on the images while plotting.
        """
        super().__init__()
        self.class_names = class_names
        self.max_predictions_plotted = max_predictions_plotted
        self.post_prediction_callback = post_prediction_callback
        self.reverse_channels = reverse_channels
        self.wandb_images = []
        self.epoch_count = 0
        self.mean_prediction_dicts = []
        self.wandb_table = wandb.Table(columns=["Epoch", "Prediction", "Mean-Confidence"])

    def on_validation_batch_end(self, context: PhaseContext) -> None:
        self.wandb_images = []
        mean_prediction_dict = {class_name: 0.0 for class_name in self.class_names}
        if isinstance(context.net, HasPredict):
            post_nms_predictions = context.net(context.inputs)
        else:
            self.post_prediction_callback = (
                unwrap_model(context.net).get_post_prediction_callback() if self.post_prediction_callback is None else self.post_prediction_callback
            )
            self.post_prediction_callback.fuse_layers = False
            post_nms_predictions = self.post_prediction_callback(context.preds, device=context.device)
        if self.max_predictions_plotted is not None:
            post_nms_predictions = post_nms_predictions[: self.max_predictions_plotted]
            input_images = context.inputs[: self.max_predictions_plotted]
        else:
            input_images = context.inputs
        for prediction, image in zip(post_nms_predictions, input_images):
            prediction = prediction if prediction is not None else torch.zeros((0, 6), dtype=torch.float32)
            prediction = prediction.detach().cpu().numpy()
            postprocessed_image = image.detach().cpu().numpy().transpose(1, 2, 0).astype(np.int32)
            image_prediction = ImageDetectionPrediction(
                image=postprocessed_image,
                class_names=self.class_names,
                prediction=DetectionPrediction(
                    bboxes=prediction[:, :4],
                    confidence=prediction[:, 4],
                    labels=prediction[:, 5],
                    bbox_format="xyxy",
                    image_shape=image.shape,
                ),
            )
            for predicted_label, prediction_confidence in zip(prediction[:, 5], prediction[:, 4]):
                mean_prediction_dict[self.class_names[int(predicted_label)]] += prediction_confidence
            mean_prediction_dict = {k: v / len(prediction[:, 4]) for k, v in mean_prediction_dict.items()}
            self.mean_prediction_dicts.append(mean_prediction_dict)
            wandb_image = visualize_image_detection_prediction_on_wandb(
                prediction=image_prediction, show_confidence=True, reverse_channels=self.reverse_channels
            )
            self.wandb_images.append(wandb_image)

    def on_validation_loader_end(self, context: PhaseContext) -> None:
        for wandb_image, mean_prediction_dict in zip(self.wandb_images, self.mean_prediction_dicts):
            self.wandb_table.add_data(self.epoch_count, wandb_image, mean_prediction_dict)
        self.wandb_images, self.mean_prediction_dicts = [], []
        self.epoch_count += 1

    def on_training_end(self, context: PhaseContext) -> None:
        wandb.log({"Validation-Prediction": self.wandb_table})

__init__(class_names, max_predictions_plotted=None, post_prediction_callback=None, reverse_channels=True)

A callback for logging object detection predictions to Weights & Biases during training. This callback is logging images on each batch in validation and accumulating generated images in a wandb.Table in the RAM. This could potentially cause OOM errors for very large datasets like COCO. In order to avoid this, it is recommended to explicitly set the parameter max_predictions_plotted to a small value, thus limiting the number of images logged in the table.

Parameters:

Name Type Description Default
class_names

A list of class names.

required
max_predictions_plotted Optional[int]

Maximum number of predictions to be plotted per epoch. This is set to None by default which means that the predictions corresponding to all images from context.inputs is logged, otherwise only max_predictions_plotted number of images is logged. Since WandBDetectionValidationPredictionLoggerCallback accumulates the generated images in the RAM, it is advisable that the value of this parameter be explicitly specified for larger datasets in order to avoid out-of-memory errors.

None
post_prediction_callback Optional[DetectionPostPredictionCallback]

DetectionPostPredictionCallback for post-processing outputs of the model.

None
reverse_channels bool

Reverse the order of channels on the images while plotting.

True
Source code in src/super_gradients/common/plugins/wandb/validation_logger.py
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
def __init__(
    self,
    class_names,
    max_predictions_plotted: Optional[int] = None,
    post_prediction_callback: Optional[DetectionPostPredictionCallback] = None,
    reverse_channels: bool = True,
) -> None:
    """A callback for logging object detection predictions to Weights & Biases during training. This callback is logging images on each batch in validation
    and accumulating generated images in a `wandb.Table` in the RAM. This could potentially cause OOM errors for very large datasets like COCO. In order to
    avoid this, it is recommended to explicitly set the parameter `max_predictions_plotted` to a small value, thus limiting the number of images logged in
    the table.

    :param class_names:                 A list of class names.
    :param max_predictions_plotted:     Maximum number of predictions to be plotted per epoch. This is set to `None` by default which means that the
                                        predictions corresponding to all images from `context.inputs` is logged, otherwise only `max_predictions_plotted`
                                        number of images is logged. Since `WandBDetectionValidationPredictionLoggerCallback` accumulates the generated
                                        images in the RAM, it is advisable that the value of this parameter be explicitly specified for larger datasets in
                                        order to avoid out-of-memory errors.
    :param post_prediction_callback:    `DetectionPostPredictionCallback` for post-processing outputs of the model.
    :param reverse_channels:            Reverse the order of channels on the images while plotting.
    """
    super().__init__()
    self.class_names = class_names
    self.max_predictions_plotted = max_predictions_plotted
    self.post_prediction_callback = post_prediction_callback
    self.reverse_channels = reverse_channels
    self.wandb_images = []
    self.epoch_count = 0
    self.mean_prediction_dicts = []
    self.wandb_table = wandb.Table(columns=["Epoch", "Prediction", "Mean-Confidence"])