From 24c1c77ef131b09e1c2bc16028377b49aeb1c8f6 Mon Sep 17 00:00:00 2001
From: evonneng <evonneng@users.noreply.github.com>
Date: Tue, 7 Jun 2022 17:35:30 -0700
Subject: [PATCH] added wandb support (#24)

---
 environment/requirements.txt |  3 ++-
 pyrad/utils/writer.py        | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/environment/requirements.txt b/environment/requirements.txt
index d690d239f..0783ed076 100644
--- a/environment/requirements.txt
+++ b/environment/requirements.txt
@@ -17,4 +17,5 @@ torch==1.11.0
 torchmetrics[image]==0.8.2
 torchtyping
 torchvision==0.12.0
-tqdm
\ No newline at end of file
+tqdm
+wandb
\ No newline at end of file
diff --git a/pyrad/utils/writer.py b/pyrad/utils/writer.py
index fbf825ac6..66f4384b6 100644
--- a/pyrad/utils/writer.py
+++ b/pyrad/utils/writer.py
@@ -11,6 +11,7 @@ from typing import Dict
 import imageio
 import numpy as np
 import torch
+import wandb
 from omegaconf import DictConfig, ListConfig
 from torch.utils.tensorboard import SummaryWriter
 from torchtyping import TensorType
@@ -189,6 +190,34 @@ class Writer:
             self.write_scalar(name, scalar, step)
 
 
+@decorate_all([check_main_thread])
+class WandbWriter(Writer):
+    """WandDB Writer Class"""
+
+    def __init__(self, save_dir: str):
+        super().__init__(save_dir)
+        wandb.init(dir=save_dir)
+
+    def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None:
+        """_summary_
+
+        Args:
+            name (str): data identifier
+            image (TensorType["H", "W", "C"]): rendered image to write
+        """
+        image = torch.permute(image, (2, 0, 1))
+        wandb.log({name: wandb.Image(image)}, step=step)
+
+    def write_scalar(self, name: str, scalar: float, step: int) -> None:
+        """Wandb method to write a single scalar value to the logger
+
+        Args:
+            name (str): data identifier
+            scalar (float): scalar value to write
+        """
+        wandb.log({name: scalar}, step=step)
+
+
 @decorate_all([check_main_thread])
 class TensorboardWriter(Writer):
     """Tensorboard Writer Class"""
@@ -212,8 +241,7 @@ class TensorboardWriter(Writer):
 
         Args:
             name (str): data identifier
-            x (float): x value to write
-            y (float): y value to write
+            scalar (float): scalar value to write
         """
         self.tb_writer.add_scalar(name, scalar, step)