diffusers 0.27.1__py3-none-any.whl → 0.28.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- diffusers/__init__.py +18 -1
- diffusers/callbacks.py +156 -0
- diffusers/commands/env.py +110 -6
- diffusers/configuration_utils.py +16 -11
- diffusers/dependency_versions_table.py +2 -1
- diffusers/image_processor.py +158 -45
- diffusers/loaders/__init__.py +2 -5
- diffusers/loaders/autoencoder.py +4 -4
- diffusers/loaders/controlnet.py +4 -4
- diffusers/loaders/ip_adapter.py +80 -22
- diffusers/loaders/lora.py +134 -20
- diffusers/loaders/lora_conversion_utils.py +46 -43
- diffusers/loaders/peft.py +4 -3
- diffusers/loaders/single_file.py +401 -170
- diffusers/loaders/single_file_model.py +290 -0
- diffusers/loaders/single_file_utils.py +616 -672
- diffusers/loaders/textual_inversion.py +41 -20
- diffusers/loaders/unet.py +168 -115
- diffusers/loaders/unet_loader_utils.py +163 -0
- diffusers/models/__init__.py +2 -0
- diffusers/models/activations.py +11 -3
- diffusers/models/attention.py +10 -11
- diffusers/models/attention_processor.py +367 -148
- diffusers/models/autoencoders/autoencoder_asym_kl.py +14 -16
- diffusers/models/autoencoders/autoencoder_kl.py +18 -19
- diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +11 -11
- diffusers/models/autoencoders/autoencoder_tiny.py +16 -16
- diffusers/models/autoencoders/consistency_decoder_vae.py +36 -11
- diffusers/models/autoencoders/vae.py +23 -24
- diffusers/models/controlnet.py +12 -9
- diffusers/models/controlnet_flax.py +4 -4
- diffusers/models/controlnet_xs.py +1915 -0
- diffusers/models/downsampling.py +17 -18
- diffusers/models/embeddings.py +147 -24
- diffusers/models/model_loading_utils.py +149 -0
- diffusers/models/modeling_flax_pytorch_utils.py +2 -1
- diffusers/models/modeling_flax_utils.py +4 -4
- diffusers/models/modeling_pytorch_flax_utils.py +1 -1
- diffusers/models/modeling_utils.py +118 -98
- diffusers/models/resnet.py +18 -23
- diffusers/models/transformer_temporal.py +3 -3
- diffusers/models/transformers/dual_transformer_2d.py +4 -4
- diffusers/models/transformers/prior_transformer.py +7 -7
- diffusers/models/transformers/t5_film_transformer.py +17 -19
- diffusers/models/transformers/transformer_2d.py +272 -156
- diffusers/models/transformers/transformer_temporal.py +10 -10
- diffusers/models/unets/unet_1d.py +5 -5
- diffusers/models/unets/unet_1d_blocks.py +29 -29
- diffusers/models/unets/unet_2d.py +6 -6
- diffusers/models/unets/unet_2d_blocks.py +137 -128
- diffusers/models/unets/unet_2d_condition.py +20 -15
- diffusers/models/unets/unet_2d_condition_flax.py +6 -5
- diffusers/models/unets/unet_3d_blocks.py +79 -77
- diffusers/models/unets/unet_3d_condition.py +13 -9
- diffusers/models/unets/unet_i2vgen_xl.py +14 -13
- diffusers/models/unets/unet_kandinsky3.py +1 -1
- diffusers/models/unets/unet_motion_model.py +114 -14
- diffusers/models/unets/unet_spatio_temporal_condition.py +15 -14
- diffusers/models/unets/unet_stable_cascade.py +16 -13
- diffusers/models/upsampling.py +17 -20
- diffusers/models/vq_model.py +16 -15
- diffusers/pipelines/__init__.py +25 -3
- diffusers/pipelines/amused/pipeline_amused.py +12 -12
- diffusers/pipelines/amused/pipeline_amused_img2img.py +14 -12
- diffusers/pipelines/amused/pipeline_amused_inpaint.py +13 -11
- diffusers/pipelines/animatediff/__init__.py +2 -0
- diffusers/pipelines/animatediff/pipeline_animatediff.py +24 -46
- diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +1284 -0
- diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +82 -72
- diffusers/pipelines/animatediff/pipeline_output.py +3 -2
- diffusers/pipelines/audioldm/pipeline_audioldm.py +14 -14
- diffusers/pipelines/audioldm2/modeling_audioldm2.py +54 -35
- diffusers/pipelines/audioldm2/pipeline_audioldm2.py +120 -36
- diffusers/pipelines/auto_pipeline.py +21 -17
- diffusers/pipelines/blip_diffusion/blip_image_processing.py +1 -1
- diffusers/pipelines/blip_diffusion/modeling_blip2.py +5 -5
- diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py +1 -1
- diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py +2 -2
- diffusers/pipelines/consistency_models/pipeline_consistency_models.py +5 -5
- diffusers/pipelines/controlnet/multicontrolnet.py +4 -8
- diffusers/pipelines/controlnet/pipeline_controlnet.py +87 -52
- diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py +2 -2
- diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +50 -43
- diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +52 -40
- diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +80 -47
- diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +147 -49
- diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +89 -55
- diffusers/pipelines/controlnet_xs/__init__.py +68 -0
- diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +911 -0
- diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +1115 -0
- diffusers/pipelines/deepfloyd_if/pipeline_if.py +14 -28
- diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +18 -33
- diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +21 -39
- diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +20 -36
- diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +23 -39
- diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +17 -32
- diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py +11 -11
- diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +43 -20
- diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +36 -18
- diffusers/pipelines/deprecated/repaint/pipeline_repaint.py +2 -2
- diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py +7 -7
- diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +12 -12
- diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +18 -21
- diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +20 -15
- diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +20 -15
- diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +30 -25
- diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +69 -59
- diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +13 -13
- diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py +10 -5
- diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +11 -6
- diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +10 -5
- diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +5 -5
- diffusers/pipelines/dit/pipeline_dit.py +3 -0
- diffusers/pipelines/free_init_utils.py +39 -38
- diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py +33 -48
- diffusers/pipelines/kandinsky/pipeline_kandinsky.py +8 -8
- diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +23 -20
- diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py +11 -11
- diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py +12 -12
- diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py +10 -10
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +6 -6
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +32 -29
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py +10 -10
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py +10 -10
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +6 -6
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +8 -8
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +7 -7
- diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py +6 -6
- diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +3 -3
- diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +20 -33
- diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +24 -35
- diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +48 -30
- diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +50 -28
- diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +11 -11
- diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +61 -67
- diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +70 -69
- diffusers/pipelines/ledits_pp/pipeline_output.py +2 -2
- diffusers/pipelines/marigold/__init__.py +50 -0
- diffusers/pipelines/marigold/marigold_image_processing.py +561 -0
- diffusers/pipelines/marigold/pipeline_marigold_depth.py +813 -0
- diffusers/pipelines/marigold/pipeline_marigold_normals.py +690 -0
- diffusers/pipelines/musicldm/pipeline_musicldm.py +14 -14
- diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +17 -12
- diffusers/pipelines/pia/pipeline_pia.py +39 -125
- diffusers/pipelines/pipeline_flax_utils.py +4 -4
- diffusers/pipelines/pipeline_loading_utils.py +268 -23
- diffusers/pipelines/pipeline_utils.py +266 -37
- diffusers/pipelines/pixart_alpha/__init__.py +8 -1
- diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +65 -75
- diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +880 -0
- diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +10 -5
- diffusers/pipelines/shap_e/pipeline_shap_e.py +3 -3
- diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +14 -14
- diffusers/pipelines/shap_e/renderer.py +1 -1
- diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +36 -22
- diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +23 -19
- diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +33 -32
- diffusers/pipelines/stable_diffusion/__init__.py +0 -1
- diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +18 -11
- diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py +2 -2
- diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py +6 -6
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +73 -39
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +24 -17
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +13 -8
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +66 -36
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +82 -46
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +123 -28
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +6 -6
- diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +16 -16
- diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +24 -19
- diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +37 -31
- diffusers/pipelines/stable_diffusion/safety_checker.py +2 -1
- diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +23 -15
- diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +44 -42
- diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +23 -18
- diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +19 -14
- diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +20 -15
- diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +24 -19
- diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +65 -32
- diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +274 -38
- diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +10 -5
- diffusers/pipelines/stable_diffusion_safe/safety_checker.py +1 -1
- diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +92 -25
- diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +88 -44
- diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +108 -56
- diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +96 -51
- diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +45 -25
- diffusers/pipelines/stable_diffusion_xl/watermark.py +9 -3
- diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +110 -57
- diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +59 -30
- diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +71 -42
- diffusers/pipelines/text_to_video_synthesis/pipeline_output.py +3 -2
- diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +18 -41
- diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +21 -85
- diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +28 -19
- diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +39 -33
- diffusers/pipelines/unclip/pipeline_unclip.py +6 -6
- diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +6 -6
- diffusers/pipelines/unidiffuser/modeling_text_decoder.py +1 -1
- diffusers/pipelines/unidiffuser/modeling_uvit.py +9 -9
- diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +23 -23
- diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +5 -5
- diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py +5 -10
- diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +4 -6
- diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +4 -4
- diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +12 -12
- diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +10 -10
- diffusers/schedulers/__init__.py +2 -2
- diffusers/schedulers/deprecated/__init__.py +1 -1
- diffusers/schedulers/deprecated/scheduling_karras_ve.py +25 -25
- diffusers/schedulers/scheduling_amused.py +5 -5
- diffusers/schedulers/scheduling_consistency_decoder.py +11 -11
- diffusers/schedulers/scheduling_consistency_models.py +23 -25
- diffusers/schedulers/scheduling_ddim.py +22 -24
- diffusers/schedulers/scheduling_ddim_flax.py +2 -1
- diffusers/schedulers/scheduling_ddim_inverse.py +16 -16
- diffusers/schedulers/scheduling_ddim_parallel.py +28 -30
- diffusers/schedulers/scheduling_ddpm.py +20 -22
- diffusers/schedulers/scheduling_ddpm_flax.py +7 -3
- diffusers/schedulers/scheduling_ddpm_parallel.py +26 -28
- diffusers/schedulers/scheduling_ddpm_wuerstchen.py +14 -14
- diffusers/schedulers/scheduling_deis_multistep.py +46 -42
- diffusers/schedulers/scheduling_dpmsolver_multistep.py +107 -77
- diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +2 -2
- diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +46 -46
- diffusers/schedulers/scheduling_dpmsolver_sde.py +26 -22
- diffusers/schedulers/scheduling_dpmsolver_singlestep.py +90 -65
- diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +78 -53
- diffusers/schedulers/scheduling_edm_euler.py +53 -30
- diffusers/schedulers/scheduling_euler_ancestral_discrete.py +26 -28
- diffusers/schedulers/scheduling_euler_discrete.py +163 -67
- diffusers/schedulers/scheduling_heun_discrete.py +60 -38
- diffusers/schedulers/scheduling_ipndm.py +8 -8
- diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +22 -18
- diffusers/schedulers/scheduling_k_dpm_2_discrete.py +22 -18
- diffusers/schedulers/scheduling_karras_ve_flax.py +6 -6
- diffusers/schedulers/scheduling_lcm.py +21 -23
- diffusers/schedulers/scheduling_lms_discrete.py +27 -25
- diffusers/schedulers/scheduling_pndm.py +20 -20
- diffusers/schedulers/scheduling_repaint.py +20 -20
- diffusers/schedulers/scheduling_sasolver.py +55 -54
- diffusers/schedulers/scheduling_sde_ve.py +19 -19
- diffusers/schedulers/scheduling_tcd.py +39 -30
- diffusers/schedulers/scheduling_unclip.py +15 -15
- diffusers/schedulers/scheduling_unipc_multistep.py +115 -41
- diffusers/schedulers/scheduling_utils.py +14 -5
- diffusers/schedulers/scheduling_utils_flax.py +3 -3
- diffusers/schedulers/scheduling_vq_diffusion.py +10 -10
- diffusers/training_utils.py +56 -1
- diffusers/utils/__init__.py +7 -0
- diffusers/utils/doc_utils.py +1 -0
- diffusers/utils/dummy_pt_objects.py +30 -0
- diffusers/utils/dummy_torch_and_transformers_objects.py +90 -0
- diffusers/utils/dynamic_modules_utils.py +24 -11
- diffusers/utils/hub_utils.py +3 -2
- diffusers/utils/import_utils.py +91 -0
- diffusers/utils/loading_utils.py +2 -2
- diffusers/utils/logging.py +1 -1
- diffusers/utils/peft_utils.py +32 -5
- diffusers/utils/state_dict_utils.py +11 -2
- diffusers/utils/testing_utils.py +71 -6
- diffusers/utils/torch_utils.py +1 -0
- diffusers/video_processor.py +113 -0
- {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/METADATA +7 -7
- diffusers-0.28.0.dist-info/RECORD +414 -0
- diffusers-0.27.1.dist-info/RECORD +0 -399
- {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/LICENSE +0 -0
- {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/WHEEL +0 -0
- {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/entry_points.txt +0 -0
- {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/top_level.txt +0 -0
@@ -35,8 +35,8 @@ class LEditsPPInversionPipelineOutput(BaseOutput):
|
|
35
35
|
List of the cropped and resized input images as PIL images of length `batch_size` or NumPy array of shape `
|
36
36
|
(batch_size, height, width, num_channels)`.
|
37
37
|
vae_reconstruction_images (`List[PIL.Image.Image]` or `np.ndarray`)
|
38
|
-
List of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape
|
39
|
-
(batch_size, height, width, num_channels)`.
|
38
|
+
List of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape
|
39
|
+
` (batch_size, height, width, num_channels)`.
|
40
40
|
"""
|
41
41
|
|
42
42
|
images: Union[List[PIL.Image.Image], np.ndarray]
|
@@ -0,0 +1,50 @@
|
|
1
|
+
from typing import TYPE_CHECKING
|
2
|
+
|
3
|
+
from ...utils import (
|
4
|
+
DIFFUSERS_SLOW_IMPORT,
|
5
|
+
OptionalDependencyNotAvailable,
|
6
|
+
_LazyModule,
|
7
|
+
get_objects_from_module,
|
8
|
+
is_torch_available,
|
9
|
+
is_transformers_available,
|
10
|
+
)
|
11
|
+
|
12
|
+
|
13
|
+
_dummy_objects = {}
|
14
|
+
_import_structure = {}
|
15
|
+
|
16
|
+
try:
|
17
|
+
if not (is_transformers_available() and is_torch_available()):
|
18
|
+
raise OptionalDependencyNotAvailable()
|
19
|
+
except OptionalDependencyNotAvailable:
|
20
|
+
from ...utils import dummy_torch_and_transformers_objects # noqa F403
|
21
|
+
|
22
|
+
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
|
23
|
+
else:
|
24
|
+
_import_structure["marigold_image_processing"] = ["MarigoldImageProcessor"]
|
25
|
+
_import_structure["pipeline_marigold_depth"] = ["MarigoldDepthOutput", "MarigoldDepthPipeline"]
|
26
|
+
_import_structure["pipeline_marigold_normals"] = ["MarigoldNormalsOutput", "MarigoldNormalsPipeline"]
|
27
|
+
|
28
|
+
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
29
|
+
try:
|
30
|
+
if not (is_transformers_available() and is_torch_available()):
|
31
|
+
raise OptionalDependencyNotAvailable()
|
32
|
+
|
33
|
+
except OptionalDependencyNotAvailable:
|
34
|
+
from ...utils.dummy_torch_and_transformers_objects import *
|
35
|
+
else:
|
36
|
+
from .marigold_image_processing import MarigoldImageProcessor
|
37
|
+
from .pipeline_marigold_depth import MarigoldDepthOutput, MarigoldDepthPipeline
|
38
|
+
from .pipeline_marigold_normals import MarigoldNormalsOutput, MarigoldNormalsPipeline
|
39
|
+
|
40
|
+
else:
|
41
|
+
import sys
|
42
|
+
|
43
|
+
sys.modules[__name__] = _LazyModule(
|
44
|
+
__name__,
|
45
|
+
globals()["__file__"],
|
46
|
+
_import_structure,
|
47
|
+
module_spec=__spec__,
|
48
|
+
)
|
49
|
+
for name, value in _dummy_objects.items():
|
50
|
+
setattr(sys.modules[__name__], name, value)
|
@@ -0,0 +1,561 @@
|
|
1
|
+
from typing import List, Optional, Tuple, Union
|
2
|
+
|
3
|
+
import numpy as np
|
4
|
+
import PIL
|
5
|
+
import torch
|
6
|
+
import torch.nn.functional as F
|
7
|
+
from PIL import Image
|
8
|
+
|
9
|
+
from ... import ConfigMixin
|
10
|
+
from ...configuration_utils import register_to_config
|
11
|
+
from ...image_processor import PipelineImageInput
|
12
|
+
from ...utils import CONFIG_NAME, logging
|
13
|
+
from ...utils.import_utils import is_matplotlib_available
|
14
|
+
|
15
|
+
|
16
|
+
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
17
|
+
|
18
|
+
|
19
|
+
class MarigoldImageProcessor(ConfigMixin):
|
20
|
+
config_name = CONFIG_NAME
|
21
|
+
|
22
|
+
@register_to_config
|
23
|
+
def __init__(
|
24
|
+
self,
|
25
|
+
vae_scale_factor: int = 8,
|
26
|
+
do_normalize: bool = True,
|
27
|
+
do_range_check: bool = True,
|
28
|
+
):
|
29
|
+
super().__init__()
|
30
|
+
|
31
|
+
@staticmethod
|
32
|
+
def expand_tensor_or_array(images: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
|
33
|
+
"""
|
34
|
+
Expand a tensor or array to a specified number of images.
|
35
|
+
"""
|
36
|
+
if isinstance(images, np.ndarray):
|
37
|
+
if images.ndim == 2: # [H,W] -> [1,H,W,1]
|
38
|
+
images = images[None, ..., None]
|
39
|
+
if images.ndim == 3: # [H,W,C] -> [1,H,W,C]
|
40
|
+
images = images[None]
|
41
|
+
elif isinstance(images, torch.Tensor):
|
42
|
+
if images.ndim == 2: # [H,W] -> [1,1,H,W]
|
43
|
+
images = images[None, None]
|
44
|
+
elif images.ndim == 3: # [1,H,W] -> [1,1,H,W]
|
45
|
+
images = images[None]
|
46
|
+
else:
|
47
|
+
raise ValueError(f"Unexpected input type: {type(images)}")
|
48
|
+
return images
|
49
|
+
|
50
|
+
@staticmethod
|
51
|
+
def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
|
52
|
+
"""
|
53
|
+
Convert a PyTorch tensor to a NumPy image.
|
54
|
+
"""
|
55
|
+
images = images.cpu().permute(0, 2, 3, 1).float().numpy()
|
56
|
+
return images
|
57
|
+
|
58
|
+
@staticmethod
|
59
|
+
def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
|
60
|
+
"""
|
61
|
+
Convert a NumPy image to a PyTorch tensor.
|
62
|
+
"""
|
63
|
+
if np.issubdtype(images.dtype, np.integer) and not np.issubdtype(images.dtype, np.unsignedinteger):
|
64
|
+
raise ValueError(f"Input image dtype={images.dtype} cannot be a signed integer.")
|
65
|
+
if np.issubdtype(images.dtype, np.complexfloating):
|
66
|
+
raise ValueError(f"Input image dtype={images.dtype} cannot be complex.")
|
67
|
+
if np.issubdtype(images.dtype, bool):
|
68
|
+
raise ValueError(f"Input image dtype={images.dtype} cannot be boolean.")
|
69
|
+
|
70
|
+
images = torch.from_numpy(images.transpose(0, 3, 1, 2))
|
71
|
+
return images
|
72
|
+
|
73
|
+
@staticmethod
|
74
|
+
def resize_antialias(
|
75
|
+
image: torch.Tensor, size: Tuple[int, int], mode: str, is_aa: Optional[bool] = None
|
76
|
+
) -> torch.Tensor:
|
77
|
+
if not torch.is_tensor(image):
|
78
|
+
raise ValueError(f"Invalid input type={type(image)}.")
|
79
|
+
if not torch.is_floating_point(image):
|
80
|
+
raise ValueError(f"Invalid input dtype={image.dtype}.")
|
81
|
+
if image.dim() != 4:
|
82
|
+
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
|
83
|
+
|
84
|
+
antialias = is_aa and mode in ("bilinear", "bicubic")
|
85
|
+
image = F.interpolate(image, size, mode=mode, antialias=antialias)
|
86
|
+
|
87
|
+
return image
|
88
|
+
|
89
|
+
@staticmethod
|
90
|
+
def resize_to_max_edge(image: torch.Tensor, max_edge_sz: int, mode: str) -> torch.Tensor:
|
91
|
+
if not torch.is_tensor(image):
|
92
|
+
raise ValueError(f"Invalid input type={type(image)}.")
|
93
|
+
if not torch.is_floating_point(image):
|
94
|
+
raise ValueError(f"Invalid input dtype={image.dtype}.")
|
95
|
+
if image.dim() != 4:
|
96
|
+
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
|
97
|
+
|
98
|
+
h, w = image.shape[-2:]
|
99
|
+
max_orig = max(h, w)
|
100
|
+
new_h = h * max_edge_sz // max_orig
|
101
|
+
new_w = w * max_edge_sz // max_orig
|
102
|
+
|
103
|
+
if new_h == 0 or new_w == 0:
|
104
|
+
raise ValueError(f"Extreme aspect ratio of the input image: [{w} x {h}]")
|
105
|
+
|
106
|
+
image = MarigoldImageProcessor.resize_antialias(image, (new_h, new_w), mode, is_aa=True)
|
107
|
+
|
108
|
+
return image
|
109
|
+
|
110
|
+
@staticmethod
|
111
|
+
def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
|
112
|
+
if not torch.is_tensor(image):
|
113
|
+
raise ValueError(f"Invalid input type={type(image)}.")
|
114
|
+
if not torch.is_floating_point(image):
|
115
|
+
raise ValueError(f"Invalid input dtype={image.dtype}.")
|
116
|
+
if image.dim() != 4:
|
117
|
+
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
|
118
|
+
|
119
|
+
h, w = image.shape[-2:]
|
120
|
+
ph, pw = -h % align, -w % align
|
121
|
+
|
122
|
+
image = F.pad(image, (0, pw, 0, ph), mode="replicate")
|
123
|
+
|
124
|
+
return image, (ph, pw)
|
125
|
+
|
126
|
+
@staticmethod
|
127
|
+
def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor:
|
128
|
+
if not torch.is_tensor(image):
|
129
|
+
raise ValueError(f"Invalid input type={type(image)}.")
|
130
|
+
if not torch.is_floating_point(image):
|
131
|
+
raise ValueError(f"Invalid input dtype={image.dtype}.")
|
132
|
+
if image.dim() != 4:
|
133
|
+
raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
|
134
|
+
|
135
|
+
ph, pw = padding
|
136
|
+
uh = None if ph == 0 else -ph
|
137
|
+
uw = None if pw == 0 else -pw
|
138
|
+
|
139
|
+
image = image[:, :, :uh, :uw]
|
140
|
+
|
141
|
+
return image
|
142
|
+
|
143
|
+
@staticmethod
|
144
|
+
def load_image_canonical(
|
145
|
+
image: Union[torch.Tensor, np.ndarray, Image.Image],
|
146
|
+
device: torch.device = torch.device("cpu"),
|
147
|
+
dtype: torch.dtype = torch.float32,
|
148
|
+
) -> Tuple[torch.Tensor, int]:
|
149
|
+
if isinstance(image, Image.Image):
|
150
|
+
image = np.array(image)
|
151
|
+
|
152
|
+
image_dtype_max = None
|
153
|
+
if isinstance(image, (np.ndarray, torch.Tensor)):
|
154
|
+
image = MarigoldImageProcessor.expand_tensor_or_array(image)
|
155
|
+
if image.ndim != 4:
|
156
|
+
raise ValueError("Input image is not 2-, 3-, or 4-dimensional.")
|
157
|
+
if isinstance(image, np.ndarray):
|
158
|
+
if np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.unsignedinteger):
|
159
|
+
raise ValueError(f"Input image dtype={image.dtype} cannot be a signed integer.")
|
160
|
+
if np.issubdtype(image.dtype, np.complexfloating):
|
161
|
+
raise ValueError(f"Input image dtype={image.dtype} cannot be complex.")
|
162
|
+
if np.issubdtype(image.dtype, bool):
|
163
|
+
raise ValueError(f"Input image dtype={image.dtype} cannot be boolean.")
|
164
|
+
if np.issubdtype(image.dtype, np.unsignedinteger):
|
165
|
+
image_dtype_max = np.iinfo(image.dtype).max
|
166
|
+
image = image.astype(np.float32) # because torch does not have unsigned dtypes beyond torch.uint8
|
167
|
+
image = MarigoldImageProcessor.numpy_to_pt(image)
|
168
|
+
|
169
|
+
if torch.is_tensor(image) and not torch.is_floating_point(image) and image_dtype_max is None:
|
170
|
+
if image.dtype != torch.uint8:
|
171
|
+
raise ValueError(f"Image dtype={image.dtype} is not supported.")
|
172
|
+
image_dtype_max = 255
|
173
|
+
|
174
|
+
if not torch.is_tensor(image):
|
175
|
+
raise ValueError(f"Input type unsupported: {type(image)}.")
|
176
|
+
|
177
|
+
if image.shape[1] == 1:
|
178
|
+
image = image.repeat(1, 3, 1, 1) # [N,1,H,W] -> [N,3,H,W]
|
179
|
+
if image.shape[1] != 3:
|
180
|
+
raise ValueError(f"Input image is not 1- or 3-channel: {image.shape}.")
|
181
|
+
|
182
|
+
image = image.to(device=device, dtype=dtype)
|
183
|
+
|
184
|
+
if image_dtype_max is not None:
|
185
|
+
image = image / image_dtype_max
|
186
|
+
|
187
|
+
return image
|
188
|
+
|
189
|
+
@staticmethod
|
190
|
+
def check_image_values_range(image: torch.Tensor) -> None:
|
191
|
+
if not torch.is_tensor(image):
|
192
|
+
raise ValueError(f"Invalid input type={type(image)}.")
|
193
|
+
if not torch.is_floating_point(image):
|
194
|
+
raise ValueError(f"Invalid input dtype={image.dtype}.")
|
195
|
+
if image.min().item() < 0.0 or image.max().item() > 1.0:
|
196
|
+
raise ValueError("Input image data is partially outside of the [0,1] range.")
|
197
|
+
|
198
|
+
def preprocess(
|
199
|
+
self,
|
200
|
+
image: PipelineImageInput,
|
201
|
+
processing_resolution: Optional[int] = None,
|
202
|
+
resample_method_input: str = "bilinear",
|
203
|
+
device: torch.device = torch.device("cpu"),
|
204
|
+
dtype: torch.dtype = torch.float32,
|
205
|
+
):
|
206
|
+
if isinstance(image, list):
|
207
|
+
images = None
|
208
|
+
for i, img in enumerate(image):
|
209
|
+
img = self.load_image_canonical(img, device, dtype) # [N,3,H,W]
|
210
|
+
if images is None:
|
211
|
+
images = img
|
212
|
+
else:
|
213
|
+
if images.shape[2:] != img.shape[2:]:
|
214
|
+
raise ValueError(
|
215
|
+
f"Input image[{i}] has incompatible dimensions {img.shape[2:]} with the previous images "
|
216
|
+
f"{images.shape[2:]}"
|
217
|
+
)
|
218
|
+
images = torch.cat((images, img), dim=0)
|
219
|
+
image = images
|
220
|
+
del images
|
221
|
+
else:
|
222
|
+
image = self.load_image_canonical(image, device, dtype) # [N,3,H,W]
|
223
|
+
|
224
|
+
original_resolution = image.shape[2:]
|
225
|
+
|
226
|
+
if self.config.do_range_check:
|
227
|
+
self.check_image_values_range(image)
|
228
|
+
|
229
|
+
if self.config.do_normalize:
|
230
|
+
image = image * 2.0 - 1.0
|
231
|
+
|
232
|
+
if processing_resolution is not None and processing_resolution > 0:
|
233
|
+
image = self.resize_to_max_edge(image, processing_resolution, resample_method_input) # [N,3,PH,PW]
|
234
|
+
|
235
|
+
image, padding = self.pad_image(image, self.config.vae_scale_factor) # [N,3,PPH,PPW]
|
236
|
+
|
237
|
+
return image, padding, original_resolution
|
238
|
+
|
239
|
+
@staticmethod
|
240
|
+
def colormap(
|
241
|
+
image: Union[np.ndarray, torch.Tensor],
|
242
|
+
cmap: str = "Spectral",
|
243
|
+
bytes: bool = False,
|
244
|
+
_force_method: Optional[str] = None,
|
245
|
+
) -> Union[np.ndarray, torch.Tensor]:
|
246
|
+
"""
|
247
|
+
Converts a monochrome image into an RGB image by applying the specified colormap. This function mimics the
|
248
|
+
behavior of matplotlib.colormaps, but allows the user to use the most discriminative color map "Spectral"
|
249
|
+
without having to install or import matplotlib. For all other cases, the function will attempt to use the
|
250
|
+
native implementation.
|
251
|
+
|
252
|
+
Args:
|
253
|
+
image: 2D tensor of values between 0 and 1, either as np.ndarray or torch.Tensor.
|
254
|
+
cmap: Colormap name.
|
255
|
+
bytes: Whether to return the output as uint8 or floating point image.
|
256
|
+
_force_method:
|
257
|
+
Can be used to specify whether to use the native implementation (`"matplotlib"`), the efficient custom
|
258
|
+
implementation of the "Spectral" color map (`"custom"`), or rely on autodetection (`None`, default).
|
259
|
+
|
260
|
+
Returns:
|
261
|
+
An RGB-colorized tensor corresponding to the input image.
|
262
|
+
"""
|
263
|
+
if not (torch.is_tensor(image) or isinstance(image, np.ndarray)):
|
264
|
+
raise ValueError("Argument must be a numpy array or torch tensor.")
|
265
|
+
if _force_method not in (None, "matplotlib", "custom"):
|
266
|
+
raise ValueError("_force_method must be either `None`, `'matplotlib'` or `'custom'`.")
|
267
|
+
|
268
|
+
def method_matplotlib(image, cmap, bytes=False):
|
269
|
+
if is_matplotlib_available():
|
270
|
+
import matplotlib
|
271
|
+
else:
|
272
|
+
return None
|
273
|
+
|
274
|
+
arg_is_pt, device = torch.is_tensor(image), None
|
275
|
+
if arg_is_pt:
|
276
|
+
image, device = image.cpu().numpy(), image.device
|
277
|
+
|
278
|
+
if cmap not in matplotlib.colormaps:
|
279
|
+
raise ValueError(
|
280
|
+
f"Unexpected color map {cmap}; available options are: {', '.join(list(matplotlib.colormaps.keys()))}"
|
281
|
+
)
|
282
|
+
|
283
|
+
cmap = matplotlib.colormaps[cmap]
|
284
|
+
out = cmap(image, bytes=bytes) # [?,4]
|
285
|
+
out = out[..., :3] # [?,3]
|
286
|
+
|
287
|
+
if arg_is_pt:
|
288
|
+
out = torch.tensor(out, device=device)
|
289
|
+
|
290
|
+
return out
|
291
|
+
|
292
|
+
def method_custom(image, cmap, bytes=False):
|
293
|
+
arg_is_np = isinstance(image, np.ndarray)
|
294
|
+
if arg_is_np:
|
295
|
+
image = torch.tensor(image)
|
296
|
+
if image.dtype == torch.uint8:
|
297
|
+
image = image.float() / 255
|
298
|
+
else:
|
299
|
+
image = image.float()
|
300
|
+
|
301
|
+
if cmap != "Spectral":
|
302
|
+
raise ValueError("Only 'Spectral' color map is available without installing matplotlib.")
|
303
|
+
|
304
|
+
_Spectral_data = ( # Taken from matplotlib/_cm.py
|
305
|
+
(0.61960784313725492, 0.003921568627450980, 0.25882352941176473), # 0.0 -> [0]
|
306
|
+
(0.83529411764705885, 0.24313725490196078, 0.30980392156862746),
|
307
|
+
(0.95686274509803926, 0.42745098039215684, 0.2627450980392157),
|
308
|
+
(0.99215686274509807, 0.68235294117647061, 0.38039215686274508),
|
309
|
+
(0.99607843137254903, 0.8784313725490196, 0.54509803921568623),
|
310
|
+
(1.0, 1.0, 0.74901960784313726),
|
311
|
+
(0.90196078431372551, 0.96078431372549022, 0.59607843137254901),
|
312
|
+
(0.6705882352941176, 0.8666666666666667, 0.64313725490196083),
|
313
|
+
(0.4, 0.76078431372549016, 0.6470588235294118),
|
314
|
+
(0.19607843137254902, 0.53333333333333333, 0.74117647058823533),
|
315
|
+
(0.36862745098039218, 0.30980392156862746, 0.63529411764705879), # 1.0 -> [K-1]
|
316
|
+
)
|
317
|
+
|
318
|
+
cmap = torch.tensor(_Spectral_data, dtype=torch.float, device=image.device) # [K,3]
|
319
|
+
K = cmap.shape[0]
|
320
|
+
|
321
|
+
pos = image.clamp(min=0, max=1) * (K - 1)
|
322
|
+
left = pos.long()
|
323
|
+
right = (left + 1).clamp(max=K - 1)
|
324
|
+
|
325
|
+
d = (pos - left.float()).unsqueeze(-1)
|
326
|
+
left_colors = cmap[left]
|
327
|
+
right_colors = cmap[right]
|
328
|
+
|
329
|
+
out = (1 - d) * left_colors + d * right_colors
|
330
|
+
|
331
|
+
if bytes:
|
332
|
+
out = (out * 255).to(torch.uint8)
|
333
|
+
|
334
|
+
if arg_is_np:
|
335
|
+
out = out.numpy()
|
336
|
+
|
337
|
+
return out
|
338
|
+
|
339
|
+
if _force_method is None and torch.is_tensor(image) and cmap == "Spectral":
|
340
|
+
return method_custom(image, cmap, bytes)
|
341
|
+
|
342
|
+
out = None
|
343
|
+
if _force_method != "custom":
|
344
|
+
out = method_matplotlib(image, cmap, bytes)
|
345
|
+
|
346
|
+
if _force_method == "matplotlib" and out is None:
|
347
|
+
raise ImportError("Make sure to install matplotlib if you want to use a color map other than 'Spectral'.")
|
348
|
+
|
349
|
+
if out is None:
|
350
|
+
out = method_custom(image, cmap, bytes)
|
351
|
+
|
352
|
+
return out
|
353
|
+
|
354
|
+
@staticmethod
|
355
|
+
def visualize_depth(
|
356
|
+
depth: Union[
|
357
|
+
PIL.Image.Image,
|
358
|
+
np.ndarray,
|
359
|
+
torch.Tensor,
|
360
|
+
List[PIL.Image.Image],
|
361
|
+
List[np.ndarray],
|
362
|
+
List[torch.Tensor],
|
363
|
+
],
|
364
|
+
val_min: float = 0.0,
|
365
|
+
val_max: float = 1.0,
|
366
|
+
color_map: str = "Spectral",
|
367
|
+
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
|
368
|
+
"""
|
369
|
+
Visualizes depth maps, such as predictions of the `MarigoldDepthPipeline`.
|
370
|
+
|
371
|
+
Args:
|
372
|
+
depth (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image], List[np.ndarray],
|
373
|
+
List[torch.Tensor]]`): Depth maps.
|
374
|
+
val_min (`float`, *optional*, defaults to `0.0`): Minimum value of the visualized depth range.
|
375
|
+
val_max (`float`, *optional*, defaults to `1.0`): Maximum value of the visualized depth range.
|
376
|
+
color_map (`str`, *optional*, defaults to `"Spectral"`): Color map used to convert a single-channel
|
377
|
+
depth prediction into colored representation.
|
378
|
+
|
379
|
+
Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with depth maps visualization.
|
380
|
+
"""
|
381
|
+
if val_max <= val_min:
|
382
|
+
raise ValueError(f"Invalid values range: [{val_min}, {val_max}].")
|
383
|
+
|
384
|
+
def visualize_depth_one(img, idx=None):
|
385
|
+
prefix = "Depth" + (f"[{idx}]" if idx else "")
|
386
|
+
if isinstance(img, PIL.Image.Image):
|
387
|
+
if img.mode != "I;16":
|
388
|
+
raise ValueError(f"{prefix}: invalid PIL mode={img.mode}.")
|
389
|
+
img = np.array(img).astype(np.float32) / (2**16 - 1)
|
390
|
+
if isinstance(img, np.ndarray) or torch.is_tensor(img):
|
391
|
+
if img.ndim != 2:
|
392
|
+
raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
|
393
|
+
if isinstance(img, np.ndarray):
|
394
|
+
img = torch.from_numpy(img)
|
395
|
+
if not torch.is_floating_point(img):
|
396
|
+
raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
|
397
|
+
else:
|
398
|
+
raise ValueError(f"{prefix}: unexpected type={type(img)}.")
|
399
|
+
if val_min != 0.0 or val_max != 1.0:
|
400
|
+
img = (img - val_min) / (val_max - val_min)
|
401
|
+
img = MarigoldImageProcessor.colormap(img, cmap=color_map, bytes=True) # [H,W,3]
|
402
|
+
img = PIL.Image.fromarray(img.cpu().numpy())
|
403
|
+
return img
|
404
|
+
|
405
|
+
if depth is None or isinstance(depth, list) and any(o is None for o in depth):
|
406
|
+
raise ValueError("Input depth is `None`")
|
407
|
+
if isinstance(depth, (np.ndarray, torch.Tensor)):
|
408
|
+
depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
|
409
|
+
if isinstance(depth, np.ndarray):
|
410
|
+
depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W]
|
411
|
+
if not (depth.ndim == 4 and depth.shape[1] == 1): # [N,1,H,W]
|
412
|
+
raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
|
413
|
+
return [visualize_depth_one(img[0], idx) for idx, img in enumerate(depth)]
|
414
|
+
elif isinstance(depth, list):
|
415
|
+
return [visualize_depth_one(img, idx) for idx, img in enumerate(depth)]
|
416
|
+
else:
|
417
|
+
raise ValueError(f"Unexpected input type: {type(depth)}")
|
418
|
+
|
419
|
+
@staticmethod
|
420
|
+
def export_depth_to_16bit_png(
|
421
|
+
depth: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
|
422
|
+
val_min: float = 0.0,
|
423
|
+
val_max: float = 1.0,
|
424
|
+
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
|
425
|
+
def export_depth_to_16bit_png_one(img, idx=None):
|
426
|
+
prefix = "Depth" + (f"[{idx}]" if idx else "")
|
427
|
+
if not isinstance(img, np.ndarray) and not torch.is_tensor(img):
|
428
|
+
raise ValueError(f"{prefix}: unexpected type={type(img)}.")
|
429
|
+
if img.ndim != 2:
|
430
|
+
raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
|
431
|
+
if torch.is_tensor(img):
|
432
|
+
img = img.cpu().numpy()
|
433
|
+
if not np.issubdtype(img.dtype, np.floating):
|
434
|
+
raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
|
435
|
+
if val_min != 0.0 or val_max != 1.0:
|
436
|
+
img = (img - val_min) / (val_max - val_min)
|
437
|
+
img = (img * (2**16 - 1)).astype(np.uint16)
|
438
|
+
img = PIL.Image.fromarray(img, mode="I;16")
|
439
|
+
return img
|
440
|
+
|
441
|
+
if depth is None or isinstance(depth, list) and any(o is None for o in depth):
|
442
|
+
raise ValueError("Input depth is `None`")
|
443
|
+
if isinstance(depth, (np.ndarray, torch.Tensor)):
|
444
|
+
depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
|
445
|
+
if isinstance(depth, np.ndarray):
|
446
|
+
depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W]
|
447
|
+
if not (depth.ndim == 4 and depth.shape[1] == 1):
|
448
|
+
raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
|
449
|
+
return [export_depth_to_16bit_png_one(img[0], idx) for idx, img in enumerate(depth)]
|
450
|
+
elif isinstance(depth, list):
|
451
|
+
return [export_depth_to_16bit_png_one(img, idx) for idx, img in enumerate(depth)]
|
452
|
+
else:
|
453
|
+
raise ValueError(f"Unexpected input type: {type(depth)}")
|
454
|
+
|
455
|
+
@staticmethod
|
456
|
+
def visualize_normals(
|
457
|
+
normals: Union[
|
458
|
+
np.ndarray,
|
459
|
+
torch.Tensor,
|
460
|
+
List[np.ndarray],
|
461
|
+
List[torch.Tensor],
|
462
|
+
],
|
463
|
+
flip_x: bool = False,
|
464
|
+
flip_y: bool = False,
|
465
|
+
flip_z: bool = False,
|
466
|
+
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
|
467
|
+
"""
|
468
|
+
Visualizes surface normals, such as predictions of the `MarigoldNormalsPipeline`.
|
469
|
+
|
470
|
+
Args:
|
471
|
+
normals (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
|
472
|
+
Surface normals.
|
473
|
+
flip_x (`bool`, *optional*, defaults to `False`): Flips the X axis of the normals frame of reference.
|
474
|
+
Default direction is right.
|
475
|
+
flip_y (`bool`, *optional*, defaults to `False`): Flips the Y axis of the normals frame of reference.
|
476
|
+
Default direction is top.
|
477
|
+
flip_z (`bool`, *optional*, defaults to `False`): Flips the Z axis of the normals frame of reference.
|
478
|
+
Default direction is facing the observer.
|
479
|
+
|
480
|
+
Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with surface normals visualization.
|
481
|
+
"""
|
482
|
+
flip_vec = None
|
483
|
+
if any((flip_x, flip_y, flip_z)):
|
484
|
+
flip_vec = torch.tensor(
|
485
|
+
[
|
486
|
+
(-1) ** flip_x,
|
487
|
+
(-1) ** flip_y,
|
488
|
+
(-1) ** flip_z,
|
489
|
+
],
|
490
|
+
dtype=torch.float32,
|
491
|
+
)
|
492
|
+
|
493
|
+
def visualize_normals_one(img, idx=None):
|
494
|
+
img = img.permute(1, 2, 0)
|
495
|
+
if flip_vec is not None:
|
496
|
+
img *= flip_vec.to(img.device)
|
497
|
+
img = (img + 1.0) * 0.5
|
498
|
+
img = (img * 255).to(dtype=torch.uint8, device="cpu").numpy()
|
499
|
+
img = PIL.Image.fromarray(img)
|
500
|
+
return img
|
501
|
+
|
502
|
+
if normals is None or isinstance(normals, list) and any(o is None for o in normals):
|
503
|
+
raise ValueError("Input normals is `None`")
|
504
|
+
if isinstance(normals, (np.ndarray, torch.Tensor)):
|
505
|
+
normals = MarigoldImageProcessor.expand_tensor_or_array(normals)
|
506
|
+
if isinstance(normals, np.ndarray):
|
507
|
+
normals = MarigoldImageProcessor.numpy_to_pt(normals) # [N,3,H,W]
|
508
|
+
if not (normals.ndim == 4 and normals.shape[1] == 3):
|
509
|
+
raise ValueError(f"Unexpected input shape={normals.shape}, expecting [N,3,H,W].")
|
510
|
+
return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
|
511
|
+
elif isinstance(normals, list):
|
512
|
+
return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
|
513
|
+
else:
|
514
|
+
raise ValueError(f"Unexpected input type: {type(normals)}")
|
515
|
+
|
516
|
+
@staticmethod
|
517
|
+
def visualize_uncertainty(
|
518
|
+
uncertainty: Union[
|
519
|
+
np.ndarray,
|
520
|
+
torch.Tensor,
|
521
|
+
List[np.ndarray],
|
522
|
+
List[torch.Tensor],
|
523
|
+
],
|
524
|
+
saturation_percentile=95,
|
525
|
+
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
|
526
|
+
"""
|
527
|
+
Visualizes dense uncertainties, such as produced by `MarigoldDepthPipeline` or `MarigoldNormalsPipeline`.
|
528
|
+
|
529
|
+
Args:
|
530
|
+
uncertainty (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
|
531
|
+
Uncertainty maps.
|
532
|
+
saturation_percentile (`int`, *optional*, defaults to `95`):
|
533
|
+
Specifies the percentile uncertainty value visualized with maximum intensity.
|
534
|
+
|
535
|
+
Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with uncertainty visualization.
|
536
|
+
"""
|
537
|
+
|
538
|
+
def visualize_uncertainty_one(img, idx=None):
|
539
|
+
prefix = "Uncertainty" + (f"[{idx}]" if idx else "")
|
540
|
+
if img.min() < 0:
|
541
|
+
raise ValueError(f"{prefix}: unexected data range, min={img.min()}.")
|
542
|
+
img = img.squeeze(0).cpu().numpy()
|
543
|
+
saturation_value = np.percentile(img, saturation_percentile)
|
544
|
+
img = np.clip(img * 255 / saturation_value, 0, 255)
|
545
|
+
img = img.astype(np.uint8)
|
546
|
+
img = PIL.Image.fromarray(img)
|
547
|
+
return img
|
548
|
+
|
549
|
+
if uncertainty is None or isinstance(uncertainty, list) and any(o is None for o in uncertainty):
|
550
|
+
raise ValueError("Input uncertainty is `None`")
|
551
|
+
if isinstance(uncertainty, (np.ndarray, torch.Tensor)):
|
552
|
+
uncertainty = MarigoldImageProcessor.expand_tensor_or_array(uncertainty)
|
553
|
+
if isinstance(uncertainty, np.ndarray):
|
554
|
+
uncertainty = MarigoldImageProcessor.numpy_to_pt(uncertainty) # [N,1,H,W]
|
555
|
+
if not (uncertainty.ndim == 4 and uncertainty.shape[1] == 1):
|
556
|
+
raise ValueError(f"Unexpected input shape={uncertainty.shape}, expecting [N,1,H,W].")
|
557
|
+
return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)]
|
558
|
+
elif isinstance(uncertainty, list):
|
559
|
+
return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)]
|
560
|
+
else:
|
561
|
+
raise ValueError(f"Unexpected input type: {type(uncertainty)}")
|