add intermediate outputs

app.py

@@ -41,6 +41,8 @@ def predict(cg: ClimateGAN, api_key):
         masked_input = output_dict["masked_input"]
         wildfire = output_dict["wildfire"]
         smog = output_dict["smog"]
+        depth = np.repeat(output_dict["depth"][..., None], 3, axis=-1)
+        segmentation = output_dict["segmentation"]
 
         climategan_flood = output_dict.get(
             "climategan_flood",
@@ -62,6 +64,8 @@ def predict(cg: ClimateGAN, api_key):
         return (
             input_image,
             masked_input,
+            segmentation,
+            depth,
             climategan_flood,
             stable_flood,
             stable_copy_flood,
@@ -127,6 +131,12 @@ if __name__ == "__main__":
             outputs.append(
                 gr.outputs.Image(type="numpy", label="Masked input image"),
             )
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Segmentation map"),
+            )
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Depth map"),
+            )
         with gr.Row():
             outputs.append(
                 gr.outputs.Image(type="numpy", label="ClimateGAN-Flooded image"),

climategan/trainer.py

@@ -22,7 +22,8 @@ from torch import autograd, sigmoid, softmax
 from torch.cuda.amp import GradScaler, autocast
 from tqdm import tqdm
 
-from climategan.data import get_all_loaders
+from climategan.data import get_all_loaders, decode_segmap_merged_labels
+
 from climategan.discriminator import OmniDiscriminator, create_discriminator
 from climategan.eval_metrics import accuracy, mIOU
 from climategan.fid import compute_val_fid
@@ -38,6 +39,7 @@ from climategan.tutils import (
     get_WGAN_gradient,
     lrgb2srgb,
     normalize,
+    normalize_tensor,
     print_num_parameters,
     shuffle_batch_tuple,
     srgb2lrgb,
@@ -226,7 +228,7 @@ class Trainer:
         cloudy=True,
         auto_resize_640=False,
         ignore_event=set(),
-        return_masks=False,
+        return_intermediates=False,
     ):
         """
         Create a dictionnary of events from a numpy or tensor,
@@ -331,10 +333,20 @@ class Trainer:
                     smog = (smog * 255).astype(np.uint8)
                     output_data["smog"] = smog
 
-        if return_masks:
+        if return_intermediates:
             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["segmentation"] = (
+                decode_segmap_merged_labels(segmentation, "r", False)
+                .cpu()
+                .permute(0, 2, 3, 1)
+                .numpy()
+                .astype(np.uint8)
+            )
 
         return output_data
 

climategan_wrapper.py

@@ -15,6 +15,7 @@ from skimage.transform import resize
 
 from climategan.trainer import Trainer
 
+
 CUDA = torch.cuda.is_available()
 
 
@@ -313,7 +314,7 @@ class ClimateGAN:
             bin_value=0.5,
             half=CUDA,
             ignore_event=ignore_event,
-            return_masks=True,
+            return_intermediates=True,
         )
 
         outputs["input"] = uint8(images, True)