libcom 0.1.8__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.
- libcom/__init__.py +23 -0
- libcom/color_transfer/__init__.py +3 -0
- libcom/color_transfer/reinhard.py +57 -0
- libcom/fopa_heat_map/__init__.py +4 -0
- libcom/fopa_heat_map/fopa_heat_map.py +212 -0
- libcom/fopa_heat_map/source/__init__.py +0 -0
- libcom/fopa_heat_map/source/backbone/ResNet.py +229 -0
- libcom/fopa_heat_map/source/backbone/__init__.py +0 -0
- libcom/fopa_heat_map/source/data/OBdataset.py +292 -0
- libcom/fopa_heat_map/source/data/__init__.py +0 -0
- libcom/fopa_heat_map/source/data/all_transforms.py +25 -0
- libcom/fopa_heat_map/source/network/BaseBlocks.py +26 -0
- libcom/fopa_heat_map/source/network/DynamicModules.py +43 -0
- libcom/fopa_heat_map/source/network/ObPlaNet_simple.py +112 -0
- libcom/fopa_heat_map/source/network/__init__.py +1 -0
- libcom/fopa_heat_map/source/network/tensor_ops.py +37 -0
- libcom/fopa_heat_map/source/prepare_multi_fg_scales.py +85 -0
- libcom/fos_score/__init__.py +4 -0
- libcom/fos_score/fos_score_prediction.py +218 -0
- libcom/fos_score/source/__init__.py +0 -0
- libcom/fos_score/source/config/__init__.py +1 -0
- libcom/fos_score/source/config/config.py +68 -0
- libcom/fos_score/source/config/config.sfosd.yaml +43 -0
- libcom/fos_score/source/config/config_rfosd.yaml +43 -0
- libcom/fos_score/source/loss_function.py +118 -0
- libcom/fos_score/source/networks.py +355 -0
- libcom/harmony_score/__init__.py +3 -0
- libcom/harmony_score/harmony_score_prediction.py +119 -0
- libcom/harmony_score/source/__init__.py +0 -0
- libcom/harmony_score/source/bargainnet.py +226 -0
- libcom/image_harmonization/__init__.py +4 -0
- libcom/image_harmonization/image_harmonization.py +193 -0
- libcom/image_harmonization/source/__init__.py +0 -0
- libcom/image_harmonization/source/functions.py +747 -0
- libcom/image_harmonization/source/pct_net.py +85 -0
- libcom/image_harmonization/source/src/__init__.py +0 -0
- libcom/image_harmonization/source/src/lbm/__init__.py +0 -0
- libcom/image_harmonization/source/src/lbm/config.py +141 -0
- libcom/image_harmonization/source/src/lbm/data/__init__.py +62 -0
- libcom/image_harmonization/source/src/lbm/data/datasets/__init__.py +9 -0
- libcom/image_harmonization/source/src/lbm/data/datasets/collation_fn.py +41 -0
- libcom/image_harmonization/source/src/lbm/data/datasets/dataset.py +243 -0
- libcom/image_harmonization/source/src/lbm/data/datasets/datasets_config.py +42 -0
- libcom/image_harmonization/source/src/lbm/data/filters/__init__.py +12 -0
- libcom/image_harmonization/source/src/lbm/data/filters/base.py +21 -0
- libcom/image_harmonization/source/src/lbm/data/filters/filter_wrapper.py +36 -0
- libcom/image_harmonization/source/src/lbm/data/filters/filters.py +33 -0
- libcom/image_harmonization/source/src/lbm/data/filters/filters_config.py +32 -0
- libcom/image_harmonization/source/src/lbm/data/mappers/__init__.py +19 -0
- libcom/image_harmonization/source/src/lbm/data/mappers/base.py +26 -0
- libcom/image_harmonization/source/src/lbm/data/mappers/mappers.py +135 -0
- libcom/image_harmonization/source/src/lbm/data/mappers/mappers_config.py +109 -0
- libcom/image_harmonization/source/src/lbm/data/mappers/mappers_wrapper.py +31 -0
- libcom/image_harmonization/source/src/lbm/inference/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/inference/inference.py +70 -0
- libcom/image_harmonization/source/src/lbm/inference/utils.py +222 -0
- libcom/image_harmonization/source/src/lbm/models/__init__.py +0 -0
- libcom/image_harmonization/source/src/lbm/models/base/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/models/base/base_model.py +66 -0
- libcom/image_harmonization/source/src/lbm/models/base/model_config.py +8 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/base/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner.py +60 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner_config.py +27 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/conditioners_wrapper.py +114 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_config.py +31 -0
- libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_model.py +80 -0
- libcom/image_harmonization/source/src/lbm/models/lbm/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/models/lbm/lbm_config.py +101 -0
- libcom/image_harmonization/source/src/lbm/models/lbm/lbm_model.py +359 -0
- libcom/image_harmonization/source/src/lbm/models/unets/__init__.py +14 -0
- libcom/image_harmonization/source/src/lbm/models/unets/unet.py +148 -0
- libcom/image_harmonization/source/src/lbm/models/utils.py +377 -0
- libcom/image_harmonization/source/src/lbm/models/vae/__init__.py +4 -0
- libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL.py +136 -0
- libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL_config.py +27 -0
- libcom/image_harmonization/source/src/lbm/trainer/__init__.py +85 -0
- libcom/image_harmonization/source/src/lbm/trainer/loggers.py +324 -0
- libcom/image_harmonization/source/src/lbm/trainer/trainer.py +199 -0
- libcom/image_harmonization/source/src/lbm/trainer/training_config.py +82 -0
- libcom/image_harmonization/source/src/lbm/trainer/utils.py +193 -0
- libcom/inharmonious_region_localization/__init__.py +4 -0
- libcom/inharmonious_region_localization/inharmonious_region_localization.py +97 -0
- libcom/inharmonious_region_localization/source/__init__.py +0 -0
- libcom/inharmonious_region_localization/source/blocks.py +368 -0
- libcom/inharmonious_region_localization/source/functions.py +679 -0
- libcom/inharmonious_region_localization/source/madis_net.py +79 -0
- libcom/kontext_blending_harmonization/__init__.py +3 -0
- libcom/kontext_blending_harmonization/kontext_blending_harmonization.py +148 -0
- libcom/kontext_blending_harmonization/source/__init__.py +0 -0
- libcom/kontext_blending_harmonization/source/pipeline.py +1135 -0
- libcom/naive_composition/__init__.py +3 -0
- libcom/naive_composition/generate_composite_image.py +116 -0
- libcom/opa_score/__init__.py +3 -0
- libcom/opa_score/opa_score_prediction.py +101 -0
- libcom/opa_score/source/__init__.py +3 -0
- libcom/opa_score/source/object_place_net.py +55 -0
- libcom/opa_score/source/oblect_place_config.py +66 -0
- libcom/opa_score/source/resnet_4ch.py +213 -0
- libcom/os_insert/__init__.py +3 -0
- libcom/os_insert/diffusers_osinsert/__init__.py +78 -0
- libcom/os_insert/diffusers_osinsert/_patched_pipeline_flux_fill.py +878 -0
- libcom/os_insert/os_insert.py +534 -0
- libcom/os_insert/source/__init__.py +1 -0
- libcom/os_insert/source/ia_utils.py +183 -0
- libcom/os_insert/source/insertanything_infer.py +529 -0
- libcom/os_insert/source/ldm/__init__.py +0 -0
- libcom/os_insert/source/ldm/data/__init__.py +0 -0
- libcom/os_insert/source/ldm/data/base.py +23 -0
- libcom/os_insert/source/ldm/data/datamodule.py +128 -0
- libcom/os_insert/source/ldm/data/imagenet.py +394 -0
- libcom/os_insert/source/ldm/data/lsun.py +92 -0
- libcom/os_insert/source/ldm/data/open_images.py +728 -0
- libcom/os_insert/source/ldm/lr_scheduler.py +98 -0
- libcom/os_insert/source/ldm/models/__init__.py +0 -0
- libcom/os_insert/source/ldm/models/autoencoder.py +443 -0
- libcom/os_insert/source/ldm/models/diffusion/__init__.py +0 -0
- libcom/os_insert/source/ldm/models/diffusion/classifier.py +267 -0
- libcom/os_insert/source/ldm/models/diffusion/ddim.py +263 -0
- libcom/os_insert/source/ldm/models/diffusion/ddpm.py +1417 -0
- libcom/os_insert/source/ldm/models/diffusion/plms.py +252 -0
- libcom/os_insert/source/ldm/modules/__init__.py +0 -0
- libcom/os_insert/source/ldm/modules/attention.py +261 -0
- libcom/os_insert/source/ldm/modules/diffusionmodules/__init__.py +0 -0
- libcom/os_insert/source/ldm/modules/diffusionmodules/model.py +835 -0
- libcom/os_insert/source/ldm/modules/diffusionmodules/openaimodel.py +1108 -0
- libcom/os_insert/source/ldm/modules/diffusionmodules/util.py +267 -0
- libcom/os_insert/source/ldm/modules/distributions/__init__.py +0 -0
- libcom/os_insert/source/ldm/modules/distributions/distributions.py +92 -0
- libcom/os_insert/source/ldm/modules/ema.py +76 -0
- libcom/os_insert/source/ldm/modules/encoders/__init__.py +0 -0
- libcom/os_insert/source/ldm/modules/encoders/modules.py +190 -0
- libcom/os_insert/source/ldm/modules/encoders/xf.py +130 -0
- libcom/os_insert/source/ldm/modules/losses/__init__.py +1 -0
- libcom/os_insert/source/ldm/modules/losses/contperceptual.py +111 -0
- libcom/os_insert/source/ldm/modules/losses/vqperceptual.py +167 -0
- libcom/os_insert/source/ldm/modules/x_transformer.py +641 -0
- libcom/os_insert/source/ldm/util.py +203 -0
- libcom/os_insert/source/objectstitch_data.py +66 -0
- libcom/os_insert/source/objectstitch_infer.py +618 -0
- libcom/os_insert/source/sam_on_objectstitch.py +253 -0
- libcom/os_insert/source/utils.py +41 -0
- libcom/os_insert/source/v1.yaml +108 -0
- libcom/painterly_image_harmonization/__init__.py +4 -0
- libcom/painterly_image_harmonization/painterly_image_harmonization.py +215 -0
- libcom/painterly_image_harmonization/source/PHDNet/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDNet/networks.py +202 -0
- libcom/painterly_image_harmonization/source/PHDNet/phdnet.py +47 -0
- libcom/painterly_image_harmonization/source/PHDNet/phdnet_config.py +18 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/inference_base.py +300 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/lr_scheduler.py +98 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/autoencoder.py +211 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/__init__.py +3 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/ddim.py +336 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/ddpm.py +1400 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/dpm_solver.py +1217 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/plms.py +243 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/scheduling_pndm.py +518 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/attention.py +378 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/model.py +857 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/openaimodel.py +823 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/util.py +270 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/distributions/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/distributions/distributions.py +92 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/ema.py +80 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/adapter.py +213 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/modules.py +82 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/__init__.py +1 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/api.py +269 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/api.py +175 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/base_model.py +16 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/blocks.py +342 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/dpt_depth.py +109 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/midas_net.py +76 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/midas_net_custom.py +128 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/transforms.py +234 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/vit.py +491 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/utils.py +189 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/model_edge.py +653 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/api.py +35 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/body.py +209 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/hand.py +75 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/model.py +178 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/util.py +203 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/utils.py +72 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/__init__.py +0 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/dual_encoder_fusion.py +93 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/__init__.py +2 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/bsrgan.py +730 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/bsrgan_light.py +651 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/utils_image.py +916 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/ldm/util.py +200 -0
- libcom/painterly_image_harmonization/source/PHDiffusion/stable_diffusion.yaml +75 -0
- libcom/painterly_image_harmonization/source/__init__.py +0 -0
- libcom/reflection_generation/__init__.py +3 -0
- libcom/reflection_generation/reflection_generation.py +174 -0
- libcom/reflection_generation/source/PostProcessModel.py +513 -0
- libcom/reflection_generation/source/__init__.py +0 -0
- libcom/reflection_generation/source/cldm/__init__.py +0 -0
- libcom/reflection_generation/source/cldm/base_network.py +166 -0
- libcom/reflection_generation/source/cldm/cldm.py +509 -0
- libcom/reflection_generation/source/cldm/ddim_hacked.py +317 -0
- libcom/reflection_generation/source/cldm/hack.py +111 -0
- libcom/reflection_generation/source/cldm/logger.py +127 -0
- libcom/reflection_generation/source/cldm/model.py +26 -0
- libcom/reflection_generation/source/cldm_v15.yaml +81 -0
- libcom/reflection_generation/source/cluster_feature.py +109 -0
- libcom/reflection_generation/source/ldm/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/models/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/models/autoencoder.py +219 -0
- libcom/reflection_generation/source/ldm/models/diffusion/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/models/diffusion/attention_processor.py +188 -0
- libcom/reflection_generation/source/ldm/models/diffusion/ddim.py +340 -0
- libcom/reflection_generation/source/ldm/models/diffusion/ddpm.py +1892 -0
- libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
- libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/dpm_solver.py +1154 -0
- libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
- libcom/reflection_generation/source/ldm/models/diffusion/ip_adapter.py +57 -0
- libcom/reflection_generation/source/ldm/models/diffusion/plms.py +244 -0
- libcom/reflection_generation/source/ldm/models/diffusion/sampling_util.py +22 -0
- libcom/reflection_generation/source/ldm/modules/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/modules/attention.py +341 -0
- libcom/reflection_generation/source/ldm/modules/diffusionmodules/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/modules/diffusionmodules/model.py +852 -0
- libcom/reflection_generation/source/ldm/modules/diffusionmodules/openaimodel.py +786 -0
- libcom/reflection_generation/source/ldm/modules/diffusionmodules/upscaling.py +81 -0
- libcom/reflection_generation/source/ldm/modules/diffusionmodules/util.py +280 -0
- libcom/reflection_generation/source/ldm/modules/distributions/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/modules/distributions/distributions.py +92 -0
- libcom/reflection_generation/source/ldm/modules/ema.py +80 -0
- libcom/reflection_generation/source/ldm/modules/encoders/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/modules/encoders/modules.py +214 -0
- libcom/reflection_generation/source/ldm/modules/image_degradation/__init__.py +0 -0
- libcom/reflection_generation/source/ldm/modules/image_degradation/bsrgan.py +730 -0
- libcom/reflection_generation/source/ldm/modules/image_degradation/bsrgan_light.py +651 -0
- libcom/reflection_generation/source/ldm/modules/image_degradation/utils_image.py +916 -0
- libcom/reflection_generation/source/ldm/util.py +198 -0
- libcom/reflection_generation/source/local_clip.py +749 -0
- libcom/reflection_generation/source/tool_add_control.py +49 -0
- libcom/shadow_generation/__init__.py +3 -0
- libcom/shadow_generation/shadow_generation.py +754 -0
- libcom/shadow_generation/source/PostProcessModel.py +511 -0
- libcom/shadow_generation/source/__init__.py +0 -0
- libcom/shadow_generation/source/attention_processor.py +234 -0
- libcom/shadow_generation/source/cldm/__init__.py +0 -0
- libcom/shadow_generation/source/cldm/base_network.py +166 -0
- libcom/shadow_generation/source/cldm/cldm.py +524 -0
- libcom/shadow_generation/source/cldm/model.py +28 -0
- libcom/shadow_generation/source/ldm/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/models/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/models/autoencoder.py +219 -0
- libcom/shadow_generation/source/ldm/models/diffusion/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/models/diffusion/attention_processor.py +188 -0
- libcom/shadow_generation/source/ldm/models/diffusion/ddim.py +340 -0
- libcom/shadow_generation/source/ldm/models/diffusion/ddpm.py +1893 -0
- libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
- libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/dpm_solver.py +1154 -0
- libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
- libcom/shadow_generation/source/ldm/models/diffusion/ip_adapter.py +57 -0
- libcom/shadow_generation/source/ldm/models/diffusion/plms.py +244 -0
- libcom/shadow_generation/source/ldm/models/diffusion/sampling_util.py +22 -0
- libcom/shadow_generation/source/ldm/modules/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/modules/attention.py +341 -0
- libcom/shadow_generation/source/ldm/modules/diffusionmodules/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/modules/diffusionmodules/model.py +852 -0
- libcom/shadow_generation/source/ldm/modules/diffusionmodules/openaimodel.py +786 -0
- libcom/shadow_generation/source/ldm/modules/diffusionmodules/upscaling.py +81 -0
- libcom/shadow_generation/source/ldm/modules/diffusionmodules/util.py +270 -0
- libcom/shadow_generation/source/ldm/modules/distributions/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/modules/distributions/distributions.py +92 -0
- libcom/shadow_generation/source/ldm/modules/ema.py +80 -0
- libcom/shadow_generation/source/ldm/modules/encoders/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/modules/encoders/modules.py +214 -0
- libcom/shadow_generation/source/ldm/modules/image_degradation/__init__.py +2 -0
- libcom/shadow_generation/source/ldm/modules/image_degradation/bsrgan.py +730 -0
- libcom/shadow_generation/source/ldm/modules/image_degradation/bsrgan_light.py +651 -0
- libcom/shadow_generation/source/ldm/modules/image_degradation/utils_image.py +916 -0
- libcom/shadow_generation/source/ldm/modules/midas/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/modules/midas/api.py +170 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/__init__.py +0 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/base_model.py +16 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/blocks.py +342 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/dpt_depth.py +109 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/midas_net.py +76 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/midas_net_custom.py +128 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/transforms.py +234 -0
- libcom/shadow_generation/source/ldm/modules/midas/midas/vit.py +491 -0
- libcom/shadow_generation/source/ldm/modules/midas/utils.py +189 -0
- libcom/shadow_generation/source/ldm/util.py +199 -0
- libcom/utils/__init__.py +0 -0
- libcom/utils/environment.py +12 -0
- libcom/utils/model_download.py +67 -0
- libcom/utils/process_image.py +108 -0
- libcom-0.1.8.dist-info/METADATA +157 -0
- libcom-0.1.8.dist-info/RECORD +310 -0
- libcom-0.1.8.dist-info/WHEEL +5 -0
- libcom-0.1.8.dist-info/licenses/LICENSE +201 -0
- libcom-0.1.8.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,916 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import math
|
|
3
|
+
import random
|
|
4
|
+
import numpy as np
|
|
5
|
+
import torch
|
|
6
|
+
import cv2
|
|
7
|
+
from torchvision.utils import make_grid
|
|
8
|
+
from datetime import datetime
|
|
9
|
+
#import matplotlib.pyplot as plt # TODO: check with Dominik, also bsrgan.py vs bsrgan_light.py
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
'''
|
|
16
|
+
# --------------------------------------------
|
|
17
|
+
# Kai Zhang (github: https://github.com/cszn)
|
|
18
|
+
# 03/Mar/2019
|
|
19
|
+
# --------------------------------------------
|
|
20
|
+
# https://github.com/twhui/SRGAN-pyTorch
|
|
21
|
+
# https://github.com/xinntao/BasicSR
|
|
22
|
+
# --------------------------------------------
|
|
23
|
+
'''
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif']
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def is_image_file(filename):
|
|
30
|
+
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def get_timestamp():
|
|
34
|
+
return datetime.now().strftime('%y%m%d-%H%M%S')
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def imshow(x, title=None, cbar=False, figsize=None):
|
|
38
|
+
plt.figure(figsize=figsize)
|
|
39
|
+
plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
|
|
40
|
+
if title:
|
|
41
|
+
plt.title(title)
|
|
42
|
+
if cbar:
|
|
43
|
+
plt.colorbar()
|
|
44
|
+
plt.show()
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def surf(Z, cmap='rainbow', figsize=None):
|
|
48
|
+
plt.figure(figsize=figsize)
|
|
49
|
+
ax3 = plt.axes(projection='3d')
|
|
50
|
+
|
|
51
|
+
w, h = Z.shape[:2]
|
|
52
|
+
xx = np.arange(0,w,1)
|
|
53
|
+
yy = np.arange(0,h,1)
|
|
54
|
+
X, Y = np.meshgrid(xx, yy)
|
|
55
|
+
ax3.plot_surface(X,Y,Z,cmap=cmap)
|
|
56
|
+
#ax3.contour(X,Y,Z, zdim='z',offset=-2,cmap=cmap)
|
|
57
|
+
plt.show()
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
'''
|
|
61
|
+
# --------------------------------------------
|
|
62
|
+
# get image pathes
|
|
63
|
+
# --------------------------------------------
|
|
64
|
+
'''
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def get_image_paths(dataroot):
|
|
68
|
+
paths = None # return None if dataroot is None
|
|
69
|
+
if dataroot is not None:
|
|
70
|
+
paths = sorted(_get_paths_from_images(dataroot))
|
|
71
|
+
return paths
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def _get_paths_from_images(path):
|
|
75
|
+
assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
|
|
76
|
+
images = []
|
|
77
|
+
for dirpath, _, fnames in sorted(os.walk(path)):
|
|
78
|
+
for fname in sorted(fnames):
|
|
79
|
+
if is_image_file(fname):
|
|
80
|
+
img_path = os.path.join(dirpath, fname)
|
|
81
|
+
images.append(img_path)
|
|
82
|
+
assert images, '{:s} has no valid image file'.format(path)
|
|
83
|
+
return images
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
'''
|
|
87
|
+
# --------------------------------------------
|
|
88
|
+
# split large images into small images
|
|
89
|
+
# --------------------------------------------
|
|
90
|
+
'''
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def patches_from_image(img, p_size=512, p_overlap=64, p_max=800):
|
|
94
|
+
w, h = img.shape[:2]
|
|
95
|
+
patches = []
|
|
96
|
+
if w > p_max and h > p_max:
|
|
97
|
+
w1 = list(np.arange(0, w-p_size, p_size-p_overlap, dtype=np.int))
|
|
98
|
+
h1 = list(np.arange(0, h-p_size, p_size-p_overlap, dtype=np.int))
|
|
99
|
+
w1.append(w-p_size)
|
|
100
|
+
h1.append(h-p_size)
|
|
101
|
+
# print(w1)
|
|
102
|
+
# print(h1)
|
|
103
|
+
for i in w1:
|
|
104
|
+
for j in h1:
|
|
105
|
+
patches.append(img[i:i+p_size, j:j+p_size,:])
|
|
106
|
+
else:
|
|
107
|
+
patches.append(img)
|
|
108
|
+
|
|
109
|
+
return patches
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def imssave(imgs, img_path):
|
|
113
|
+
"""
|
|
114
|
+
imgs: list, N images of size WxHxC
|
|
115
|
+
"""
|
|
116
|
+
img_name, ext = os.path.splitext(os.path.basename(img_path))
|
|
117
|
+
|
|
118
|
+
for i, img in enumerate(imgs):
|
|
119
|
+
if img.ndim == 3:
|
|
120
|
+
img = img[:, :, [2, 1, 0]]
|
|
121
|
+
new_path = os.path.join(os.path.dirname(img_path), img_name+str('_s{:04d}'.format(i))+'.png')
|
|
122
|
+
cv2.imwrite(new_path, img)
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=800, p_overlap=96, p_max=1000):
|
|
126
|
+
"""
|
|
127
|
+
split the large images from original_dataroot into small overlapped images with size (p_size)x(p_size),
|
|
128
|
+
and save them into taget_dataroot; only the images with larger size than (p_max)x(p_max)
|
|
129
|
+
will be splitted.
|
|
130
|
+
Args:
|
|
131
|
+
original_dataroot:
|
|
132
|
+
taget_dataroot:
|
|
133
|
+
p_size: size of small images
|
|
134
|
+
p_overlap: patch size in training is a good choice
|
|
135
|
+
p_max: images with smaller size than (p_max)x(p_max) keep unchanged.
|
|
136
|
+
"""
|
|
137
|
+
paths = get_image_paths(original_dataroot)
|
|
138
|
+
for img_path in paths:
|
|
139
|
+
# img_name, ext = os.path.splitext(os.path.basename(img_path))
|
|
140
|
+
img = imread_uint(img_path, n_channels=n_channels)
|
|
141
|
+
patches = patches_from_image(img, p_size, p_overlap, p_max)
|
|
142
|
+
imssave(patches, os.path.join(taget_dataroot,os.path.basename(img_path)))
|
|
143
|
+
#if original_dataroot == taget_dataroot:
|
|
144
|
+
#del img_path
|
|
145
|
+
|
|
146
|
+
'''
|
|
147
|
+
# --------------------------------------------
|
|
148
|
+
# makedir
|
|
149
|
+
# --------------------------------------------
|
|
150
|
+
'''
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def mkdir(path):
|
|
154
|
+
if not os.path.exists(path):
|
|
155
|
+
os.makedirs(path)
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def mkdirs(paths):
|
|
159
|
+
if isinstance(paths, str):
|
|
160
|
+
mkdir(paths)
|
|
161
|
+
else:
|
|
162
|
+
for path in paths:
|
|
163
|
+
mkdir(path)
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
def mkdir_and_rename(path):
|
|
167
|
+
if os.path.exists(path):
|
|
168
|
+
new_name = path + '_archived_' + get_timestamp()
|
|
169
|
+
print('Path already exists. Rename it to [{:s}]'.format(new_name))
|
|
170
|
+
os.rename(path, new_name)
|
|
171
|
+
os.makedirs(path)
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
'''
|
|
175
|
+
# --------------------------------------------
|
|
176
|
+
# read image from path
|
|
177
|
+
# opencv is fast, but read BGR numpy image
|
|
178
|
+
# --------------------------------------------
|
|
179
|
+
'''
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
# --------------------------------------------
|
|
183
|
+
# get uint8 image of size HxWxn_channles (RGB)
|
|
184
|
+
# --------------------------------------------
|
|
185
|
+
def imread_uint(path, n_channels=3):
|
|
186
|
+
# input: path
|
|
187
|
+
# output: HxWx3(RGB or GGG), or HxWx1 (G)
|
|
188
|
+
if n_channels == 1:
|
|
189
|
+
img = cv2.imread(path, 0) # cv2.IMREAD_GRAYSCALE
|
|
190
|
+
img = np.expand_dims(img, axis=2) # HxWx1
|
|
191
|
+
elif n_channels == 3:
|
|
192
|
+
img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGR or G
|
|
193
|
+
if img.ndim == 2:
|
|
194
|
+
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # GGG
|
|
195
|
+
else:
|
|
196
|
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB
|
|
197
|
+
return img
|
|
198
|
+
|
|
199
|
+
|
|
200
|
+
# --------------------------------------------
|
|
201
|
+
# matlab's imwrite
|
|
202
|
+
# --------------------------------------------
|
|
203
|
+
def imsave(img, img_path):
|
|
204
|
+
img = np.squeeze(img)
|
|
205
|
+
if img.ndim == 3:
|
|
206
|
+
img = img[:, :, [2, 1, 0]]
|
|
207
|
+
cv2.imwrite(img_path, img)
|
|
208
|
+
|
|
209
|
+
def imwrite(img, img_path):
|
|
210
|
+
img = np.squeeze(img)
|
|
211
|
+
if img.ndim == 3:
|
|
212
|
+
img = img[:, :, [2, 1, 0]]
|
|
213
|
+
cv2.imwrite(img_path, img)
|
|
214
|
+
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
# --------------------------------------------
|
|
218
|
+
# get single image of size HxWxn_channles (BGR)
|
|
219
|
+
# --------------------------------------------
|
|
220
|
+
def read_img(path):
|
|
221
|
+
# read image by cv2
|
|
222
|
+
# return: Numpy float32, HWC, BGR, [0,1]
|
|
223
|
+
img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # cv2.IMREAD_GRAYSCALE
|
|
224
|
+
img = img.astype(np.float32) / 255.
|
|
225
|
+
if img.ndim == 2:
|
|
226
|
+
img = np.expand_dims(img, axis=2)
|
|
227
|
+
# some images have 4 channels
|
|
228
|
+
if img.shape[2] > 3:
|
|
229
|
+
img = img[:, :, :3]
|
|
230
|
+
return img
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
'''
|
|
234
|
+
# --------------------------------------------
|
|
235
|
+
# image format conversion
|
|
236
|
+
# --------------------------------------------
|
|
237
|
+
# numpy(single) <---> numpy(unit)
|
|
238
|
+
# numpy(single) <---> tensor
|
|
239
|
+
# numpy(unit) <---> tensor
|
|
240
|
+
# --------------------------------------------
|
|
241
|
+
'''
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
# --------------------------------------------
|
|
245
|
+
# numpy(single) [0, 1] <---> numpy(unit)
|
|
246
|
+
# --------------------------------------------
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
def uint2single(img):
|
|
250
|
+
|
|
251
|
+
return np.float32(img/255.)
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
def single2uint(img):
|
|
255
|
+
|
|
256
|
+
return np.uint8((img.clip(0, 1)*255.).round())
|
|
257
|
+
|
|
258
|
+
|
|
259
|
+
def uint162single(img):
|
|
260
|
+
|
|
261
|
+
return np.float32(img/65535.)
|
|
262
|
+
|
|
263
|
+
|
|
264
|
+
def single2uint16(img):
|
|
265
|
+
|
|
266
|
+
return np.uint16((img.clip(0, 1)*65535.).round())
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
# --------------------------------------------
|
|
270
|
+
# numpy(unit) (HxWxC or HxW) <---> tensor
|
|
271
|
+
# --------------------------------------------
|
|
272
|
+
|
|
273
|
+
|
|
274
|
+
# convert uint to 4-dimensional torch tensor
|
|
275
|
+
def uint2tensor4(img):
|
|
276
|
+
if img.ndim == 2:
|
|
277
|
+
img = np.expand_dims(img, axis=2)
|
|
278
|
+
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0)
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
# convert uint to 3-dimensional torch tensor
|
|
282
|
+
def uint2tensor3(img):
|
|
283
|
+
if img.ndim == 2:
|
|
284
|
+
img = np.expand_dims(img, axis=2)
|
|
285
|
+
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.)
|
|
286
|
+
|
|
287
|
+
|
|
288
|
+
# convert 2/3/4-dimensional torch tensor to uint
|
|
289
|
+
def tensor2uint(img):
|
|
290
|
+
img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
|
|
291
|
+
if img.ndim == 3:
|
|
292
|
+
img = np.transpose(img, (1, 2, 0))
|
|
293
|
+
return np.uint8((img*255.0).round())
|
|
294
|
+
|
|
295
|
+
|
|
296
|
+
# --------------------------------------------
|
|
297
|
+
# numpy(single) (HxWxC) <---> tensor
|
|
298
|
+
# --------------------------------------------
|
|
299
|
+
|
|
300
|
+
|
|
301
|
+
# convert single (HxWxC) to 3-dimensional torch tensor
|
|
302
|
+
def single2tensor3(img):
|
|
303
|
+
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float()
|
|
304
|
+
|
|
305
|
+
|
|
306
|
+
# convert single (HxWxC) to 4-dimensional torch tensor
|
|
307
|
+
def single2tensor4(img):
|
|
308
|
+
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0)
|
|
309
|
+
|
|
310
|
+
|
|
311
|
+
# convert torch tensor to single
|
|
312
|
+
def tensor2single(img):
|
|
313
|
+
img = img.data.squeeze().float().cpu().numpy()
|
|
314
|
+
if img.ndim == 3:
|
|
315
|
+
img = np.transpose(img, (1, 2, 0))
|
|
316
|
+
|
|
317
|
+
return img
|
|
318
|
+
|
|
319
|
+
# convert torch tensor to single
|
|
320
|
+
def tensor2single3(img):
|
|
321
|
+
img = img.data.squeeze().float().cpu().numpy()
|
|
322
|
+
if img.ndim == 3:
|
|
323
|
+
img = np.transpose(img, (1, 2, 0))
|
|
324
|
+
elif img.ndim == 2:
|
|
325
|
+
img = np.expand_dims(img, axis=2)
|
|
326
|
+
return img
|
|
327
|
+
|
|
328
|
+
|
|
329
|
+
def single2tensor5(img):
|
|
330
|
+
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float().unsqueeze(0)
|
|
331
|
+
|
|
332
|
+
|
|
333
|
+
def single32tensor5(img):
|
|
334
|
+
return torch.from_numpy(np.ascontiguousarray(img)).float().unsqueeze(0).unsqueeze(0)
|
|
335
|
+
|
|
336
|
+
|
|
337
|
+
def single42tensor4(img):
|
|
338
|
+
return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float()
|
|
339
|
+
|
|
340
|
+
|
|
341
|
+
# from skimage.io import imread, imsave
|
|
342
|
+
def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
|
|
343
|
+
'''
|
|
344
|
+
Converts a torch Tensor into an image Numpy array of BGR channel order
|
|
345
|
+
Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
|
|
346
|
+
Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
|
|
347
|
+
'''
|
|
348
|
+
tensor = tensor.squeeze().float().cpu().clamp_(*min_max) # squeeze first, then clamp
|
|
349
|
+
tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) # to range [0,1]
|
|
350
|
+
n_dim = tensor.dim()
|
|
351
|
+
if n_dim == 4:
|
|
352
|
+
n_img = len(tensor)
|
|
353
|
+
img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
|
|
354
|
+
img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
|
|
355
|
+
elif n_dim == 3:
|
|
356
|
+
img_np = tensor.numpy()
|
|
357
|
+
img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
|
|
358
|
+
elif n_dim == 2:
|
|
359
|
+
img_np = tensor.numpy()
|
|
360
|
+
else:
|
|
361
|
+
raise TypeError(
|
|
362
|
+
'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
|
|
363
|
+
if out_type == np.uint8:
|
|
364
|
+
img_np = (img_np * 255.0).round()
|
|
365
|
+
# Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
|
|
366
|
+
return img_np.astype(out_type)
|
|
367
|
+
|
|
368
|
+
|
|
369
|
+
'''
|
|
370
|
+
# --------------------------------------------
|
|
371
|
+
# Augmentation, flipe and/or rotate
|
|
372
|
+
# --------------------------------------------
|
|
373
|
+
# The following two are enough.
|
|
374
|
+
# (1) augmet_img: numpy image of WxHxC or WxH
|
|
375
|
+
# (2) augment_img_tensor4: tensor image 1xCxWxH
|
|
376
|
+
# --------------------------------------------
|
|
377
|
+
'''
|
|
378
|
+
|
|
379
|
+
|
|
380
|
+
def augment_img(img, mode=0):
|
|
381
|
+
'''Kai Zhang (github: https://github.com/cszn)
|
|
382
|
+
'''
|
|
383
|
+
if mode == 0:
|
|
384
|
+
return img
|
|
385
|
+
elif mode == 1:
|
|
386
|
+
return np.flipud(np.rot90(img))
|
|
387
|
+
elif mode == 2:
|
|
388
|
+
return np.flipud(img)
|
|
389
|
+
elif mode == 3:
|
|
390
|
+
return np.rot90(img, k=3)
|
|
391
|
+
elif mode == 4:
|
|
392
|
+
return np.flipud(np.rot90(img, k=2))
|
|
393
|
+
elif mode == 5:
|
|
394
|
+
return np.rot90(img)
|
|
395
|
+
elif mode == 6:
|
|
396
|
+
return np.rot90(img, k=2)
|
|
397
|
+
elif mode == 7:
|
|
398
|
+
return np.flipud(np.rot90(img, k=3))
|
|
399
|
+
|
|
400
|
+
|
|
401
|
+
def augment_img_tensor4(img, mode=0):
|
|
402
|
+
'''Kai Zhang (github: https://github.com/cszn)
|
|
403
|
+
'''
|
|
404
|
+
if mode == 0:
|
|
405
|
+
return img
|
|
406
|
+
elif mode == 1:
|
|
407
|
+
return img.rot90(1, [2, 3]).flip([2])
|
|
408
|
+
elif mode == 2:
|
|
409
|
+
return img.flip([2])
|
|
410
|
+
elif mode == 3:
|
|
411
|
+
return img.rot90(3, [2, 3])
|
|
412
|
+
elif mode == 4:
|
|
413
|
+
return img.rot90(2, [2, 3]).flip([2])
|
|
414
|
+
elif mode == 5:
|
|
415
|
+
return img.rot90(1, [2, 3])
|
|
416
|
+
elif mode == 6:
|
|
417
|
+
return img.rot90(2, [2, 3])
|
|
418
|
+
elif mode == 7:
|
|
419
|
+
return img.rot90(3, [2, 3]).flip([2])
|
|
420
|
+
|
|
421
|
+
|
|
422
|
+
def augment_img_tensor(img, mode=0):
|
|
423
|
+
'''Kai Zhang (github: https://github.com/cszn)
|
|
424
|
+
'''
|
|
425
|
+
img_size = img.size()
|
|
426
|
+
img_np = img.data.cpu().numpy()
|
|
427
|
+
if len(img_size) == 3:
|
|
428
|
+
img_np = np.transpose(img_np, (1, 2, 0))
|
|
429
|
+
elif len(img_size) == 4:
|
|
430
|
+
img_np = np.transpose(img_np, (2, 3, 1, 0))
|
|
431
|
+
img_np = augment_img(img_np, mode=mode)
|
|
432
|
+
img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
|
|
433
|
+
if len(img_size) == 3:
|
|
434
|
+
img_tensor = img_tensor.permute(2, 0, 1)
|
|
435
|
+
elif len(img_size) == 4:
|
|
436
|
+
img_tensor = img_tensor.permute(3, 2, 0, 1)
|
|
437
|
+
|
|
438
|
+
return img_tensor.type_as(img)
|
|
439
|
+
|
|
440
|
+
|
|
441
|
+
def augment_img_np3(img, mode=0):
|
|
442
|
+
if mode == 0:
|
|
443
|
+
return img
|
|
444
|
+
elif mode == 1:
|
|
445
|
+
return img.transpose(1, 0, 2)
|
|
446
|
+
elif mode == 2:
|
|
447
|
+
return img[::-1, :, :]
|
|
448
|
+
elif mode == 3:
|
|
449
|
+
img = img[::-1, :, :]
|
|
450
|
+
img = img.transpose(1, 0, 2)
|
|
451
|
+
return img
|
|
452
|
+
elif mode == 4:
|
|
453
|
+
return img[:, ::-1, :]
|
|
454
|
+
elif mode == 5:
|
|
455
|
+
img = img[:, ::-1, :]
|
|
456
|
+
img = img.transpose(1, 0, 2)
|
|
457
|
+
return img
|
|
458
|
+
elif mode == 6:
|
|
459
|
+
img = img[:, ::-1, :]
|
|
460
|
+
img = img[::-1, :, :]
|
|
461
|
+
return img
|
|
462
|
+
elif mode == 7:
|
|
463
|
+
img = img[:, ::-1, :]
|
|
464
|
+
img = img[::-1, :, :]
|
|
465
|
+
img = img.transpose(1, 0, 2)
|
|
466
|
+
return img
|
|
467
|
+
|
|
468
|
+
|
|
469
|
+
def augment_imgs(img_list, hflip=True, rot=True):
|
|
470
|
+
# horizontal flip OR rotate
|
|
471
|
+
hflip = hflip and random.random() < 0.5
|
|
472
|
+
vflip = rot and random.random() < 0.5
|
|
473
|
+
rot90 = rot and random.random() < 0.5
|
|
474
|
+
|
|
475
|
+
def _augment(img):
|
|
476
|
+
if hflip:
|
|
477
|
+
img = img[:, ::-1, :]
|
|
478
|
+
if vflip:
|
|
479
|
+
img = img[::-1, :, :]
|
|
480
|
+
if rot90:
|
|
481
|
+
img = img.transpose(1, 0, 2)
|
|
482
|
+
return img
|
|
483
|
+
|
|
484
|
+
return [_augment(img) for img in img_list]
|
|
485
|
+
|
|
486
|
+
|
|
487
|
+
'''
|
|
488
|
+
# --------------------------------------------
|
|
489
|
+
# modcrop and shave
|
|
490
|
+
# --------------------------------------------
|
|
491
|
+
'''
|
|
492
|
+
|
|
493
|
+
|
|
494
|
+
def modcrop(img_in, scale):
|
|
495
|
+
# img_in: Numpy, HWC or HW
|
|
496
|
+
img = np.copy(img_in)
|
|
497
|
+
if img.ndim == 2:
|
|
498
|
+
H, W = img.shape
|
|
499
|
+
H_r, W_r = H % scale, W % scale
|
|
500
|
+
img = img[:H - H_r, :W - W_r]
|
|
501
|
+
elif img.ndim == 3:
|
|
502
|
+
H, W, C = img.shape
|
|
503
|
+
H_r, W_r = H % scale, W % scale
|
|
504
|
+
img = img[:H - H_r, :W - W_r, :]
|
|
505
|
+
else:
|
|
506
|
+
raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
|
|
507
|
+
return img
|
|
508
|
+
|
|
509
|
+
|
|
510
|
+
def shave(img_in, border=0):
|
|
511
|
+
# img_in: Numpy, HWC or HW
|
|
512
|
+
img = np.copy(img_in)
|
|
513
|
+
h, w = img.shape[:2]
|
|
514
|
+
img = img[border:h-border, border:w-border]
|
|
515
|
+
return img
|
|
516
|
+
|
|
517
|
+
|
|
518
|
+
'''
|
|
519
|
+
# --------------------------------------------
|
|
520
|
+
# image processing process on numpy image
|
|
521
|
+
# channel_convert(in_c, tar_type, img_list):
|
|
522
|
+
# rgb2ycbcr(img, only_y=True):
|
|
523
|
+
# bgr2ycbcr(img, only_y=True):
|
|
524
|
+
# ycbcr2rgb(img):
|
|
525
|
+
# --------------------------------------------
|
|
526
|
+
'''
|
|
527
|
+
|
|
528
|
+
|
|
529
|
+
def rgb2ycbcr(img, only_y=True):
|
|
530
|
+
'''same as matlab rgb2ycbcr
|
|
531
|
+
only_y: only return Y channel
|
|
532
|
+
Input:
|
|
533
|
+
uint8, [0, 255]
|
|
534
|
+
float, [0, 1]
|
|
535
|
+
'''
|
|
536
|
+
in_img_type = img.dtype
|
|
537
|
+
img.astype(np.float32)
|
|
538
|
+
if in_img_type != np.uint8:
|
|
539
|
+
img *= 255.
|
|
540
|
+
# convert
|
|
541
|
+
if only_y:
|
|
542
|
+
rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
|
|
543
|
+
else:
|
|
544
|
+
rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
|
|
545
|
+
[24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
|
|
546
|
+
if in_img_type == np.uint8:
|
|
547
|
+
rlt = rlt.round()
|
|
548
|
+
else:
|
|
549
|
+
rlt /= 255.
|
|
550
|
+
return rlt.astype(in_img_type)
|
|
551
|
+
|
|
552
|
+
|
|
553
|
+
def ycbcr2rgb(img):
|
|
554
|
+
'''same as matlab ycbcr2rgb
|
|
555
|
+
Input:
|
|
556
|
+
uint8, [0, 255]
|
|
557
|
+
float, [0, 1]
|
|
558
|
+
'''
|
|
559
|
+
in_img_type = img.dtype
|
|
560
|
+
img.astype(np.float32)
|
|
561
|
+
if in_img_type != np.uint8:
|
|
562
|
+
img *= 255.
|
|
563
|
+
# convert
|
|
564
|
+
rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
|
|
565
|
+
[0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
|
|
566
|
+
if in_img_type == np.uint8:
|
|
567
|
+
rlt = rlt.round()
|
|
568
|
+
else:
|
|
569
|
+
rlt /= 255.
|
|
570
|
+
return rlt.astype(in_img_type)
|
|
571
|
+
|
|
572
|
+
|
|
573
|
+
def bgr2ycbcr(img, only_y=True):
|
|
574
|
+
'''bgr version of rgb2ycbcr
|
|
575
|
+
only_y: only return Y channel
|
|
576
|
+
Input:
|
|
577
|
+
uint8, [0, 255]
|
|
578
|
+
float, [0, 1]
|
|
579
|
+
'''
|
|
580
|
+
in_img_type = img.dtype
|
|
581
|
+
img.astype(np.float32)
|
|
582
|
+
if in_img_type != np.uint8:
|
|
583
|
+
img *= 255.
|
|
584
|
+
# convert
|
|
585
|
+
if only_y:
|
|
586
|
+
rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
|
|
587
|
+
else:
|
|
588
|
+
rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
|
|
589
|
+
[65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
|
|
590
|
+
if in_img_type == np.uint8:
|
|
591
|
+
rlt = rlt.round()
|
|
592
|
+
else:
|
|
593
|
+
rlt /= 255.
|
|
594
|
+
return rlt.astype(in_img_type)
|
|
595
|
+
|
|
596
|
+
|
|
597
|
+
def channel_convert(in_c, tar_type, img_list):
|
|
598
|
+
# conversion among BGR, gray and y
|
|
599
|
+
if in_c == 3 and tar_type == 'gray': # BGR to gray
|
|
600
|
+
gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
|
|
601
|
+
return [np.expand_dims(img, axis=2) for img in gray_list]
|
|
602
|
+
elif in_c == 3 and tar_type == 'y': # BGR to y
|
|
603
|
+
y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
|
|
604
|
+
return [np.expand_dims(img, axis=2) for img in y_list]
|
|
605
|
+
elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR
|
|
606
|
+
return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
|
|
607
|
+
else:
|
|
608
|
+
return img_list
|
|
609
|
+
|
|
610
|
+
|
|
611
|
+
'''
|
|
612
|
+
# --------------------------------------------
|
|
613
|
+
# metric, PSNR and SSIM
|
|
614
|
+
# --------------------------------------------
|
|
615
|
+
'''
|
|
616
|
+
|
|
617
|
+
|
|
618
|
+
# --------------------------------------------
|
|
619
|
+
# PSNR
|
|
620
|
+
# --------------------------------------------
|
|
621
|
+
def calculate_psnr(img1, img2, border=0):
|
|
622
|
+
# img1 and img2 have range [0, 255]
|
|
623
|
+
#img1 = img1.squeeze()
|
|
624
|
+
#img2 = img2.squeeze()
|
|
625
|
+
if not img1.shape == img2.shape:
|
|
626
|
+
raise ValueError('Input images must have the same dimensions.')
|
|
627
|
+
h, w = img1.shape[:2]
|
|
628
|
+
img1 = img1[border:h-border, border:w-border]
|
|
629
|
+
img2 = img2[border:h-border, border:w-border]
|
|
630
|
+
|
|
631
|
+
img1 = img1.astype(np.float64)
|
|
632
|
+
img2 = img2.astype(np.float64)
|
|
633
|
+
mse = np.mean((img1 - img2)**2)
|
|
634
|
+
if mse == 0:
|
|
635
|
+
return float('inf')
|
|
636
|
+
return 20 * math.log10(255.0 / math.sqrt(mse))
|
|
637
|
+
|
|
638
|
+
|
|
639
|
+
# --------------------------------------------
|
|
640
|
+
# SSIM
|
|
641
|
+
# --------------------------------------------
|
|
642
|
+
def calculate_ssim(img1, img2, border=0):
|
|
643
|
+
'''calculate SSIM
|
|
644
|
+
the same outputs as MATLAB's
|
|
645
|
+
img1, img2: [0, 255]
|
|
646
|
+
'''
|
|
647
|
+
#img1 = img1.squeeze()
|
|
648
|
+
#img2 = img2.squeeze()
|
|
649
|
+
if not img1.shape == img2.shape:
|
|
650
|
+
raise ValueError('Input images must have the same dimensions.')
|
|
651
|
+
h, w = img1.shape[:2]
|
|
652
|
+
img1 = img1[border:h-border, border:w-border]
|
|
653
|
+
img2 = img2[border:h-border, border:w-border]
|
|
654
|
+
|
|
655
|
+
if img1.ndim == 2:
|
|
656
|
+
return ssim(img1, img2)
|
|
657
|
+
elif img1.ndim == 3:
|
|
658
|
+
if img1.shape[2] == 3:
|
|
659
|
+
ssims = []
|
|
660
|
+
for i in range(3):
|
|
661
|
+
ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
|
|
662
|
+
return np.array(ssims).mean()
|
|
663
|
+
elif img1.shape[2] == 1:
|
|
664
|
+
return ssim(np.squeeze(img1), np.squeeze(img2))
|
|
665
|
+
else:
|
|
666
|
+
raise ValueError('Wrong input image dimensions.')
|
|
667
|
+
|
|
668
|
+
|
|
669
|
+
def ssim(img1, img2):
|
|
670
|
+
C1 = (0.01 * 255)**2
|
|
671
|
+
C2 = (0.03 * 255)**2
|
|
672
|
+
|
|
673
|
+
img1 = img1.astype(np.float64)
|
|
674
|
+
img2 = img2.astype(np.float64)
|
|
675
|
+
kernel = cv2.getGaussianKernel(11, 1.5)
|
|
676
|
+
window = np.outer(kernel, kernel.transpose())
|
|
677
|
+
|
|
678
|
+
mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid
|
|
679
|
+
mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
|
|
680
|
+
mu1_sq = mu1**2
|
|
681
|
+
mu2_sq = mu2**2
|
|
682
|
+
mu1_mu2 = mu1 * mu2
|
|
683
|
+
sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
|
|
684
|
+
sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
|
|
685
|
+
sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
|
|
686
|
+
|
|
687
|
+
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
|
|
688
|
+
(sigma1_sq + sigma2_sq + C2))
|
|
689
|
+
return ssim_map.mean()
|
|
690
|
+
|
|
691
|
+
|
|
692
|
+
'''
|
|
693
|
+
# --------------------------------------------
|
|
694
|
+
# matlab's bicubic imresize (numpy and torch) [0, 1]
|
|
695
|
+
# --------------------------------------------
|
|
696
|
+
'''
|
|
697
|
+
|
|
698
|
+
|
|
699
|
+
# matlab 'imresize' function, now only support 'bicubic'
|
|
700
|
+
def cubic(x):
|
|
701
|
+
absx = torch.abs(x)
|
|
702
|
+
absx2 = absx**2
|
|
703
|
+
absx3 = absx**3
|
|
704
|
+
return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
|
|
705
|
+
(-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
|
|
706
|
+
|
|
707
|
+
|
|
708
|
+
def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
|
|
709
|
+
if (scale < 1) and (antialiasing):
|
|
710
|
+
# Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
|
|
711
|
+
kernel_width = kernel_width / scale
|
|
712
|
+
|
|
713
|
+
# Output-space coordinates
|
|
714
|
+
x = torch.linspace(1, out_length, out_length)
|
|
715
|
+
|
|
716
|
+
# Input-space coordinates. Calculate the inverse mapping such that 0.5
|
|
717
|
+
# in output space maps to 0.5 in input space, and 0.5+scale in output
|
|
718
|
+
# space maps to 1.5 in input space.
|
|
719
|
+
u = x / scale + 0.5 * (1 - 1 / scale)
|
|
720
|
+
|
|
721
|
+
# What is the left-most pixel that can be involved in the computation?
|
|
722
|
+
left = torch.floor(u - kernel_width / 2)
|
|
723
|
+
|
|
724
|
+
# What is the maximum number of pixels that can be involved in the
|
|
725
|
+
# computation? Note: it's OK to use an extra pixel here; if the
|
|
726
|
+
# corresponding weights are all zero, it will be eliminated at the end
|
|
727
|
+
# of this function.
|
|
728
|
+
P = math.ceil(kernel_width) + 2
|
|
729
|
+
|
|
730
|
+
# The indices of the input pixels involved in computing the k-th output
|
|
731
|
+
# pixel are in row k of the indices matrix.
|
|
732
|
+
indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
|
|
733
|
+
1, P).expand(out_length, P)
|
|
734
|
+
|
|
735
|
+
# The weights used to compute the k-th output pixel are in row k of the
|
|
736
|
+
# weights matrix.
|
|
737
|
+
distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
|
|
738
|
+
# apply cubic kernel
|
|
739
|
+
if (scale < 1) and (antialiasing):
|
|
740
|
+
weights = scale * cubic(distance_to_center * scale)
|
|
741
|
+
else:
|
|
742
|
+
weights = cubic(distance_to_center)
|
|
743
|
+
# Normalize the weights matrix so that each row sums to 1.
|
|
744
|
+
weights_sum = torch.sum(weights, 1).view(out_length, 1)
|
|
745
|
+
weights = weights / weights_sum.expand(out_length, P)
|
|
746
|
+
|
|
747
|
+
# If a column in weights is all zero, get rid of it. only consider the first and last column.
|
|
748
|
+
weights_zero_tmp = torch.sum((weights == 0), 0)
|
|
749
|
+
if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
|
|
750
|
+
indices = indices.narrow(1, 1, P - 2)
|
|
751
|
+
weights = weights.narrow(1, 1, P - 2)
|
|
752
|
+
if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
|
|
753
|
+
indices = indices.narrow(1, 0, P - 2)
|
|
754
|
+
weights = weights.narrow(1, 0, P - 2)
|
|
755
|
+
weights = weights.contiguous()
|
|
756
|
+
indices = indices.contiguous()
|
|
757
|
+
sym_len_s = -indices.min() + 1
|
|
758
|
+
sym_len_e = indices.max() - in_length
|
|
759
|
+
indices = indices + sym_len_s - 1
|
|
760
|
+
return weights, indices, int(sym_len_s), int(sym_len_e)
|
|
761
|
+
|
|
762
|
+
|
|
763
|
+
# --------------------------------------------
|
|
764
|
+
# imresize for tensor image [0, 1]
|
|
765
|
+
# --------------------------------------------
|
|
766
|
+
def imresize(img, scale, antialiasing=True):
|
|
767
|
+
# Now the scale should be the same for H and W
|
|
768
|
+
# input: img: pytorch tensor, CHW or HW [0,1]
|
|
769
|
+
# output: CHW or HW [0,1] w/o round
|
|
770
|
+
need_squeeze = True if img.dim() == 2 else False
|
|
771
|
+
if need_squeeze:
|
|
772
|
+
img.unsqueeze_(0)
|
|
773
|
+
in_C, in_H, in_W = img.size()
|
|
774
|
+
out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
|
|
775
|
+
kernel_width = 4
|
|
776
|
+
kernel = 'cubic'
|
|
777
|
+
|
|
778
|
+
# Return the desired dimension order for performing the resize. The
|
|
779
|
+
# strategy is to perform the resize first along the dimension with the
|
|
780
|
+
# smallest scale factor.
|
|
781
|
+
# Now we do not support this.
|
|
782
|
+
|
|
783
|
+
# get weights and indices
|
|
784
|
+
weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
|
|
785
|
+
in_H, out_H, scale, kernel, kernel_width, antialiasing)
|
|
786
|
+
weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
|
|
787
|
+
in_W, out_W, scale, kernel, kernel_width, antialiasing)
|
|
788
|
+
# process H dimension
|
|
789
|
+
# symmetric copying
|
|
790
|
+
img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
|
|
791
|
+
img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
|
|
792
|
+
|
|
793
|
+
sym_patch = img[:, :sym_len_Hs, :]
|
|
794
|
+
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
|
795
|
+
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
|
796
|
+
img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
|
|
797
|
+
|
|
798
|
+
sym_patch = img[:, -sym_len_He:, :]
|
|
799
|
+
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
|
800
|
+
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
|
801
|
+
img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
|
|
802
|
+
|
|
803
|
+
out_1 = torch.FloatTensor(in_C, out_H, in_W)
|
|
804
|
+
kernel_width = weights_H.size(1)
|
|
805
|
+
for i in range(out_H):
|
|
806
|
+
idx = int(indices_H[i][0])
|
|
807
|
+
for j in range(out_C):
|
|
808
|
+
out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
|
|
809
|
+
|
|
810
|
+
# process W dimension
|
|
811
|
+
# symmetric copying
|
|
812
|
+
out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
|
|
813
|
+
out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
|
|
814
|
+
|
|
815
|
+
sym_patch = out_1[:, :, :sym_len_Ws]
|
|
816
|
+
inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
|
|
817
|
+
sym_patch_inv = sym_patch.index_select(2, inv_idx)
|
|
818
|
+
out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
|
|
819
|
+
|
|
820
|
+
sym_patch = out_1[:, :, -sym_len_We:]
|
|
821
|
+
inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
|
|
822
|
+
sym_patch_inv = sym_patch.index_select(2, inv_idx)
|
|
823
|
+
out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
|
|
824
|
+
|
|
825
|
+
out_2 = torch.FloatTensor(in_C, out_H, out_W)
|
|
826
|
+
kernel_width = weights_W.size(1)
|
|
827
|
+
for i in range(out_W):
|
|
828
|
+
idx = int(indices_W[i][0])
|
|
829
|
+
for j in range(out_C):
|
|
830
|
+
out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i])
|
|
831
|
+
if need_squeeze:
|
|
832
|
+
out_2.squeeze_()
|
|
833
|
+
return out_2
|
|
834
|
+
|
|
835
|
+
|
|
836
|
+
# --------------------------------------------
|
|
837
|
+
# imresize for numpy image [0, 1]
|
|
838
|
+
# --------------------------------------------
|
|
839
|
+
def imresize_np(img, scale, antialiasing=True):
|
|
840
|
+
# Now the scale should be the same for H and W
|
|
841
|
+
# input: img: Numpy, HWC or HW [0,1]
|
|
842
|
+
# output: HWC or HW [0,1] w/o round
|
|
843
|
+
img = torch.from_numpy(img)
|
|
844
|
+
need_squeeze = True if img.dim() == 2 else False
|
|
845
|
+
if need_squeeze:
|
|
846
|
+
img.unsqueeze_(2)
|
|
847
|
+
|
|
848
|
+
in_H, in_W, in_C = img.size()
|
|
849
|
+
out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
|
|
850
|
+
kernel_width = 4
|
|
851
|
+
kernel = 'cubic'
|
|
852
|
+
|
|
853
|
+
# Return the desired dimension order for performing the resize. The
|
|
854
|
+
# strategy is to perform the resize first along the dimension with the
|
|
855
|
+
# smallest scale factor.
|
|
856
|
+
# Now we do not support this.
|
|
857
|
+
|
|
858
|
+
# get weights and indices
|
|
859
|
+
weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
|
|
860
|
+
in_H, out_H, scale, kernel, kernel_width, antialiasing)
|
|
861
|
+
weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
|
|
862
|
+
in_W, out_W, scale, kernel, kernel_width, antialiasing)
|
|
863
|
+
# process H dimension
|
|
864
|
+
# symmetric copying
|
|
865
|
+
img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
|
|
866
|
+
img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
|
|
867
|
+
|
|
868
|
+
sym_patch = img[:sym_len_Hs, :, :]
|
|
869
|
+
inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
|
|
870
|
+
sym_patch_inv = sym_patch.index_select(0, inv_idx)
|
|
871
|
+
img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
|
|
872
|
+
|
|
873
|
+
sym_patch = img[-sym_len_He:, :, :]
|
|
874
|
+
inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
|
|
875
|
+
sym_patch_inv = sym_patch.index_select(0, inv_idx)
|
|
876
|
+
img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
|
|
877
|
+
|
|
878
|
+
out_1 = torch.FloatTensor(out_H, in_W, in_C)
|
|
879
|
+
kernel_width = weights_H.size(1)
|
|
880
|
+
for i in range(out_H):
|
|
881
|
+
idx = int(indices_H[i][0])
|
|
882
|
+
for j in range(out_C):
|
|
883
|
+
out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i])
|
|
884
|
+
|
|
885
|
+
# process W dimension
|
|
886
|
+
# symmetric copying
|
|
887
|
+
out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
|
|
888
|
+
out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
|
|
889
|
+
|
|
890
|
+
sym_patch = out_1[:, :sym_len_Ws, :]
|
|
891
|
+
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
|
892
|
+
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
|
893
|
+
out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
|
|
894
|
+
|
|
895
|
+
sym_patch = out_1[:, -sym_len_We:, :]
|
|
896
|
+
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
|
897
|
+
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
|
898
|
+
out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
|
|
899
|
+
|
|
900
|
+
out_2 = torch.FloatTensor(out_H, out_W, in_C)
|
|
901
|
+
kernel_width = weights_W.size(1)
|
|
902
|
+
for i in range(out_W):
|
|
903
|
+
idx = int(indices_W[i][0])
|
|
904
|
+
for j in range(out_C):
|
|
905
|
+
out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i])
|
|
906
|
+
if need_squeeze:
|
|
907
|
+
out_2.squeeze_()
|
|
908
|
+
|
|
909
|
+
return out_2.numpy()
|
|
910
|
+
|
|
911
|
+
|
|
912
|
+
if __name__ == '__main__':
|
|
913
|
+
print('---')
|
|
914
|
+
# img = imread_uint('test.bmp', 3)
|
|
915
|
+
# img = uint2single(img)
|
|
916
|
+
# img_bicubic = imresize_np(img, 1/4)
|