Merge branch 'main' into fix_gan_vis

Author: PavelCz

src/reward_preprocessing/common/utils.py

@@ -141,7 +141,10 @@ def visualize_samples(samples: np.ndarray, save_dir):
 
 
 def process_image_array(img: np.ndarray) -> np.ndarray:
-    """Process a numpy array for feeding into PIL.Image.fromarray."""
+    """Process a numpy array for feeding into PIL.Image.fromarray.
+
+    Should already be in (h,w,c) format.
+    """
     up_multiplied = img * 255
     clipped = np.clip(up_multiplied, 0, 255)
     cast = clipped.astype(np.uint8)

src/reward_preprocessing/interpret.py

@@ -381,10 +381,13 @@ def param_f():
         for feature_i in range(num_features):
             custom_logger.log(f"Feature {feature_i}")
 
-            np_trans_tens, indices = nmf.vis_dataset_thumbnail(
+            dataset_thumbnails, indices = nmf.vis_dataset_thumbnail(
                 feature=feature_i, num_mult=4, expand_mult=1
             )
 
+            # remove opacity channel from dataset thumbnails
+            np_trans_tens = dataset_thumbnails[:-1, :, :]
+
             obs, _, next_obs = ndarray_to_transition(np_trans_tens)
 
             _log_single_transition_wandb(
@@ -399,6 +402,15 @@ def param_f():
                 rows,
             )
 
+            if img_save_path is not None:
+                obs_PIL = array_to_image(obs, vis_scale)
+                obs_PIL.save(img_save_path + f"{feature_i}_obs.png")
+                next_obs_PIL = array_to_image(next_obs, vis_scale)
+                next_obs_PIL.save(img_save_path + f"{feature_i}_next_obs.png")
+                custom_logger.log(
+                    f"Saved feature {feature_i} viz in dir {img_save_path}."
+                )
+
     if pyplot:
         plt.show()
     custom_logger.log("Done with visualization.")