Plugins

DeciClient

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

Source code in latest/src/super_gradients/common/plugins/deci_client.py

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 or deci-common are not installed. Model cannot be loaded from deci lab."
                "Please install deci-platform-client>=5.0.0 and deci-common>=12.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 FileDownloadFailedException("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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

@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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/deci_client.py

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 latest/src/super_gradients/common/plugins/wandb/log_predictions.py

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 latest/src/super_gradients/common/plugins/wandb/log_predictions.py

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 latest/src/super_gradients/common/plugins/wandb/log_predictions.py

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 latest/src/super_gradients/common/plugins/wandb/validation_logger.py

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 latest/src/super_gradients/common/plugins/wandb/validation_logger.py

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"])