From a81fddbff5f07c96df353fc2d0a40b54d22cbff3 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Fri, 9 Jun 2023 17:09:31 +0530
Subject: [PATCH 01/23] add: plot_detection_dataset_on_wandb

---
 .../common/plugins/wandb/__init__.py          |  4 +-
 .../common/plugins/wandb/log_predictions.py   | 38 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/__init__.py b/src/super_gradients/common/plugins/wandb/__init__.py
index 8614e61d98..91507b8f5d 100644
--- a/src/super_gradients/common/plugins/wandb/__init__.py
+++ b/src/super_gradients/common/plugins/wandb/__init__.py
@@ -1,4 +1,4 @@
-from super_gradients.common.plugins.wandb.log_predictions import log_detection_results_to_wandb
+from super_gradients.common.plugins.wandb.log_predictions import log_detection_results_to_wandb, plot_detection_dataset_on_wandb
 
 
-__all__ = ["log_detection_results_to_wandb"]
+__all__ = ["log_detection_results_to_wandb", "plot_detection_dataset_on_wandb"]
diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index 04fa007aa4..ee7f3e31c6 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -3,7 +3,13 @@
 except (ModuleNotFoundError, ImportError, NameError):
     pass  # no action or logging - this is normal in most cases
 
+import numpy as np
+from tqdm import tqdm
+
 from super_gradients.training.models.prediction_results import ImageDetectionPrediction, ImagesDetectionPrediction
+from super_gradients.training.transforms.transforms import DetectionTargetsFormatTransform
+from super_gradients.training.datasets.data_formats.default_formats import XYXY_LABEL
+from super_gradients.training.datasets.detection_datasets import DetectionDataset
 
 
 def _visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool):
@@ -43,3 +49,35 @@ def log_detection_results_to_wandb(prediction: ImagesDetectionPrediction, show_c
         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:
         _visualize_image_detection_prediction_on_wandb(prediction=prediction, show_confidence=show_confidence)
+
+
+def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_examples: int = None, dataset_name: str = None):
+    """Log a detection dataset"""
+    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)}
+    wandb_images = []
+    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 = []
+        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])]),
+                }
+            )
+        wandb_images.append(wandb.Image(image, boxes={"predictions": {"box_data": wandb_boxes, "class_labels": class_id_to_labels}}))
+    dataset_name = "Dataset" if dataset_name is None else dataset_name
+    wandb.log({dataset_name: wandb_images})

From 76dc7b3b18647b49c47a7f9c6939955e57257ca6 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Fri, 9 Jun 2023 17:34:59 +0530
Subject: [PATCH 02/23] add: doctring for plot_detection_dataset_on_wandb

---
 .../common/plugins/wandb/log_predictions.py       | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index ee7f3e31c6..555772a611 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -52,11 +52,16 @@ def log_detection_results_to_wandb(prediction: ImagesDetectionPrediction, show_c
 
 
 def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_examples: int = None, dataset_name: str = None):
-    """Log a detection dataset"""
+    """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`).
+    """
+    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)}
-    wandb_images = []
     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)
@@ -65,6 +70,7 @@ def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_exa
         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(
                 {
@@ -78,6 +84,7 @@ def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_exa
                     "box_caption": str(class_id_to_labels[int(classes[idx])]),
                 }
             )
-        wandb_images.append(wandb.Image(image, boxes={"predictions": {"box_data": wandb_boxes, "class_labels": class_id_to_labels}}))
+            class_frequencies[str(class_id_to_labels[int(classes[idx])])] += 1
+        wandb_table.add_data(wandb.Image(image, boxes={"predictions": {"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_images})
+    wandb.log({dataset_name: wandb_table})

From 07f8c0ccfa1ae2c5a25fc50f7366e2d5471b9298 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Fri, 9 Jun 2023 17:49:00 +0530
Subject: [PATCH 03/23] update: bbox type

---
 src/super_gradients/common/plugins/wandb/log_predictions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index 555772a611..c24756469d 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -85,6 +85,6 @@ def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_exa
                 }
             )
             class_frequencies[str(class_id_to_labels[int(classes[idx])])] += 1
-        wandb_table.add_data(wandb.Image(image, boxes={"predictions": {"box_data": wandb_boxes, "class_labels": class_id_to_labels}}), class_frequencies)
+        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})

From 8f7b5c68d3906e0ad5abce9d9b154aaa39ca138c Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Sat, 10 Jun 2023 03:30:10 +0530
Subject: [PATCH 04/23] update: visualize_image_detection_prediction_on_wandb

---
 .../common/plugins/wandb/log_predictions.py        | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index c24756469d..ad094b7d46 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -12,7 +12,12 @@
 from super_gradients.training.datasets.detection_datasets import DetectionDataset
 
 
-def _visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool):
+def visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool):
+    """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.
+    """
     boxes = []
     image = prediction.image.copy()
     height, width, _ = image.shape
@@ -34,9 +39,7 @@ def _visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPre
             box["scores"] = {"confidence": float(round(prediction.prediction.confidence[pred_i], 2))}
         boxes.append(box)
 
-    wandb_image = wandb.Image(image, boxes={"predictions": {"box_data": boxes, "class_labels": class_id_to_labels}})
-
-    wandb.log({"Predictions": wandb_image})
+    return wandb.Image(image, boxes={"predictions": {"box_data": boxes, "class_labels": class_id_to_labels}})
 
 
 def log_detection_results_to_wandb(prediction: ImagesDetectionPrediction, show_confidence: bool = True):
@@ -48,7 +51,8 @@ def log_detection_results_to_wandb(prediction: ImagesDetectionPrediction, show_c
     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:
-        _visualize_image_detection_prediction_on_wandb(prediction=prediction, show_confidence=show_confidence)
+        wandb_image = visualize_image_detection_prediction_on_wandb(prediction=prediction, show_confidence=show_confidence)
+        wandb.log({"Predictions": wandb_image})
 
 
 def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_examples: int = None, dataset_name: str = None):

From 71a52af39d2dea8880bbe67c819056903e2cadf9 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Sat, 10 Jun 2023 03:32:26 +0530
Subject: [PATCH 05/23] update: fix wandb module and linting

---
 src/super_gradients/common/plugins/wandb/__init__.py      | 8 ++++++--
 .../common/plugins/wandb/log_predictions.py               | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/__init__.py b/src/super_gradients/common/plugins/wandb/__init__.py
index 91507b8f5d..faf704b52b 100644
--- a/src/super_gradients/common/plugins/wandb/__init__.py
+++ b/src/super_gradients/common/plugins/wandb/__init__.py
@@ -1,4 +1,8 @@
-from super_gradients.common.plugins.wandb.log_predictions import log_detection_results_to_wandb, plot_detection_dataset_on_wandb
+from super_gradients.common.plugins.wandb.log_predictions import (
+    visualize_image_detection_prediction_on_wandb,
+    log_detection_results_to_wandb,
+    plot_detection_dataset_on_wandb,
+)
 
 
-__all__ = ["log_detection_results_to_wandb", "plot_detection_dataset_on_wandb"]
+__all__ = ["visualize_image_detection_prediction_on_wandb", "log_detection_results_to_wandb", "plot_detection_dataset_on_wandb"]
diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index ad094b7d46..8055c08d02 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -14,7 +14,7 @@
 
 def visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool):
     """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.
     """

From de567246c9460a8820008c7fc4d47a771f995bff Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Sat, 10 Jun 2023 03:39:00 +0530
Subject: [PATCH 06/23] add: WandBDetectionValidationPredictionLoggerCallback

---
 .../common/plugins/wandb/__init__.py          |  8 ++-
 .../common/plugins/wandb/validation_logger.py | 64 +++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 src/super_gradients/common/plugins/wandb/validation_logger.py

diff --git a/src/super_gradients/common/plugins/wandb/__init__.py b/src/super_gradients/common/plugins/wandb/__init__.py
index faf704b52b..f50cd75fad 100644
--- a/src/super_gradients/common/plugins/wandb/__init__.py
+++ b/src/super_gradients/common/plugins/wandb/__init__.py
@@ -3,6 +3,12 @@
     log_detection_results_to_wandb,
     plot_detection_dataset_on_wandb,
 )
+from super_gradients.common.plugins.wandb.validation_logger import WandBDetectionValidationPredictionLoggerCallback
 
 
-__all__ = ["visualize_image_detection_prediction_on_wandb", "log_detection_results_to_wandb", "plot_detection_dataset_on_wandb"]
+__all__ = [
+    "visualize_image_detection_prediction_on_wandb",
+    "log_detection_results_to_wandb",
+    "plot_detection_dataset_on_wandb",
+    "WandBDetectionValidationPredictionLoggerCallback",
+]
diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
new file mode 100644
index 0000000000..38ca10e50c
--- /dev/null
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -0,0 +1,64 @@
+import wandb
+import torch
+import numpy as np
+
+from super_gradients.training.utils.callbacks import Callback, PhaseContext
+from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
+from super_gradients.common.plugins.wandb.log_predictions import visualize_image_detection_prediction_on_wandb
+from super_gradients.training.models.predictions import DetectionPrediction
+from super_gradients.training.models.prediction_results import ImageDetectionPrediction
+
+
+class WandBDetectionValidationPredictionLoggerCallback(Callback):
+    def __init__(
+        self,
+        class_names,
+        score_threshold: float = 0.001,
+        nms_threshold: float = 0.6,
+        nms_top_k: int = 1000,
+        max_predictions: int = 300,
+        multi_label_per_box: bool = True,
+    ) -> None:
+        """A callback for logging object detection predictions to Weights & Biases during training.
+
+        :param class_names:             A list of class names.
+        :param score_threshold: Predictions confidence threshold. Predictions with score lower than score_threshold will not participate in Top-K & NMS
+        :param iou: IoU threshold for NMS step.
+        :param nms_top_k: Number of predictions participating in NMS step
+        :param max_predictions: maximum number of boxes to return after NMS step
+        """
+        super().__init__()
+        self.class_names = class_names
+        self.post_prediction_callback = PPYoloEPostPredictionCallback(
+            score_threshold=score_threshold,
+            nms_threshold=nms_threshold,
+            nms_top_k=nms_top_k,
+            max_predictions=max_predictions,
+            multi_label_per_box=multi_label_per_box,
+        )
+        self.wandb_images = []
+
+    def on_validation_batch_end(self, context: PhaseContext) -> None:
+        self.wandb_images = []
+        post_nms_predictions = self.post_prediction_callback(context.preds, device=context.device)
+        for prediction, image in zip(post_nms_predictions, context.inputs):
+            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,
+                ),
+            )
+            wandb_image = visualize_image_detection_prediction_on_wandb(prediction=image_prediction, show_confidence=True)
+            self.wandb_images.append(wandb_image)
+
+    def on_validation_loader_end(self, context: PhaseContext) -> None:
+        wandb.log({"Validation-Predictions": self.wandb_images})
+        self.wandb_images = []

From facad14136780ca440b78305dbd18a5145e6abea Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Sat, 10 Jun 2023 03:41:31 +0530
Subject: [PATCH 07/23] update: docstring for
 WandBDetectionValidationPredictionLoggerCallback

---
 .../common/plugins/wandb/validation_logger.py          | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 38ca10e50c..48a352e074 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -21,11 +21,11 @@ def __init__(
     ) -> None:
         """A callback for logging object detection predictions to Weights & Biases during training.
 
-        :param class_names:             A list of class names.
-        :param score_threshold: Predictions confidence threshold. Predictions with score lower than score_threshold will not participate in Top-K & NMS
-        :param iou: IoU threshold for NMS step.
-        :param nms_top_k: Number of predictions participating in NMS step
-        :param max_predictions: maximum number of boxes to return after NMS step
+        :param class_names:         A list of class names.
+        :param score_threshold:     Predictions confidence threshold. Predictions with score lower than score_threshold will not participate in Top-K & NMS
+        :param iou:                 IoU threshold for NMS step.
+        :param nms_top_k:           Number of predictions participating in NMS step
+        :param max_predictions:     maximum number of boxes to return after NMS step
         """
         super().__init__()
         self.class_names = class_names

From 4c3e328ec48aad816bd14165f298b221491e1b8c Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Tue, 13 Jun 2023 00:24:57 +0530
Subject: [PATCH 08/23] update:
 WandBDetectionValidationPredictionLoggerCallback

---
 .../common/plugins/wandb/validation_logger.py  | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 48a352e074..b2aa098712 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -1,3 +1,4 @@
+from super_gradients.training.utils.callbacks.base_callbacks import PhaseContext
 import wandb
 import torch
 import numpy as np
@@ -37,9 +38,13 @@ def __init__(
             multi_label_per_box=multi_label_per_box,
         )
         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}
         post_nms_predictions = self.post_prediction_callback(context.preds, device=context.device)
         for prediction, image in zip(post_nms_predictions, context.inputs):
             prediction = prediction if prediction is not None else torch.zeros((0, 6), dtype=torch.float32)
@@ -56,9 +61,18 @@ def on_validation_batch_end(self, context: PhaseContext) -> None:
                     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)
             self.wandb_images.append(wandb_image)
 
     def on_validation_loader_end(self, context: PhaseContext) -> None:
-        wandb.log({"Validation-Predictions": self.wandb_images})
-        self.wandb_images = []
+        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})

From ddcd6308515e994d28619bab47a3585c34f61e3b Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Tue, 13 Jun 2023 11:03:51 +0530
Subject: [PATCH 09/23] update: plot_detection_dataset_on_wandb

---
 src/super_gradients/common/plugins/wandb/log_predictions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index 8055c08d02..8f9dd24784 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -91,4 +91,4 @@ def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_exa
             class_frequencies[str(class_id_to_labels[int(classes[idx])])] += 1
         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})
+    wandb.log({dataset_name: wandb_table}, commit=False)

From 65fb5218777339a7b8677f4947026fdab00c52bf Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Tue, 13 Jun 2023 21:14:43 +0530
Subject: [PATCH 10/23] update: imports

---
 src/super_gradients/common/plugins/wandb/log_predictions.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index 8f9dd24784..b2f44906be 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -6,11 +6,12 @@
 import numpy as np
 from tqdm import tqdm
 
-from super_gradients.training.models.prediction_results import ImageDetectionPrediction, ImagesDetectionPrediction
 from super_gradients.training.transforms.transforms import DetectionTargetsFormatTransform
 from super_gradients.training.datasets.data_formats.default_formats import XYXY_LABEL
 from super_gradients.training.datasets.detection_datasets import DetectionDataset
 
+from super_gradients.training.models.prediction_results import ImageDetectionPrediction, ImagesDetectionPrediction
+
 
 def visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool):
     """Visualize detection results on a single image.

From 59edcb207bae29e3f9bb4b0ec00913d838fc674c Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Wed, 14 Jun 2023 17:58:13 +0530
Subject: [PATCH 11/23] fix: linting

---
 src/super_gradients/common/plugins/wandb/validation_logger.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index b2aa098712..29e5f7d9cc 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -73,6 +73,7 @@ def on_validation_loader_end(self, context: PhaseContext) -> None:
             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:
+        _ = context
         wandb.log({"Validation-Prediction": self.wandb_table})

From 7271db44f73baa7e441792042ebf027752f64c92 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Wed, 14 Jun 2023 18:00:34 +0530
Subject: [PATCH 12/23] fix: linting

---
 src/super_gradients/common/plugins/wandb/validation_logger.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 29e5f7d9cc..a5c3039bcb 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -69,6 +69,7 @@ def on_validation_batch_end(self, context: PhaseContext) -> None:
             self.wandb_images.append(wandb_image)
 
     def on_validation_loader_end(self, context: PhaseContext) -> None:
+        _ = context
         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 = [], []

From abf274d2e01aa7393f32f6cd13adfa363fc0af55 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Wed, 14 Jun 2023 18:05:36 +0530
Subject: [PATCH 13/23] fix: linting

---
 src/super_gradients/common/plugins/wandb/validation_logger.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index a5c3039bcb..188669a74f 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -1,4 +1,3 @@
-from super_gradients.training.utils.callbacks.base_callbacks import PhaseContext
 import wandb
 import torch
 import numpy as np
@@ -69,12 +68,10 @@ def on_validation_batch_end(self, context: PhaseContext) -> None:
             self.wandb_images.append(wandb_image)
 
     def on_validation_loader_end(self, context: PhaseContext) -> None:
-        _ = context
         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:
-        _ = context
         wandb.log({"Validation-Prediction": self.wandb_table})

From ae339d682e579739204065e6446757d690a669a6 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 22 Jun 2023 19:46:41 +0530
Subject: [PATCH 14/23] fix: imports in validation_logger.py

---
 src/super_gradients/common/plugins/wandb/validation_logger.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 188669a74f..9d5b0e5d76 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -6,7 +6,7 @@
 from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
 from super_gradients.common.plugins.wandb.log_predictions import visualize_image_detection_prediction_on_wandb
 from super_gradients.training.models.predictions import DetectionPrediction
-from super_gradients.training.models.prediction_results import ImageDetectionPrediction
+from super_gradients.training.utils.predict import ImageDetectionPrediction
 
 
 class WandBDetectionValidationPredictionLoggerCallback(Callback):

From d9b5bf3074a0e945b787b52846e381e7a2c04099 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 22 Jun 2023 19:51:28 +0530
Subject: [PATCH 15/23] Update log_predictions.py

---
 src/super_gradients/common/plugins/wandb/log_predictions.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index 311272a3f0..d3ec702a66 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -63,6 +63,7 @@ def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_exa
     :param max_examples:            Maximum number of examples from the detection dataset to plot (an `int`).
     :param dataset_name:            Name of the dataset (a `str`).
     """
+    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)

From 344c4e80c21d0536bf203e0f96884587a16070ff Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 29 Jun 2023 01:31:22 +0530
Subject: [PATCH 16/23] update: add max_predictions_plotted parameter to
 WandBDetectionValidationPredictionLoggerCallback

---
 .../common/plugins/wandb/validation_logger.py | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 9d5b0e5d76..49fb64df16 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 import wandb
 import torch
 import numpy as np
@@ -17,18 +19,24 @@ def __init__(
         nms_threshold: float = 0.6,
         nms_top_k: int = 1000,
         max_predictions: int = 300,
+        max_predictions_plotted: Optional[int] = None,
         multi_label_per_box: bool = True,
     ) -> None:
         """A callback for logging object detection predictions to Weights & Biases during training.
 
-        :param class_names:         A list of class names.
-        :param score_threshold:     Predictions confidence threshold. Predictions with score lower than score_threshold will not participate in Top-K & NMS
-        :param iou:                 IoU threshold for NMS step.
-        :param nms_top_k:           Number of predictions participating in NMS step
-        :param max_predictions:     maximum number of boxes to return after NMS step
+        :param class_names:             A list of class names.
+        :param score_threshold:         Predictions confidence threshold. Predictions with score lower than score_threshold will not participate in Top-K & NMS
+        :param iou:                     IoU threshold for NMS step.
+        :param nms_top_k:               Number of predictions participating in NMS step
+        :param max_predictions:         Maximum number of boxes to return after NMS step
+        :param max_predictions_plotted: Maximum number of predictions to be plotted per epoch. This is set to `None` by default which means thatthe 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.
         """
         super().__init__()
         self.class_names = class_names
+        self.max_predictions_plotted = max_predictions_plotted
         self.post_prediction_callback = PPYoloEPostPredictionCallback(
             score_threshold=score_threshold,
             nms_threshold=nms_threshold,
@@ -45,7 +53,12 @@ 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}
         post_nms_predictions = self.post_prediction_callback(context.preds, device=context.device)
-        for prediction, image in zip(post_nms_predictions, context.inputs):
+        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)

From 8522be21750fb6aa9cdbe42762aedff12d189b0a Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 29 Jun 2023 01:34:29 +0530
Subject: [PATCH 17/23] update: docstring for
 WandBDetectionValidationPredictionLoggerCallback

---
 .../common/plugins/wandb/validation_logger.py              | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 49fb64df16..75b1a8a561 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -30,9 +30,10 @@ def __init__(
         :param nms_top_k:               Number of predictions participating in NMS step
         :param max_predictions:         Maximum number of boxes to return after NMS step
         :param max_predictions_plotted: Maximum number of predictions to be plotted per epoch. This is set to `None` by default which means thatthe 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.
+                                        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.
         """
         super().__init__()
         self.class_names = class_names

From b7225bec18052ca233c7a83895dd5f21319c60a0 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 27 Jul 2023 20:24:09 +0530
Subject: [PATCH 18/23] update: validation logger

---
 .../common/plugins/wandb/validation_logger.py | 43 ++++++++-----------
 1 file changed, 19 insertions(+), 24 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 75b1a8a561..5b47b24fe9 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -5,46 +5,35 @@
 import numpy as np
 
 from super_gradients.training.utils.callbacks import Callback, PhaseContext
-from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
 from super_gradients.common.plugins.wandb.log_predictions import visualize_image_detection_prediction_on_wandb
 from super_gradients.training.models.predictions import DetectionPrediction
 from super_gradients.training.utils.predict import ImageDetectionPrediction
+from super_gradients.training.utils.detection_utils import DetectionPostPredictionCallback
+from super_gradients.module_interfaces import HasPredict
+from super_gradients.training.utils.utils import unwrap_model
 
 
 class WandBDetectionValidationPredictionLoggerCallback(Callback):
     def __init__(
         self,
         class_names,
-        score_threshold: float = 0.001,
-        nms_threshold: float = 0.6,
-        nms_top_k: int = 1000,
-        max_predictions: int = 300,
         max_predictions_plotted: Optional[int] = None,
-        multi_label_per_box: bool = True,
+        post_prediction_callback: Optional[DetectionPostPredictionCallback] = None,
     ) -> None:
         """A callback for logging object detection predictions to Weights & Biases during training.
 
-        :param class_names:             A list of class names.
-        :param score_threshold:         Predictions confidence threshold. Predictions with score lower than score_threshold will not participate in Top-K & NMS
-        :param iou:                     IoU threshold for NMS step.
-        :param nms_top_k:               Number of predictions participating in NMS step
-        :param max_predictions:         Maximum number of boxes to return after NMS step
-        :param max_predictions_plotted: Maximum number of predictions to be plotted per epoch. This is set to `None` by default which means thatthe 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 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 thatthe 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.
         """
         super().__init__()
         self.class_names = class_names
         self.max_predictions_plotted = max_predictions_plotted
-        self.post_prediction_callback = PPYoloEPostPredictionCallback(
-            score_threshold=score_threshold,
-            nms_threshold=nms_threshold,
-            nms_top_k=nms_top_k,
-            max_predictions=max_predictions,
-            multi_label_per_box=multi_label_per_box,
-        )
+        self.post_prediction_callback = post_prediction_callback
         self.wandb_images = []
         self.epoch_count = 0
         self.mean_prediction_dicts = []
@@ -53,7 +42,13 @@ def __init__(
     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}
-        post_nms_predictions = self.post_prediction_callback(context.preds, device=context.device)
+        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
+            )
+            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]

From caef860bc32224a548cb59be360bdd3be128415c Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 27 Jul 2023 21:57:15 +0530
Subject: [PATCH 19/23] update:
 WandBDetectionValidationPredictionLoggerCallback +
 visualize_image_detection_prediction_on_wandb

---
 .../common/plugins/wandb/log_predictions.py                | 3 ++-
 .../common/plugins/wandb/validation_logger.py              | 7 ++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index d3ec702a66..df31eead96 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -13,7 +13,7 @@
 from super_gradients.training.utils.predict import ImageDetectionPrediction, ImagesDetectionPrediction
 
 
-def visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPrediction, show_confidence: bool):
+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)
@@ -21,6 +21,7 @@ def visualize_image_detection_prediction_on_wandb(prediction: ImageDetectionPred
     """
     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)}
 
diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 5b47b24fe9..4dc8ed8354 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -19,6 +19,7 @@ def __init__(
         class_names,
         max_predictions_plotted: Optional[int] = None,
         post_prediction_callback: Optional[DetectionPostPredictionCallback] = None,
+        reverse_channel_order: bool = True,
     ) -> None:
         """A callback for logging object detection predictions to Weights & Biases during training.
 
@@ -34,6 +35,7 @@ def __init__(
         self.class_names = class_names
         self.max_predictions_plotted = max_predictions_plotted
         self.post_prediction_callback = post_prediction_callback
+        self.reverse_channel_order = reverse_channel_order
         self.wandb_images = []
         self.epoch_count = 0
         self.mean_prediction_dicts = []
@@ -48,6 +50,7 @@ def on_validation_batch_end(self, context: PhaseContext) -> None:
             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]
@@ -73,7 +76,9 @@ def on_validation_batch_end(self, context: PhaseContext) -> None:
                 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)
+            wandb_image = visualize_image_detection_prediction_on_wandb(
+                prediction=image_prediction, show_confidence=True, reverse_channel_order=self.reverse_channel_order
+            )
             self.wandb_images.append(wandb_image)
 
     def on_validation_loader_end(self, context: PhaseContext) -> None:

From 880f55d66be4fb68d301b40648585aef8a35ec1d Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 27 Jul 2023 22:04:37 +0530
Subject: [PATCH 20/23] update: wandb import + docstring

---
 .../common/plugins/wandb/validation_logger.py       | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 4dc8ed8354..3a65ec8ded 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -1,6 +1,5 @@
 from typing import Optional
 
-import wandb
 import torch
 import numpy as np
 
@@ -12,6 +11,11 @@
 from super_gradients.module_interfaces import HasPredict
 from super_gradients.training.utils.utils import unwrap_model
 
+try:
+    import wandb
+except (ModuleNotFoundError, ImportError, NameError):
+    pass  # no action or logging - this is normal in most cases
+
 
 class WandBDetectionValidationPredictionLoggerCallback(Callback):
     def __init__(
@@ -21,10 +25,13 @@ def __init__(
         post_prediction_callback: Optional[DetectionPostPredictionCallback] = None,
         reverse_channel_order: bool = True,
     ) -> None:
-        """A callback for logging object detection predictions to Weights & Biases during training.
+        """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 thatthe predictions
+        :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

From a5c4328aede80aa694590ec02c132efd3e3f46af Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Thu, 27 Jul 2023 22:07:46 +0530
Subject: [PATCH 21/23] update: doctring

---
 .../common/plugins/wandb/validation_logger.py        | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 3a65ec8ded..3382cbf274 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -30,12 +30,12 @@ def __init__(
         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 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.
         """
         super().__init__()

From 49ba7030ca579276e8c73d50da15ae60d87fcb23 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Fri, 28 Jul 2023 16:03:02 +0530
Subject: [PATCH 22/23] update: docstrings + channel reversal

---
 .../common/plugins/wandb/log_predictions.py   | 10 ++++++---
 .../common/plugins/wandb/validation_logger.py | 21 ++++++++++---------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/log_predictions.py b/src/super_gradients/common/plugins/wandb/log_predictions.py
index df31eead96..a242349b58 100644
--- a/src/super_gradients/common/plugins/wandb/log_predictions.py
+++ b/src/super_gradients/common/plugins/wandb/log_predictions.py
@@ -16,8 +16,10 @@
 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 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()
@@ -57,12 +59,13 @@ def log_detection_results_to_wandb(prediction: ImagesDetectionPrediction, show_c
         wandb.log({"Predictions": wandb_image})
 
 
-def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_examples: int = None, dataset_name: str = None):
+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"])
@@ -92,6 +95,7 @@ def plot_detection_dataset_on_wandb(detection_dataset: DetectionDataset, max_exa
                 }
             )
             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)
diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index 3382cbf274..a3fcdcf577 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -23,26 +23,27 @@ def __init__(
         class_names,
         max_predictions_plotted: Optional[int] = None,
         post_prediction_callback: Optional[DetectionPostPredictionCallback] = None,
-        reverse_channel_order: bool = True,
+        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 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_channel_order = reverse_channel_order
+        self.reverse_channels = reverse_channels
         self.wandb_images = []
         self.epoch_count = 0
         self.mean_prediction_dicts = []
@@ -84,7 +85,7 @@ def on_validation_batch_end(self, context: PhaseContext) -> None:
             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_channel_order=self.reverse_channel_order
+                prediction=image_prediction, show_confidence=True, reverse_channels=self.reverse_channels
             )
             self.wandb_images.append(wandb_image)
 

From b6acd23868a5e178a67e894e26a3bbd6e9bee051 Mon Sep 17 00:00:00 2001
From: Soumik Rakshit <19soumik.rakshit96@gmail.com>
Date: Fri, 28 Jul 2023 16:05:25 +0530
Subject: [PATCH 23/23] update: make ci happy

---
 .../common/plugins/wandb/validation_logger.py               | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/super_gradients/common/plugins/wandb/validation_logger.py b/src/super_gradients/common/plugins/wandb/validation_logger.py
index a3fcdcf577..357e39f86d 100644
--- a/src/super_gradients/common/plugins/wandb/validation_logger.py
+++ b/src/super_gradients/common/plugins/wandb/validation_logger.py
@@ -33,9 +33,9 @@ def __init__(
         :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.
+                                            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.
         """