new `tensor_to_uint8_numpy_image` tensor util

climategan/trainer.py

@@ -39,10 +39,10 @@ from climategan.tutils import (
     get_WGAN_gradient,
     lrgb2srgb,
     normalize,
-    normalize_tensor,
     print_num_parameters,
     shuffle_batch_tuple,
     srgb2lrgb,
+    tensor_to_uint8_numpy_image,
     vgg_preprocess,
     zero_grad,
 )
@@ -231,12 +231,15 @@ class Trainer:
         return_intermediates=False,
     ):
         """
-        Create a dictionnary of events from a numpy or tensor,
+        Create a dictionary of events from a numpy or tensor,
         single or batch image data.
 
-        stores is a dictionnary of times for the Timer class.
+        stores is a dictionary of times for the Timer class.
 
         bin_value is used to binarize (or not) flood masks
+
+        all values in the output dictionary have 4 dimensions:
+            BxHxWxC if numpy else BxCxHxW
         """
         assert self.is_setup
         assert len(x.shape) in {3, 4}, f"Unknown Data shape {x.shape}"
@@ -316,21 +319,14 @@ class Trainer:
             with Timer(store=stores.get("numpy", [])):
                 if "flood" not in ignore_event:
                     # normalize to 0-1
-                    flood = normalize(flood).cpu()
-                    # convert to numpy
-                    flood = flood.permute(0, 2, 3, 1).numpy()
+                    flood = tensor_to_uint8_numpy_image(flood)
                     # convert to 0-255 uint8
-                    flood = (flood * 255).astype(np.uint8)
                     output_data["flood"] = flood
                 if "wildfire" not in ignore_event:
-                    wildfire = normalize(wildfire).cpu()
-                    wildfire = wildfire.permute(0, 2, 3, 1).numpy()
-                    wildfire = (wildfire * 255).astype(np.uint8)
+                    wildfire = tensor_to_uint8_numpy_image(wildfire)
                     output_data["wildfire"] = wildfire
                 if "smog" not in ignore_event:
-                    smog = normalize(smog).cpu()
-                    smog = smog.permute(0, 2, 3, 1).numpy()
-                    smog = (smog * 255).astype(np.uint8)
+                    smog = tensor_to_uint8_numpy_image(smog)
                     output_data["smog"] = smog
 
         if return_intermediates:
@@ -338,9 +334,7 @@ class Trainer:
                 output_data["mask"] = (
                     ((mask > bin_value) * 255).cpu().numpy().astype(np.uint8)
                 )
-                output_data["depth"] = (
-                    normalize_tensor(depth).cpu().squeeze(1).numpy().astype(np.uint8) * 255
-                )
+                output_data["depth"] = tensor_to_uint8_numpy_image(depth)
                 output_data["segmentation"] = (
                     decode_segmap_merged_labels(segmentation, "r", False)
                     .cpu()

climategan/tutils.py

@@ -564,14 +564,29 @@ def lrgb2srgb(ims):
     return outs[0]
 
 
-def normalize(t, mini=0, maxi=1):
+def normalize(t, mini=0.0, maxi=1.0):
+    """
+    Normalizes a tensor to [0, 1].
+    If the tensor has more than 3 dimensions, the first one
+    is assumed to be the batch dimension and the tensor is
+    normalized per batch element, not across the batches.
+
+    Args:
+        t (torch.Tensor): Tensor to normalize
+        mini (float, optional): Min allowed value. Defaults to 0.
+        maxi (float, optional): Max allowed value. Defaults to 1.
+
+    Returns:
+        torch.Tensor: The normalized tensor
+    """
     if len(t.shape) == 3:
         return mini + (maxi - mini) * (t - t.min()) / (t.max() - t.min())
 
     batch_size = t.shape[0]
-    min_t = t.reshape(batch_size, -1).min(1)[0].reshape(batch_size, 1, 1, 1)
+    extra_dims = [1] * (t.ndim - 1)
+    min_t = t.reshape(batch_size, -1).min(1)[0].reshape(batch_size, *extra_dims)
     t = t - min_t
-    max_t = t.reshape(batch_size, -1).max(1)[0].reshape(batch_size, 1, 1, 1)
+    max_t = t.reshape(batch_size, -1).max(1)[0].reshape(batch_size, *extra_dims)
     t = t / max_t
     return mini + (maxi - mini) * t
 
@@ -644,7 +659,7 @@ def write_architecture(trainer):
             f.write(output)
 
 
-def rand_perlin_2d(shape, res, fade=lambda t: 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3):
+def rand_perlin_2d(shape, res, fade=lambda t: 6 * t**5 - 15 * t**4 + 10 * t**3):
     delta = (res[0] / shape[0], res[1] / shape[1])
     d = (shape[0] // res[0], shape[1] // res[1])
 
@@ -719,3 +734,30 @@ def tensor_ims_to_np_uint8s(ims):
         nps.append(n.astype(np.uint8))
 
     return nps[0] if len(nps) == 1 else nps
+
+
+def tensor_to_uint8_numpy_image(tensor):
+    """
+    Turns a BxCxHxW tensor into a numpy image:
+    * normalize
+    * to [0, 255]
+    * detach
+    * channels last
+    * to uin8
+    * to cpu
+    * to numpy
+
+    Args:
+        tensor (torch.Tensor): Tensor to transform
+
+    Returns:
+        np.array: BxHxWxC np.uint8 array in [0, 255]
+    """
+    return (
+        normalize(tensor, 0, 255)  # [0, 255]
+        .detach()  # detach from graph if needed
+        .permute(0, 2, 3, 1)  # BxHxWxC
+        .to(torch.uint8)  # uint8
+        .cpu()  # cpu
+        .numpy()  # numpy array
+    )