Added wandb support (#481)

* added wandb support

* add wandb_img import

* adds wandb as a soft dependency

* fixes formatting

* updates changelog

Co-authored-by: lgvaz <[email protected]>

CHANGELOG.md

@@ -9,6 +9,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - **Switched from poetry to setuptools**
 
 ### Added
+- Function `wandb_img_preds` to help logging bboxes to wandb
+- wandb as a soft dependency
 - Template code for `parsers.SizeMixin` if `parsers.FilepathMixin` is used
 - Get image size without opening image with `get_image_size`
 - Ability to skip record while parsing with `AbortParseRecord`

icevision/imports.py

@@ -50,6 +50,9 @@
 if SoftDependencies.pytorch_lightning:
     import pytorch_lightning as pl
 
+if SoftDependencies.wandb:
+    import wandb
+
 
 # TODO: Stop importing partial from fastcore and move this to utils
 class partial:

icevision/soft_dependencies.py

@@ -20,6 +20,7 @@ def __init__(self):
         self.pytorch_lightning = soft_import("pytorch_lightning")
         self.albumentations = soft_import("albumentations")
         self.effdet = soft_import("effdet")
+        self.wandb = soft_import("wandb")
 
     def check(self) -> Dict[str, bool]:
         return self.__dict__.copy()

icevision/visualize/__init__.py

@@ -1,3 +1,4 @@
 from icevision.visualize.utils import *
 from icevision.visualize.draw_data import *
 from icevision.visualize.show_data import *
+from icevision.visualize.wandb_img import *

icevision/visualize/wandb_img.py

@@ -0,0 +1,78 @@
+__all__ = [
+    "wandb_img_preds",
+]
+
+from icevision.imports import *
+from icevision.data import *
+from icevision.core import *
+
+
+def bbox_wandb(bbox: BBox, label: int, class_id_to_label, pred_score=None):
+    """Creates a wandb compatible dictionary with bbox, label and score"""
+    xmin, ymin, xmax, ymax = map(int, bbox.xyxy)
+
+    box_data = {
+        "position": {"minX": xmin, "maxX": xmax, "minY": ymin, "maxY": ymax},
+        "class_id": label,
+        "domain": "pixel",
+    }
+
+    if pred_score:
+        score = int(pred_score * 100)
+        box_caption = f"{class_id_to_label[label]} ({score}%)"
+        box_data["score"] = score
+    else:
+        box_caption = f"{class_id_to_label[label]}"
+
+    box_data["box_caption"] = box_caption
+
+    return box_data
+
+
+def wandb_image(sample, pred, class_id_to_label, add_ground_truth=False):
+    raw_image = sample["img"]
+    true_bboxes = sample["bboxes"]
+    true_labels = sample["labels"]
+
+    pred_bboxes = pred["bboxes"]
+    pred_labels = pred["labels"].tolist()
+    pred_scores = pred["scores"]
+
+    pred_all_boxes = []
+    # Collect predicted bounding boxes for this image
+    for b_i, bbox in enumerate(pred_bboxes):
+        box_data = bbox_wandb(
+            bbox, pred_labels[b_i], class_id_to_label, pred_score=pred_scores[b_i]
+        )
+        pred_all_boxes.append(box_data)
+
+    # log to wandb: raw image, predictions, and dictionary of class labels for each class id
+    boxes = {
+        "predictions": {"box_data": pred_all_boxes, "class_labels": class_id_to_label}
+    }
+
+    if add_ground_truth:
+        true_all_boxes = []
+        # Collect ground truth bounding boxes for this image
+        for b_i, bbox in enumerate(true_bboxes):
+            box_data = bbox_wandb(bbox, true_labels[b_i], class_id_to_label)
+            true_all_boxes.append(box_data)
+
+        boxes["ground_truth"] = {
+            "box_data": true_all_boxes,
+            "class_labels": class_id_to_label,
+        }
+
+    return wandb.Image(raw_image, boxes=boxes)
+
+
+def wandb_img_preds(samples, preds, class_map, add_ground_truth=False):
+    class_id_to_label = {int(v): k for k, v in class_map.class2id.items()}
+
+    wandb_imgs = []
+    for (sample, pred) in zip(samples, preds):
+        img_wandb = wandb_image(
+            sample, pred, class_id_to_label, add_ground_truth=add_ground_truth
+        )
+        wandb_imgs.append(img_wandb)
+    return wandb_imgs

setup.cfg

@@ -41,6 +41,7 @@ all =
     effdet >=0.1.6,<0.2
     omegaconf >=2,<3
     dataclasses ==0.6
+    wandb >=0.10.7
 inference =
     effdet >=0.1.6,<0.2
     omegaconf >=2,<3

tests/utils/test_soft_dependencies.py

@@ -17,4 +17,5 @@ def test_soft_dependencies():
         "pytorch_lightning": True,
         "albumentations": True,
         "effdet": True,
+        "wandb": True,
     }