[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`
         """
@@ -233,12 +234,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,