[Community Pipeline] Skip Marigold depth_colored generation by passing color_map=None

Author: qqii

examples/community/README.md

@@ -105,7 +105,7 @@ pipeline_output = pipe(
     # processing_res=768,     # (optional) Maximum resolution of processing. If set to 0: will not resize at all. Defaults to 768.
     # match_input_res=True,   # (optional) Resize depth prediction to match input resolution.
     # batch_size=0,           # (optional) Inference batch size, no bigger than `num_ensemble`. If set to 0, the script will automatically decide the proper batch size. Defaults to 0.
-    # color_map="Spectral",   # (optional) Colormap used to colorize the depth map. Defaults to "Spectral".
+    # color_map="Spectral",   # (optional) Colormap used to colorize the depth map. Defaults to "Spectral". Set to `None` to skip colormap generation.
     # show_progress_bar=True, # (optional) If true, will show progress bars of the inference progress.
 )
 

examples/community/marigold_depth_estimation.py

@@ -50,14 +50,14 @@ class MarigoldDepthOutput(BaseOutput):
     Args:
         depth_np (`np.ndarray`):
             Predicted depth map, with depth values in the range of [0, 1].
-        depth_colored (`PIL.Image.Image`):
+        depth_colored (`None` or `PIL.Image.Image`):
             Colorized depth map, with the shape of [3, H, W] and values in [0, 1].
         uncertainty (`None` or `np.ndarray`):
             Uncalibrated uncertainty(MAD, median absolute deviation) coming from ensembling.
     """
 
     depth_np: np.ndarray
-    depth_colored: Image.Image
+    depth_colored: Union[None, Image.Image]
     uncertainty: Union[None, np.ndarray]
 
 
@@ -139,14 +139,15 @@ def __call__(
                 If set to 0, the script will automatically decide the proper batch size.
             show_progress_bar (`bool`, *optional*, defaults to `True`):
                 Display a progress bar of diffusion denoising.
-            color_map (`str`, *optional*, defaults to `"Spectral"`):
+            color_map (`str`, *optional*, defaults to `"Spectral"`, pass `None` to skip colorized depth map generation):
                 Colormap used to colorize the depth map.
             ensemble_kwargs (`dict`, *optional*, defaults to `None`):
                 Arguments for detailed ensembling settings.
         Returns:
             `MarigoldDepthOutput`: Output class for Marigold monocular depth prediction pipeline, including:
             - **depth_np** (`np.ndarray`) Predicted depth map, with depth values in the range of [0, 1]
-            - **depth_colored** (`PIL.Image.Image`) Colorized depth map, with the shape of [3, H, W] and values in [0, 1]
+            - **depth_colored** (`None` or `PIL.Image.Image`) Colorized depth map, with the shape of [3, H, W] and
+                    values in [0, 1]. None if `color_map` is `None`
             - **uncertainty** (`None` or `np.ndarray`) Uncalibrated uncertainty(MAD, median absolute deviation)
                     coming from ensembling. None if `ensemble_size = 1`
         """
@@ -155,15 +156,19 @@ def __call__(
         input_size = input_image.size
 
         if not match_input_res:
-            assert processing_res is not None, "Value error: `resize_output_back` is only valid with "
+            assert (
+                processing_res is not None
+            ), "Value error: `resize_output_back` is only valid with "
         assert processing_res >= 0
         assert denoising_steps >= 1
         assert ensemble_size >= 1
 
         # ----------------- Image Preprocess -----------------
         # Resize image
         if processing_res > 0:
-            input_image = self.resize_max_res(input_image, max_edge_resolution=processing_res)
+            input_image = self.resize_max_res(
+                input_image, max_edge_resolution=processing_res
+            )
         # Convert the image to RGB, to 1.remove the alpha channel 2.convert B&W to 3-channel
         input_image = input_image.convert("RGB")
         image = np.asarray(input_image)
@@ -188,12 +193,16 @@ def __call__(
                 dtype=self.dtype,
             )
 
-        single_rgb_loader = DataLoader(single_rgb_dataset, batch_size=_bs, shuffle=False)
+        single_rgb_loader = DataLoader(
+            single_rgb_dataset, batch_size=_bs, shuffle=False
+        )
 
         # Predict depth maps (batched)
         depth_pred_ls = []
         if show_progress_bar:
-            iterable = tqdm(single_rgb_loader, desc=" " * 2 + "Inference batches", leave=False)
+            iterable = tqdm(
+                single_rgb_loader, desc=" " * 2 + "Inference batches", leave=False
+            )
         else:
             iterable = single_rgb_loader
         for batch in iterable:
@@ -209,7 +218,9 @@ def __call__(
 
         # ----------------- Test-time ensembling -----------------
         if ensemble_size > 1:
-            depth_pred, pred_uncert = self.ensemble_depths(depth_preds, **(ensemble_kwargs or {}))
+            depth_pred, pred_uncert = self.ensemble_depths(
+                depth_preds, **(ensemble_kwargs or {})
+            )
         else:
             depth_pred = depth_preds
             pred_uncert = None
@@ -233,12 +244,15 @@ def __call__(
         depth_pred = depth_pred.clip(0, 1)
 
         # Colorize
-        depth_colored = self.colorize_depth_maps(
-            depth_pred, 0, 1, cmap=color_map
-        ).squeeze()  # [3, H, W], value in (0, 1)
-        depth_colored = (depth_colored * 255).astype(np.uint8)
-        depth_colored_hwc = self.chw2hwc(depth_colored)
-        depth_colored_img = Image.fromarray(depth_colored_hwc)
+        if color_map is not None:
+            depth_colored = self.colorize_depth_maps(
+                depth_pred, 0, 1, cmap=color_map
+            ).squeeze()  # [3, H, W], value in (0, 1)
+            depth_colored = (depth_colored * 255).astype(np.uint8)
+            depth_colored_hwc = self.chw2hwc(depth_colored)
+            depth_colored_img = Image.fromarray(depth_colored_hwc)
+        else:
+            depth_colored_img = None
         return MarigoldDepthOutput(
             depth_np=depth_pred,
             depth_colored=depth_colored_img,
@@ -261,7 +275,9 @@ def _encode_empty_text(self):
         self.empty_text_embed = self.text_encoder(text_input_ids)[0].to(self.dtype)
 
     @torch.no_grad()
-    def single_infer(self, rgb_in: torch.Tensor, num_inference_steps: int, show_pbar: bool) -> torch.Tensor:
+    def single_infer(
+        self, rgb_in: torch.Tensor, num_inference_steps: int, show_pbar: bool
+    ) -> torch.Tensor:
         """
         Perform an individual depth prediction without ensembling.
 
@@ -285,12 +301,16 @@ def single_infer(self, rgb_in: torch.Tensor, num_inference_steps: int, show_pbar
         rgb_latent = self._encode_rgb(rgb_in)
 
         # Initial depth map (noise)
-        depth_latent = torch.randn(rgb_latent.shape, device=device, dtype=self.dtype)  # [B, 4, h, w]
+        depth_latent = torch.randn(
+            rgb_latent.shape, device=device, dtype=self.dtype
+        )  # [B, 4, h, w]
 
         # Batched empty text embedding
         if self.empty_text_embed is None:
             self._encode_empty_text()
-        batch_empty_text_embed = self.empty_text_embed.repeat((rgb_latent.shape[0], 1, 1))  # [B, 2, 1024]
+        batch_empty_text_embed = self.empty_text_embed.repeat(
+            (rgb_latent.shape[0], 1, 1)
+        )  # [B, 2, 1024]
 
         # Denoising loop
         if show_pbar:
@@ -304,10 +324,14 @@ def single_infer(self, rgb_in: torch.Tensor, num_inference_steps: int, show_pbar
             iterable = enumerate(timesteps)
 
         for i, t in iterable:
-            unet_input = torch.cat([rgb_latent, depth_latent], dim=1)  # this order is important
+            unet_input = torch.cat(
+                [rgb_latent, depth_latent], dim=1
+            )  # this order is important
 
             # predict the noise residual
-            noise_pred = self.unet(unet_input, t, encoder_hidden_states=batch_empty_text_embed).sample  # [B, 4, h, w]
+            noise_pred = self.unet(
+                unet_input, t, encoder_hidden_states=batch_empty_text_embed
+            ).sample  # [B, 4, h, w]
 
             # compute the previous noisy sample x_t -> x_t-1
             depth_latent = self.scheduler.step(noise_pred, t, depth_latent).prev_sample
@@ -375,7 +399,9 @@ def resize_max_res(img: Image.Image, max_edge_resolution: int) -> Image.Image:
             `Image.Image`: Resized image.
         """
         original_width, original_height = img.size
-        downscale_factor = min(max_edge_resolution / original_width, max_edge_resolution / original_height)
+        downscale_factor = min(
+            max_edge_resolution / original_width, max_edge_resolution / original_height
+        )
 
         new_width = int(original_width * downscale_factor)
         new_height = int(original_height * downscale_factor)
@@ -384,7 +410,9 @@ def resize_max_res(img: Image.Image, max_edge_resolution: int) -> Image.Image:
         return resized_img
 
     @staticmethod
-    def colorize_depth_maps(depth_map, min_depth, max_depth, cmap="Spectral", valid_mask=None):
+    def colorize_depth_maps(
+        depth_map, min_depth, max_depth, cmap="Spectral", valid_mask=None
+    ):
         """
         Colorize depth maps.
         """
@@ -526,7 +554,9 @@ def inter_distances(tensors: torch.Tensor):
         if max_res is not None:
             scale_factor = torch.min(max_res / torch.tensor(ori_shape[-2:]))
             if scale_factor < 1:
-                downscaler = torch.nn.Upsample(scale_factor=scale_factor, mode="nearest")
+                downscaler = torch.nn.Upsample(
+                    scale_factor=scale_factor, mode="nearest"
+                )
                 input_images = downscaler(torch.from_numpy(input_images)).numpy()
 
         # init guess