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
|
File without changes
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
from abc import abstractmethod
|
|
2
|
+
from torch.utils.data import Dataset, ConcatDataset, ChainDataset, IterableDataset
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class Txt2ImgIterableBaseDataset(IterableDataset):
|
|
6
|
+
'''
|
|
7
|
+
Define an interface to make the IterableDatasets for text2img data chainable
|
|
8
|
+
'''
|
|
9
|
+
def __init__(self, num_records=0, valid_ids=None, size=256):
|
|
10
|
+
super().__init__()
|
|
11
|
+
self.num_records = num_records
|
|
12
|
+
self.valid_ids = valid_ids
|
|
13
|
+
self.sample_ids = valid_ids
|
|
14
|
+
self.size = size
|
|
15
|
+
|
|
16
|
+
print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
|
|
17
|
+
|
|
18
|
+
def __len__(self):
|
|
19
|
+
return self.num_records
|
|
20
|
+
|
|
21
|
+
@abstractmethod
|
|
22
|
+
def __iter__(self):
|
|
23
|
+
pass
|
|
@@ -0,0 +1,128 @@
|
|
|
1
|
+
import argparse, os, sys, datetime, glob, importlib, csv
|
|
2
|
+
import numpy as np
|
|
3
|
+
import time
|
|
4
|
+
import torch
|
|
5
|
+
import torchvision
|
|
6
|
+
import pytorch_lightning as pl
|
|
7
|
+
|
|
8
|
+
from packaging import version
|
|
9
|
+
from omegaconf import OmegaConf
|
|
10
|
+
from torch.utils.data import random_split, DataLoader, Dataset, Subset
|
|
11
|
+
from functools import partial
|
|
12
|
+
from PIL import Image
|
|
13
|
+
|
|
14
|
+
from pytorch_lightning import seed_everything
|
|
15
|
+
from pytorch_lightning.trainer import Trainer
|
|
16
|
+
from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
|
|
17
|
+
from pytorch_lightning.utilities.rank_zero import rank_zero_only
|
|
18
|
+
from pytorch_lightning.utilities import rank_zero_info
|
|
19
|
+
|
|
20
|
+
from ldm.data.base import Txt2ImgIterableBaseDataset
|
|
21
|
+
from ldm.util import instantiate_from_config
|
|
22
|
+
import socket
|
|
23
|
+
from pytorch_lightning.plugins.environments import ClusterEnvironment,SLURMEnvironment
|
|
24
|
+
|
|
25
|
+
class WrappedDataset(Dataset):
|
|
26
|
+
"""Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""
|
|
27
|
+
|
|
28
|
+
def __init__(self, dataset):
|
|
29
|
+
self.data = dataset
|
|
30
|
+
|
|
31
|
+
def __len__(self):
|
|
32
|
+
return len(self.data)
|
|
33
|
+
|
|
34
|
+
def __getitem__(self, idx):
|
|
35
|
+
return self.data[idx]
|
|
36
|
+
|
|
37
|
+
def worker_init_fn(_):
|
|
38
|
+
worker_info = torch.utils.data.get_worker_info()
|
|
39
|
+
|
|
40
|
+
dataset = worker_info.dataset
|
|
41
|
+
worker_id = worker_info.id
|
|
42
|
+
|
|
43
|
+
if isinstance(dataset, Txt2ImgIterableBaseDataset):
|
|
44
|
+
split_size = dataset.num_records // worker_info.num_workers
|
|
45
|
+
# reset num_records to the true number to retain reliable length information
|
|
46
|
+
dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size]
|
|
47
|
+
current_id = np.random.choice(len(np.random.get_state()[1]), 1)
|
|
48
|
+
return np.random.seed(np.random.get_state()[1][current_id] + worker_id)
|
|
49
|
+
else:
|
|
50
|
+
return np.random.seed(np.random.get_state()[1][0] + worker_id)
|
|
51
|
+
|
|
52
|
+
class DataModuleFromConfig(pl.LightningDataModule):
|
|
53
|
+
def __init__(self, batch_size, num_workers=None, train=None, validation=None, test=None, predict=None,
|
|
54
|
+
wrap=False, shuffle_test_loader=False, use_worker_init_fn=False,
|
|
55
|
+
shuffle_val_dataloader=False):
|
|
56
|
+
super().__init__()
|
|
57
|
+
self.batch_size = batch_size
|
|
58
|
+
self.dataset_configs = dict()
|
|
59
|
+
self.num_workers = num_workers if num_workers is not None else batch_size
|
|
60
|
+
self.use_worker_init_fn = use_worker_init_fn
|
|
61
|
+
if train is not None:
|
|
62
|
+
self.dataset_configs["train"] = train
|
|
63
|
+
self.train_dataloader = self._train_dataloader
|
|
64
|
+
if validation is not None:
|
|
65
|
+
self.dataset_configs["validation"] = validation
|
|
66
|
+
self.val_dataloader = partial(self._val_dataloader, shuffle=shuffle_val_dataloader)
|
|
67
|
+
if test is not None:
|
|
68
|
+
self.dataset_configs["test"] = test
|
|
69
|
+
self.test_dataloader = partial(self._test_dataloader, shuffle=shuffle_test_loader)
|
|
70
|
+
if predict is not None:
|
|
71
|
+
self.dataset_configs["predict"] = predict
|
|
72
|
+
self.predict_dataloader = self._predict_dataloader
|
|
73
|
+
self.wrap = wrap
|
|
74
|
+
|
|
75
|
+
def prepare_data(self):
|
|
76
|
+
for data_cfg in self.dataset_configs.values():
|
|
77
|
+
instantiate_from_config(data_cfg)
|
|
78
|
+
|
|
79
|
+
def setup(self, stage=None):
|
|
80
|
+
self.datasets = dict(
|
|
81
|
+
(k, instantiate_from_config(self.dataset_configs[k]))
|
|
82
|
+
for k in self.dataset_configs)
|
|
83
|
+
if self.wrap:
|
|
84
|
+
for k in self.datasets:
|
|
85
|
+
self.datasets[k] = WrappedDataset(self.datasets[k])
|
|
86
|
+
|
|
87
|
+
def _train_dataloader(self):
|
|
88
|
+
is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)
|
|
89
|
+
if is_iterable_dataset or self.use_worker_init_fn:
|
|
90
|
+
init_fn = worker_init_fn
|
|
91
|
+
else:
|
|
92
|
+
init_fn = None
|
|
93
|
+
return DataLoader(self.datasets["train"], batch_size=self.batch_size,
|
|
94
|
+
num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True,
|
|
95
|
+
worker_init_fn=init_fn)
|
|
96
|
+
|
|
97
|
+
def _val_dataloader(self, shuffle=False):
|
|
98
|
+
if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:
|
|
99
|
+
init_fn = worker_init_fn
|
|
100
|
+
else:
|
|
101
|
+
init_fn = None
|
|
102
|
+
return DataLoader(self.datasets["validation"],
|
|
103
|
+
batch_size=self.batch_size,
|
|
104
|
+
num_workers=self.num_workers,
|
|
105
|
+
worker_init_fn=init_fn,
|
|
106
|
+
shuffle=shuffle)
|
|
107
|
+
|
|
108
|
+
def _test_dataloader(self, shuffle=False):
|
|
109
|
+
is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)
|
|
110
|
+
if is_iterable_dataset or self.use_worker_init_fn:
|
|
111
|
+
init_fn = worker_init_fn
|
|
112
|
+
else:
|
|
113
|
+
init_fn = None
|
|
114
|
+
|
|
115
|
+
# do not shuffle dataloader for iterable dataset
|
|
116
|
+
shuffle = shuffle and (not is_iterable_dataset)
|
|
117
|
+
|
|
118
|
+
return DataLoader(self.datasets["test"], batch_size=self.batch_size,
|
|
119
|
+
num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle)
|
|
120
|
+
|
|
121
|
+
def _predict_dataloader(self, shuffle=False):
|
|
122
|
+
if isinstance(self.datasets['predict'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:
|
|
123
|
+
init_fn = worker_init_fn
|
|
124
|
+
else:
|
|
125
|
+
init_fn = None
|
|
126
|
+
return DataLoader(self.datasets["predict"], batch_size=self.batch_size,
|
|
127
|
+
num_workers=self.num_workers, worker_init_fn=init_fn)
|
|
128
|
+
|
|
@@ -0,0 +1,394 @@
|
|
|
1
|
+
import os, yaml, pickle, shutil, tarfile, glob
|
|
2
|
+
import cv2
|
|
3
|
+
import albumentations
|
|
4
|
+
import PIL
|
|
5
|
+
import numpy as np
|
|
6
|
+
import torchvision.transforms.functional as TF
|
|
7
|
+
from omegaconf import OmegaConf
|
|
8
|
+
from functools import partial
|
|
9
|
+
from PIL import Image
|
|
10
|
+
from tqdm import tqdm
|
|
11
|
+
from torch.utils.data import Dataset, Subset
|
|
12
|
+
|
|
13
|
+
import taming.data.utils as tdu
|
|
14
|
+
from taming.data.imagenet import str_to_indices, give_synsets_from_indices, download, retrieve
|
|
15
|
+
from taming.data.imagenet import ImagePaths
|
|
16
|
+
|
|
17
|
+
from ldm.modules.image_degradation import degradation_fn_bsr, degradation_fn_bsr_light
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def synset2idx(path_to_yaml="data/index_synset.yaml"):
|
|
21
|
+
with open(path_to_yaml) as f:
|
|
22
|
+
di2s = yaml.load(f)
|
|
23
|
+
return dict((v,k) for k,v in di2s.items())
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class ImageNetBase(Dataset):
|
|
27
|
+
def __init__(self, config=None):
|
|
28
|
+
self.config = config or OmegaConf.create()
|
|
29
|
+
if not type(self.config)==dict:
|
|
30
|
+
self.config = OmegaConf.to_container(self.config)
|
|
31
|
+
self.keep_orig_class_label = self.config.get("keep_orig_class_label", False)
|
|
32
|
+
self.process_images = True # if False we skip loading & processing images and self.data contains filepaths
|
|
33
|
+
self._prepare()
|
|
34
|
+
self._prepare_synset_to_human()
|
|
35
|
+
self._prepare_idx_to_synset()
|
|
36
|
+
self._prepare_human_to_integer_label()
|
|
37
|
+
self._load()
|
|
38
|
+
|
|
39
|
+
def __len__(self):
|
|
40
|
+
return len(self.data)
|
|
41
|
+
|
|
42
|
+
def __getitem__(self, i):
|
|
43
|
+
return self.data[i]
|
|
44
|
+
|
|
45
|
+
def _prepare(self):
|
|
46
|
+
raise NotImplementedError()
|
|
47
|
+
|
|
48
|
+
def _filter_relpaths(self, relpaths):
|
|
49
|
+
ignore = set([
|
|
50
|
+
"n06596364_9591.JPEG",
|
|
51
|
+
])
|
|
52
|
+
relpaths = [rpath for rpath in relpaths if not rpath.split("/")[-1] in ignore]
|
|
53
|
+
if "sub_indices" in self.config:
|
|
54
|
+
indices = str_to_indices(self.config["sub_indices"])
|
|
55
|
+
synsets = give_synsets_from_indices(indices, path_to_yaml=self.idx2syn) # returns a list of strings
|
|
56
|
+
self.synset2idx = synset2idx(path_to_yaml=self.idx2syn)
|
|
57
|
+
files = []
|
|
58
|
+
for rpath in relpaths:
|
|
59
|
+
syn = rpath.split("/")[0]
|
|
60
|
+
if syn in synsets:
|
|
61
|
+
files.append(rpath)
|
|
62
|
+
return files
|
|
63
|
+
else:
|
|
64
|
+
return relpaths
|
|
65
|
+
|
|
66
|
+
def _prepare_synset_to_human(self):
|
|
67
|
+
SIZE = 2655750
|
|
68
|
+
URL = "https://heibox.uni-heidelberg.de/f/9f28e956cd304264bb82/?dl=1"
|
|
69
|
+
self.human_dict = os.path.join(self.root, "synset_human.txt")
|
|
70
|
+
if (not os.path.exists(self.human_dict) or
|
|
71
|
+
not os.path.getsize(self.human_dict)==SIZE):
|
|
72
|
+
download(URL, self.human_dict)
|
|
73
|
+
|
|
74
|
+
def _prepare_idx_to_synset(self):
|
|
75
|
+
URL = "https://heibox.uni-heidelberg.de/f/d835d5b6ceda4d3aa910/?dl=1"
|
|
76
|
+
self.idx2syn = os.path.join(self.root, "index_synset.yaml")
|
|
77
|
+
if (not os.path.exists(self.idx2syn)):
|
|
78
|
+
download(URL, self.idx2syn)
|
|
79
|
+
|
|
80
|
+
def _prepare_human_to_integer_label(self):
|
|
81
|
+
URL = "https://heibox.uni-heidelberg.de/f/2362b797d5be43b883f6/?dl=1"
|
|
82
|
+
self.human2integer = os.path.join(self.root, "imagenet1000_clsidx_to_labels.txt")
|
|
83
|
+
if (not os.path.exists(self.human2integer)):
|
|
84
|
+
download(URL, self.human2integer)
|
|
85
|
+
with open(self.human2integer, "r") as f:
|
|
86
|
+
lines = f.read().splitlines()
|
|
87
|
+
assert len(lines) == 1000
|
|
88
|
+
self.human2integer_dict = dict()
|
|
89
|
+
for line in lines:
|
|
90
|
+
value, key = line.split(":")
|
|
91
|
+
self.human2integer_dict[key] = int(value)
|
|
92
|
+
|
|
93
|
+
def _load(self):
|
|
94
|
+
with open(self.txt_filelist, "r") as f:
|
|
95
|
+
self.relpaths = f.read().splitlines()
|
|
96
|
+
l1 = len(self.relpaths)
|
|
97
|
+
self.relpaths = self._filter_relpaths(self.relpaths)
|
|
98
|
+
print("Removed {} files from filelist during filtering.".format(l1 - len(self.relpaths)))
|
|
99
|
+
|
|
100
|
+
self.synsets = [p.split("/")[0] for p in self.relpaths]
|
|
101
|
+
self.abspaths = [os.path.join(self.datadir, p) for p in self.relpaths]
|
|
102
|
+
|
|
103
|
+
unique_synsets = np.unique(self.synsets)
|
|
104
|
+
class_dict = dict((synset, i) for i, synset in enumerate(unique_synsets))
|
|
105
|
+
if not self.keep_orig_class_label:
|
|
106
|
+
self.class_labels = [class_dict[s] for s in self.synsets]
|
|
107
|
+
else:
|
|
108
|
+
self.class_labels = [self.synset2idx[s] for s in self.synsets]
|
|
109
|
+
|
|
110
|
+
with open(self.human_dict, "r") as f:
|
|
111
|
+
human_dict = f.read().splitlines()
|
|
112
|
+
human_dict = dict(line.split(maxsplit=1) for line in human_dict)
|
|
113
|
+
|
|
114
|
+
self.human_labels = [human_dict[s] for s in self.synsets]
|
|
115
|
+
|
|
116
|
+
labels = {
|
|
117
|
+
"relpath": np.array(self.relpaths),
|
|
118
|
+
"synsets": np.array(self.synsets),
|
|
119
|
+
"class_label": np.array(self.class_labels),
|
|
120
|
+
"human_label": np.array(self.human_labels),
|
|
121
|
+
}
|
|
122
|
+
|
|
123
|
+
if self.process_images:
|
|
124
|
+
self.size = retrieve(self.config, "size", default=256)
|
|
125
|
+
self.data = ImagePaths(self.abspaths,
|
|
126
|
+
labels=labels,
|
|
127
|
+
size=self.size,
|
|
128
|
+
random_crop=self.random_crop,
|
|
129
|
+
)
|
|
130
|
+
else:
|
|
131
|
+
self.data = self.abspaths
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
class ImageNetTrain(ImageNetBase):
|
|
135
|
+
NAME = "ILSVRC2012_train"
|
|
136
|
+
URL = "http://www.image-net.org/challenges/LSVRC/2012/"
|
|
137
|
+
AT_HASH = "a306397ccf9c2ead27155983c254227c0fd938e2"
|
|
138
|
+
FILES = [
|
|
139
|
+
"ILSVRC2012_img_train.tar",
|
|
140
|
+
]
|
|
141
|
+
SIZES = [
|
|
142
|
+
147897477120,
|
|
143
|
+
]
|
|
144
|
+
|
|
145
|
+
def __init__(self, process_images=True, data_root=None, **kwargs):
|
|
146
|
+
self.process_images = process_images
|
|
147
|
+
self.data_root = data_root
|
|
148
|
+
super().__init__(**kwargs)
|
|
149
|
+
|
|
150
|
+
def _prepare(self):
|
|
151
|
+
if self.data_root:
|
|
152
|
+
self.root = os.path.join(self.data_root, self.NAME)
|
|
153
|
+
else:
|
|
154
|
+
cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
|
155
|
+
self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
|
|
156
|
+
|
|
157
|
+
self.datadir = os.path.join(self.root, "data")
|
|
158
|
+
self.txt_filelist = os.path.join(self.root, "filelist.txt")
|
|
159
|
+
self.expected_length = 1281167
|
|
160
|
+
self.random_crop = retrieve(self.config, "ImageNetTrain/random_crop",
|
|
161
|
+
default=True)
|
|
162
|
+
if not tdu.is_prepared(self.root):
|
|
163
|
+
# prep
|
|
164
|
+
print("Preparing dataset {} in {}".format(self.NAME, self.root))
|
|
165
|
+
|
|
166
|
+
datadir = self.datadir
|
|
167
|
+
if not os.path.exists(datadir):
|
|
168
|
+
path = os.path.join(self.root, self.FILES[0])
|
|
169
|
+
if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
|
|
170
|
+
import academictorrents as at
|
|
171
|
+
atpath = at.get(self.AT_HASH, datastore=self.root)
|
|
172
|
+
assert atpath == path
|
|
173
|
+
|
|
174
|
+
print("Extracting {} to {}".format(path, datadir))
|
|
175
|
+
os.makedirs(datadir, exist_ok=True)
|
|
176
|
+
with tarfile.open(path, "r:") as tar:
|
|
177
|
+
tar.extractall(path=datadir)
|
|
178
|
+
|
|
179
|
+
print("Extracting sub-tars.")
|
|
180
|
+
subpaths = sorted(glob.glob(os.path.join(datadir, "*.tar")))
|
|
181
|
+
for subpath in tqdm(subpaths):
|
|
182
|
+
subdir = subpath[:-len(".tar")]
|
|
183
|
+
os.makedirs(subdir, exist_ok=True)
|
|
184
|
+
with tarfile.open(subpath, "r:") as tar:
|
|
185
|
+
tar.extractall(path=subdir)
|
|
186
|
+
|
|
187
|
+
filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
|
|
188
|
+
filelist = [os.path.relpath(p, start=datadir) for p in filelist]
|
|
189
|
+
filelist = sorted(filelist)
|
|
190
|
+
filelist = "\n".join(filelist)+"\n"
|
|
191
|
+
with open(self.txt_filelist, "w") as f:
|
|
192
|
+
f.write(filelist)
|
|
193
|
+
|
|
194
|
+
tdu.mark_prepared(self.root)
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
class ImageNetValidation(ImageNetBase):
|
|
198
|
+
NAME = "ILSVRC2012_validation"
|
|
199
|
+
URL = "http://www.image-net.org/challenges/LSVRC/2012/"
|
|
200
|
+
AT_HASH = "5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5"
|
|
201
|
+
VS_URL = "https://heibox.uni-heidelberg.de/f/3e0f6e9c624e45f2bd73/?dl=1"
|
|
202
|
+
FILES = [
|
|
203
|
+
"ILSVRC2012_img_val.tar",
|
|
204
|
+
"validation_synset.txt",
|
|
205
|
+
]
|
|
206
|
+
SIZES = [
|
|
207
|
+
6744924160,
|
|
208
|
+
1950000,
|
|
209
|
+
]
|
|
210
|
+
|
|
211
|
+
def __init__(self, process_images=True, data_root=None, **kwargs):
|
|
212
|
+
self.data_root = data_root
|
|
213
|
+
self.process_images = process_images
|
|
214
|
+
super().__init__(**kwargs)
|
|
215
|
+
|
|
216
|
+
def _prepare(self):
|
|
217
|
+
if self.data_root:
|
|
218
|
+
self.root = os.path.join(self.data_root, self.NAME)
|
|
219
|
+
else:
|
|
220
|
+
cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
|
221
|
+
self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
|
|
222
|
+
self.datadir = os.path.join(self.root, "data")
|
|
223
|
+
self.txt_filelist = os.path.join(self.root, "filelist.txt")
|
|
224
|
+
self.expected_length = 50000
|
|
225
|
+
self.random_crop = retrieve(self.config, "ImageNetValidation/random_crop",
|
|
226
|
+
default=False)
|
|
227
|
+
if not tdu.is_prepared(self.root):
|
|
228
|
+
# prep
|
|
229
|
+
print("Preparing dataset {} in {}".format(self.NAME, self.root))
|
|
230
|
+
|
|
231
|
+
datadir = self.datadir
|
|
232
|
+
if not os.path.exists(datadir):
|
|
233
|
+
path = os.path.join(self.root, self.FILES[0])
|
|
234
|
+
if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
|
|
235
|
+
import academictorrents as at
|
|
236
|
+
atpath = at.get(self.AT_HASH, datastore=self.root)
|
|
237
|
+
assert atpath == path
|
|
238
|
+
|
|
239
|
+
print("Extracting {} to {}".format(path, datadir))
|
|
240
|
+
os.makedirs(datadir, exist_ok=True)
|
|
241
|
+
with tarfile.open(path, "r:") as tar:
|
|
242
|
+
tar.extractall(path=datadir)
|
|
243
|
+
|
|
244
|
+
vspath = os.path.join(self.root, self.FILES[1])
|
|
245
|
+
if not os.path.exists(vspath) or not os.path.getsize(vspath)==self.SIZES[1]:
|
|
246
|
+
download(self.VS_URL, vspath)
|
|
247
|
+
|
|
248
|
+
with open(vspath, "r") as f:
|
|
249
|
+
synset_dict = f.read().splitlines()
|
|
250
|
+
synset_dict = dict(line.split() for line in synset_dict)
|
|
251
|
+
|
|
252
|
+
print("Reorganizing into synset folders")
|
|
253
|
+
synsets = np.unique(list(synset_dict.values()))
|
|
254
|
+
for s in synsets:
|
|
255
|
+
os.makedirs(os.path.join(datadir, s), exist_ok=True)
|
|
256
|
+
for k, v in synset_dict.items():
|
|
257
|
+
src = os.path.join(datadir, k)
|
|
258
|
+
dst = os.path.join(datadir, v)
|
|
259
|
+
shutil.move(src, dst)
|
|
260
|
+
|
|
261
|
+
filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
|
|
262
|
+
filelist = [os.path.relpath(p, start=datadir) for p in filelist]
|
|
263
|
+
filelist = sorted(filelist)
|
|
264
|
+
filelist = "\n".join(filelist)+"\n"
|
|
265
|
+
with open(self.txt_filelist, "w") as f:
|
|
266
|
+
f.write(filelist)
|
|
267
|
+
|
|
268
|
+
tdu.mark_prepared(self.root)
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
|
|
272
|
+
class ImageNetSR(Dataset):
|
|
273
|
+
def __init__(self, size=None,
|
|
274
|
+
degradation=None, downscale_f=4, min_crop_f=0.5, max_crop_f=1.,
|
|
275
|
+
random_crop=True):
|
|
276
|
+
"""
|
|
277
|
+
Imagenet Superresolution Dataloader
|
|
278
|
+
Performs following ops in order:
|
|
279
|
+
1. crops a crop of size s from image either as random or center crop
|
|
280
|
+
2. resizes crop to size with cv2.area_interpolation
|
|
281
|
+
3. degrades resized crop with degradation_fn
|
|
282
|
+
|
|
283
|
+
:param size: resizing to size after cropping
|
|
284
|
+
:param degradation: degradation_fn, e.g. cv_bicubic or bsrgan_light
|
|
285
|
+
:param downscale_f: Low Resolution Downsample factor
|
|
286
|
+
:param min_crop_f: determines crop size s,
|
|
287
|
+
where s = c * min_img_side_len with c sampled from interval (min_crop_f, max_crop_f)
|
|
288
|
+
:param max_crop_f: ""
|
|
289
|
+
:param data_root:
|
|
290
|
+
:param random_crop:
|
|
291
|
+
"""
|
|
292
|
+
self.base = self.get_base()
|
|
293
|
+
assert size
|
|
294
|
+
assert (size / downscale_f).is_integer()
|
|
295
|
+
self.size = size
|
|
296
|
+
self.LR_size = int(size / downscale_f)
|
|
297
|
+
self.min_crop_f = min_crop_f
|
|
298
|
+
self.max_crop_f = max_crop_f
|
|
299
|
+
assert(max_crop_f <= 1.)
|
|
300
|
+
self.center_crop = not random_crop
|
|
301
|
+
|
|
302
|
+
self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA)
|
|
303
|
+
|
|
304
|
+
self.pil_interpolation = False # gets reset later if incase interp_op is from pillow
|
|
305
|
+
|
|
306
|
+
if degradation == "bsrgan":
|
|
307
|
+
self.degradation_process = partial(degradation_fn_bsr, sf=downscale_f)
|
|
308
|
+
|
|
309
|
+
elif degradation == "bsrgan_light":
|
|
310
|
+
self.degradation_process = partial(degradation_fn_bsr_light, sf=downscale_f)
|
|
311
|
+
|
|
312
|
+
else:
|
|
313
|
+
interpolation_fn = {
|
|
314
|
+
"cv_nearest": cv2.INTER_NEAREST,
|
|
315
|
+
"cv_bilinear": cv2.INTER_LINEAR,
|
|
316
|
+
"cv_bicubic": cv2.INTER_CUBIC,
|
|
317
|
+
"cv_area": cv2.INTER_AREA,
|
|
318
|
+
"cv_lanczos": cv2.INTER_LANCZOS4,
|
|
319
|
+
"pil_nearest": PIL.Image.NEAREST,
|
|
320
|
+
"pil_bilinear": PIL.Image.BILINEAR,
|
|
321
|
+
"pil_bicubic": PIL.Image.BICUBIC,
|
|
322
|
+
"pil_box": PIL.Image.BOX,
|
|
323
|
+
"pil_hamming": PIL.Image.HAMMING,
|
|
324
|
+
"pil_lanczos": PIL.Image.LANCZOS,
|
|
325
|
+
}[degradation]
|
|
326
|
+
|
|
327
|
+
self.pil_interpolation = degradation.startswith("pil_")
|
|
328
|
+
|
|
329
|
+
if self.pil_interpolation:
|
|
330
|
+
self.degradation_process = partial(TF.resize, size=self.LR_size, interpolation=interpolation_fn)
|
|
331
|
+
|
|
332
|
+
else:
|
|
333
|
+
self.degradation_process = albumentations.SmallestMaxSize(max_size=self.LR_size,
|
|
334
|
+
interpolation=interpolation_fn)
|
|
335
|
+
|
|
336
|
+
def __len__(self):
|
|
337
|
+
return len(self.base)
|
|
338
|
+
|
|
339
|
+
def __getitem__(self, i):
|
|
340
|
+
example = self.base[i]
|
|
341
|
+
image = Image.open(example["file_path_"])
|
|
342
|
+
|
|
343
|
+
if not image.mode == "RGB":
|
|
344
|
+
image = image.convert("RGB")
|
|
345
|
+
|
|
346
|
+
image = np.array(image).astype(np.uint8)
|
|
347
|
+
|
|
348
|
+
min_side_len = min(image.shape[:2])
|
|
349
|
+
crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None)
|
|
350
|
+
crop_side_len = int(crop_side_len)
|
|
351
|
+
|
|
352
|
+
if self.center_crop:
|
|
353
|
+
self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len)
|
|
354
|
+
|
|
355
|
+
else:
|
|
356
|
+
self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len)
|
|
357
|
+
|
|
358
|
+
image = self.cropper(image=image)["image"]
|
|
359
|
+
image = self.image_rescaler(image=image)["image"]
|
|
360
|
+
|
|
361
|
+
if self.pil_interpolation:
|
|
362
|
+
image_pil = PIL.Image.fromarray(image)
|
|
363
|
+
LR_image = self.degradation_process(image_pil)
|
|
364
|
+
LR_image = np.array(LR_image).astype(np.uint8)
|
|
365
|
+
|
|
366
|
+
else:
|
|
367
|
+
LR_image = self.degradation_process(image=image)["image"]
|
|
368
|
+
|
|
369
|
+
example["image"] = (image/127.5 - 1.0).astype(np.float32)
|
|
370
|
+
example["LR_image"] = (LR_image/127.5 - 1.0).astype(np.float32)
|
|
371
|
+
|
|
372
|
+
return example
|
|
373
|
+
|
|
374
|
+
|
|
375
|
+
class ImageNetSRTrain(ImageNetSR):
|
|
376
|
+
def __init__(self, **kwargs):
|
|
377
|
+
super().__init__(**kwargs)
|
|
378
|
+
|
|
379
|
+
def get_base(self):
|
|
380
|
+
with open("data/imagenet_train_hr_indices.p", "rb") as f:
|
|
381
|
+
indices = pickle.load(f)
|
|
382
|
+
dset = ImageNetTrain(process_images=False,)
|
|
383
|
+
return Subset(dset, indices)
|
|
384
|
+
|
|
385
|
+
|
|
386
|
+
class ImageNetSRValidation(ImageNetSR):
|
|
387
|
+
def __init__(self, **kwargs):
|
|
388
|
+
super().__init__(**kwargs)
|
|
389
|
+
|
|
390
|
+
def get_base(self):
|
|
391
|
+
with open("data/imagenet_val_hr_indices.p", "rb") as f:
|
|
392
|
+
indices = pickle.load(f)
|
|
393
|
+
dset = ImageNetValidation(process_images=False,)
|
|
394
|
+
return Subset(dset, indices)
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import numpy as np
|
|
3
|
+
import PIL
|
|
4
|
+
from PIL import Image
|
|
5
|
+
from torch.utils.data import Dataset
|
|
6
|
+
from torchvision import transforms
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class LSUNBase(Dataset):
|
|
10
|
+
def __init__(self,
|
|
11
|
+
txt_file,
|
|
12
|
+
data_root,
|
|
13
|
+
size=None,
|
|
14
|
+
interpolation="bicubic",
|
|
15
|
+
flip_p=0.5
|
|
16
|
+
):
|
|
17
|
+
self.data_paths = txt_file
|
|
18
|
+
self.data_root = data_root
|
|
19
|
+
with open(self.data_paths, "r") as f:
|
|
20
|
+
self.image_paths = f.read().splitlines()
|
|
21
|
+
self._length = len(self.image_paths)
|
|
22
|
+
self.labels = {
|
|
23
|
+
"relative_file_path_": [l for l in self.image_paths],
|
|
24
|
+
"file_path_": [os.path.join(self.data_root, l)
|
|
25
|
+
for l in self.image_paths],
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
self.size = size
|
|
29
|
+
self.interpolation = {"linear": PIL.Image.LINEAR,
|
|
30
|
+
"bilinear": PIL.Image.BILINEAR,
|
|
31
|
+
"bicubic": PIL.Image.BICUBIC,
|
|
32
|
+
"lanczos": PIL.Image.LANCZOS,
|
|
33
|
+
}[interpolation]
|
|
34
|
+
self.flip = transforms.RandomHorizontalFlip(p=flip_p)
|
|
35
|
+
|
|
36
|
+
def __len__(self):
|
|
37
|
+
return self._length
|
|
38
|
+
|
|
39
|
+
def __getitem__(self, i):
|
|
40
|
+
example = dict((k, self.labels[k][i]) for k in self.labels)
|
|
41
|
+
image = Image.open(example["file_path_"])
|
|
42
|
+
if not image.mode == "RGB":
|
|
43
|
+
image = image.convert("RGB")
|
|
44
|
+
|
|
45
|
+
# default to score-sde preprocessing
|
|
46
|
+
img = np.array(image).astype(np.uint8)
|
|
47
|
+
crop = min(img.shape[0], img.shape[1])
|
|
48
|
+
h, w, = img.shape[0], img.shape[1]
|
|
49
|
+
img = img[(h - crop) // 2:(h + crop) // 2,
|
|
50
|
+
(w - crop) // 2:(w + crop) // 2]
|
|
51
|
+
|
|
52
|
+
image = Image.fromarray(img)
|
|
53
|
+
if self.size is not None:
|
|
54
|
+
image = image.resize((self.size, self.size), resample=self.interpolation)
|
|
55
|
+
|
|
56
|
+
image = self.flip(image)
|
|
57
|
+
image = np.array(image).astype(np.uint8)
|
|
58
|
+
example["image"] = (image / 127.5 - 1.0).astype(np.float32)
|
|
59
|
+
return example
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
class LSUNChurchesTrain(LSUNBase):
|
|
63
|
+
def __init__(self, **kwargs):
|
|
64
|
+
super().__init__(txt_file="data/lsun/church_outdoor_train.txt", data_root="data/lsun/churches", **kwargs)
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
class LSUNChurchesValidation(LSUNBase):
|
|
68
|
+
def __init__(self, flip_p=0., **kwargs):
|
|
69
|
+
super().__init__(txt_file="data/lsun/church_outdoor_val.txt", data_root="data/lsun/churches",
|
|
70
|
+
flip_p=flip_p, **kwargs)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class LSUNBedroomsTrain(LSUNBase):
|
|
74
|
+
def __init__(self, **kwargs):
|
|
75
|
+
super().__init__(txt_file="data/lsun/bedrooms_train.txt", data_root="data/lsun/bedrooms", **kwargs)
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
class LSUNBedroomsValidation(LSUNBase):
|
|
79
|
+
def __init__(self, flip_p=0.0, **kwargs):
|
|
80
|
+
super().__init__(txt_file="data/lsun/bedrooms_val.txt", data_root="data/lsun/bedrooms",
|
|
81
|
+
flip_p=flip_p, **kwargs)
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
class LSUNCatsTrain(LSUNBase):
|
|
85
|
+
def __init__(self, **kwargs):
|
|
86
|
+
super().__init__(txt_file="data/lsun/cat_train.txt", data_root="data/lsun/cats", **kwargs)
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
class LSUNCatsValidation(LSUNBase):
|
|
90
|
+
def __init__(self, flip_p=0., **kwargs):
|
|
91
|
+
super().__init__(txt_file="data/lsun/cat_val.txt", data_root="data/lsun/cats",
|
|
92
|
+
flip_p=flip_p, **kwargs)
|