Merge pull request #22 from HumanCompatibleAI/try-feature-vis

Action labels and overview visualization

Author: dfilan

src/reward_preprocessing/common/utils.py

@@ -234,15 +234,16 @@ def forward(self, latent_tens: th.Tensor):
         return self.reward_net.forward(obs, action_vec, next_obs, done)
 
 
-def log_np_img_wandb(
-    arr: np.ndarray,
+def log_img_wandb(
+    img: Union[np.ndarray, PIL.Image.Image],
     caption: str,
     wandb_key: str,
     logger: HierarchicalLogger,
     scale: int = 1,
     step: Optional[int] = None,
 ) -> None:
-    """Log visualized np.ndarray to wandb using given logger.
+    """Log np.ndarray as image or PIL image. Logs to wandb using given logger.
+    If scale is provided, image will be scaled in both cases.
 
     Args:
         - arr: Array to turn into image, save.
@@ -253,8 +254,17 @@ def log_np_img_wandb(
         - step: Step for logging. If not provided, the logger dumping will be skipped.
             In that case logs will be dumped with the next dump().
     """
-
-    pil_img = array_to_image(arr, scale)
+    if isinstance(img, np.ndarray):
+        pil_img = array_to_image(img, scale)
+    elif isinstance(img, PIL.Image.Image):
+        pil_img = img.resize(
+            # PIL expects tuple of (width, height), as opposed to numpy's
+            # (height, width).
+            size=(img.width * scale, img.height * scale),
+            resample=Image.NEAREST,
+        )
+    else:
+        raise ValueError(f"img must be np.ndarray or PIL.Image.Image, {type(img)=}")
     wb_img = wandb.Image(pil_img, caption=caption)
     logger.record(wandb_key, wb_img)
     if step is not None:
@@ -264,7 +274,8 @@ def log_np_img_wandb(
 def array_to_image(arr: np.ndarray, scale: int) -> PIL.Image.Image:
     """Take numpy array on [0,1] scale, return PIL image."""
     return Image.fromarray(np.uint8(arr * 255), mode="RGB").resize(
-        # PIL expects tuple of (width, height), numpy's index 1 is width, 0 height.
+        # PIL expects tuple of (width, height), numpy's dimension 1 is width, and
+        # dimension 0 height.
         size=(arr.shape[1] * scale, arr.shape[0] * scale),
         resample=Image.NEAREST,
     )