From f3480c3285db40995f519816b69c3a3ee9aaa608 Mon Sep 17 00:00:00 2001 From: qqii Date: Fri, 9 Feb 2024 23:33:20 +0000 Subject: [PATCH] [Community Pipeline] Skip Marigold depth_colored generation by passing color_map=None --- examples/community/README.md | 2 +- .../community/marigold_depth_estimation.py | 24 +++++++++++-------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/examples/community/README.md b/examples/community/README.md index d1da1b231faf..9f9e118ccf3f 100755 --- a/examples/community/README.md +++ b/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. ) diff --git a/examples/community/marigold_depth_estimation.py b/examples/community/marigold_depth_estimation.py index c6bfce27b577..756b62c95b8a 100644 --- a/examples/community/marigold_depth_estimation.py +++ b/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,