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.
Files changed (310) hide show
  1. libcom/__init__.py +23 -0
  2. libcom/color_transfer/__init__.py +3 -0
  3. libcom/color_transfer/reinhard.py +57 -0
  4. libcom/fopa_heat_map/__init__.py +4 -0
  5. libcom/fopa_heat_map/fopa_heat_map.py +212 -0
  6. libcom/fopa_heat_map/source/__init__.py +0 -0
  7. libcom/fopa_heat_map/source/backbone/ResNet.py +229 -0
  8. libcom/fopa_heat_map/source/backbone/__init__.py +0 -0
  9. libcom/fopa_heat_map/source/data/OBdataset.py +292 -0
  10. libcom/fopa_heat_map/source/data/__init__.py +0 -0
  11. libcom/fopa_heat_map/source/data/all_transforms.py +25 -0
  12. libcom/fopa_heat_map/source/network/BaseBlocks.py +26 -0
  13. libcom/fopa_heat_map/source/network/DynamicModules.py +43 -0
  14. libcom/fopa_heat_map/source/network/ObPlaNet_simple.py +112 -0
  15. libcom/fopa_heat_map/source/network/__init__.py +1 -0
  16. libcom/fopa_heat_map/source/network/tensor_ops.py +37 -0
  17. libcom/fopa_heat_map/source/prepare_multi_fg_scales.py +85 -0
  18. libcom/fos_score/__init__.py +4 -0
  19. libcom/fos_score/fos_score_prediction.py +218 -0
  20. libcom/fos_score/source/__init__.py +0 -0
  21. libcom/fos_score/source/config/__init__.py +1 -0
  22. libcom/fos_score/source/config/config.py +68 -0
  23. libcom/fos_score/source/config/config.sfosd.yaml +43 -0
  24. libcom/fos_score/source/config/config_rfosd.yaml +43 -0
  25. libcom/fos_score/source/loss_function.py +118 -0
  26. libcom/fos_score/source/networks.py +355 -0
  27. libcom/harmony_score/__init__.py +3 -0
  28. libcom/harmony_score/harmony_score_prediction.py +119 -0
  29. libcom/harmony_score/source/__init__.py +0 -0
  30. libcom/harmony_score/source/bargainnet.py +226 -0
  31. libcom/image_harmonization/__init__.py +4 -0
  32. libcom/image_harmonization/image_harmonization.py +193 -0
  33. libcom/image_harmonization/source/__init__.py +0 -0
  34. libcom/image_harmonization/source/functions.py +747 -0
  35. libcom/image_harmonization/source/pct_net.py +85 -0
  36. libcom/image_harmonization/source/src/__init__.py +0 -0
  37. libcom/image_harmonization/source/src/lbm/__init__.py +0 -0
  38. libcom/image_harmonization/source/src/lbm/config.py +141 -0
  39. libcom/image_harmonization/source/src/lbm/data/__init__.py +62 -0
  40. libcom/image_harmonization/source/src/lbm/data/datasets/__init__.py +9 -0
  41. libcom/image_harmonization/source/src/lbm/data/datasets/collation_fn.py +41 -0
  42. libcom/image_harmonization/source/src/lbm/data/datasets/dataset.py +243 -0
  43. libcom/image_harmonization/source/src/lbm/data/datasets/datasets_config.py +42 -0
  44. libcom/image_harmonization/source/src/lbm/data/filters/__init__.py +12 -0
  45. libcom/image_harmonization/source/src/lbm/data/filters/base.py +21 -0
  46. libcom/image_harmonization/source/src/lbm/data/filters/filter_wrapper.py +36 -0
  47. libcom/image_harmonization/source/src/lbm/data/filters/filters.py +33 -0
  48. libcom/image_harmonization/source/src/lbm/data/filters/filters_config.py +32 -0
  49. libcom/image_harmonization/source/src/lbm/data/mappers/__init__.py +19 -0
  50. libcom/image_harmonization/source/src/lbm/data/mappers/base.py +26 -0
  51. libcom/image_harmonization/source/src/lbm/data/mappers/mappers.py +135 -0
  52. libcom/image_harmonization/source/src/lbm/data/mappers/mappers_config.py +109 -0
  53. libcom/image_harmonization/source/src/lbm/data/mappers/mappers_wrapper.py +31 -0
  54. libcom/image_harmonization/source/src/lbm/inference/__init__.py +4 -0
  55. libcom/image_harmonization/source/src/lbm/inference/inference.py +70 -0
  56. libcom/image_harmonization/source/src/lbm/inference/utils.py +222 -0
  57. libcom/image_harmonization/source/src/lbm/models/__init__.py +0 -0
  58. libcom/image_harmonization/source/src/lbm/models/base/__init__.py +4 -0
  59. libcom/image_harmonization/source/src/lbm/models/base/base_model.py +66 -0
  60. libcom/image_harmonization/source/src/lbm/models/base/model_config.py +8 -0
  61. libcom/image_harmonization/source/src/lbm/models/embedders/__init__.py +4 -0
  62. libcom/image_harmonization/source/src/lbm/models/embedders/base/__init__.py +4 -0
  63. libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner.py +60 -0
  64. libcom/image_harmonization/source/src/lbm/models/embedders/base/base_conditioner_config.py +27 -0
  65. libcom/image_harmonization/source/src/lbm/models/embedders/conditioners_wrapper.py +114 -0
  66. libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/__init__.py +4 -0
  67. libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_config.py +31 -0
  68. libcom/image_harmonization/source/src/lbm/models/embedders/latents_concat/latents_concat_embedder_model.py +80 -0
  69. libcom/image_harmonization/source/src/lbm/models/lbm/__init__.py +4 -0
  70. libcom/image_harmonization/source/src/lbm/models/lbm/lbm_config.py +101 -0
  71. libcom/image_harmonization/source/src/lbm/models/lbm/lbm_model.py +359 -0
  72. libcom/image_harmonization/source/src/lbm/models/unets/__init__.py +14 -0
  73. libcom/image_harmonization/source/src/lbm/models/unets/unet.py +148 -0
  74. libcom/image_harmonization/source/src/lbm/models/utils.py +377 -0
  75. libcom/image_harmonization/source/src/lbm/models/vae/__init__.py +4 -0
  76. libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL.py +136 -0
  77. libcom/image_harmonization/source/src/lbm/models/vae/autoencoderKL_config.py +27 -0
  78. libcom/image_harmonization/source/src/lbm/trainer/__init__.py +85 -0
  79. libcom/image_harmonization/source/src/lbm/trainer/loggers.py +324 -0
  80. libcom/image_harmonization/source/src/lbm/trainer/trainer.py +199 -0
  81. libcom/image_harmonization/source/src/lbm/trainer/training_config.py +82 -0
  82. libcom/image_harmonization/source/src/lbm/trainer/utils.py +193 -0
  83. libcom/inharmonious_region_localization/__init__.py +4 -0
  84. libcom/inharmonious_region_localization/inharmonious_region_localization.py +97 -0
  85. libcom/inharmonious_region_localization/source/__init__.py +0 -0
  86. libcom/inharmonious_region_localization/source/blocks.py +368 -0
  87. libcom/inharmonious_region_localization/source/functions.py +679 -0
  88. libcom/inharmonious_region_localization/source/madis_net.py +79 -0
  89. libcom/kontext_blending_harmonization/__init__.py +3 -0
  90. libcom/kontext_blending_harmonization/kontext_blending_harmonization.py +148 -0
  91. libcom/kontext_blending_harmonization/source/__init__.py +0 -0
  92. libcom/kontext_blending_harmonization/source/pipeline.py +1135 -0
  93. libcom/naive_composition/__init__.py +3 -0
  94. libcom/naive_composition/generate_composite_image.py +116 -0
  95. libcom/opa_score/__init__.py +3 -0
  96. libcom/opa_score/opa_score_prediction.py +101 -0
  97. libcom/opa_score/source/__init__.py +3 -0
  98. libcom/opa_score/source/object_place_net.py +55 -0
  99. libcom/opa_score/source/oblect_place_config.py +66 -0
  100. libcom/opa_score/source/resnet_4ch.py +213 -0
  101. libcom/os_insert/__init__.py +3 -0
  102. libcom/os_insert/diffusers_osinsert/__init__.py +78 -0
  103. libcom/os_insert/diffusers_osinsert/_patched_pipeline_flux_fill.py +878 -0
  104. libcom/os_insert/os_insert.py +534 -0
  105. libcom/os_insert/source/__init__.py +1 -0
  106. libcom/os_insert/source/ia_utils.py +183 -0
  107. libcom/os_insert/source/insertanything_infer.py +529 -0
  108. libcom/os_insert/source/ldm/__init__.py +0 -0
  109. libcom/os_insert/source/ldm/data/__init__.py +0 -0
  110. libcom/os_insert/source/ldm/data/base.py +23 -0
  111. libcom/os_insert/source/ldm/data/datamodule.py +128 -0
  112. libcom/os_insert/source/ldm/data/imagenet.py +394 -0
  113. libcom/os_insert/source/ldm/data/lsun.py +92 -0
  114. libcom/os_insert/source/ldm/data/open_images.py +728 -0
  115. libcom/os_insert/source/ldm/lr_scheduler.py +98 -0
  116. libcom/os_insert/source/ldm/models/__init__.py +0 -0
  117. libcom/os_insert/source/ldm/models/autoencoder.py +443 -0
  118. libcom/os_insert/source/ldm/models/diffusion/__init__.py +0 -0
  119. libcom/os_insert/source/ldm/models/diffusion/classifier.py +267 -0
  120. libcom/os_insert/source/ldm/models/diffusion/ddim.py +263 -0
  121. libcom/os_insert/source/ldm/models/diffusion/ddpm.py +1417 -0
  122. libcom/os_insert/source/ldm/models/diffusion/plms.py +252 -0
  123. libcom/os_insert/source/ldm/modules/__init__.py +0 -0
  124. libcom/os_insert/source/ldm/modules/attention.py +261 -0
  125. libcom/os_insert/source/ldm/modules/diffusionmodules/__init__.py +0 -0
  126. libcom/os_insert/source/ldm/modules/diffusionmodules/model.py +835 -0
  127. libcom/os_insert/source/ldm/modules/diffusionmodules/openaimodel.py +1108 -0
  128. libcom/os_insert/source/ldm/modules/diffusionmodules/util.py +267 -0
  129. libcom/os_insert/source/ldm/modules/distributions/__init__.py +0 -0
  130. libcom/os_insert/source/ldm/modules/distributions/distributions.py +92 -0
  131. libcom/os_insert/source/ldm/modules/ema.py +76 -0
  132. libcom/os_insert/source/ldm/modules/encoders/__init__.py +0 -0
  133. libcom/os_insert/source/ldm/modules/encoders/modules.py +190 -0
  134. libcom/os_insert/source/ldm/modules/encoders/xf.py +130 -0
  135. libcom/os_insert/source/ldm/modules/losses/__init__.py +1 -0
  136. libcom/os_insert/source/ldm/modules/losses/contperceptual.py +111 -0
  137. libcom/os_insert/source/ldm/modules/losses/vqperceptual.py +167 -0
  138. libcom/os_insert/source/ldm/modules/x_transformer.py +641 -0
  139. libcom/os_insert/source/ldm/util.py +203 -0
  140. libcom/os_insert/source/objectstitch_data.py +66 -0
  141. libcom/os_insert/source/objectstitch_infer.py +618 -0
  142. libcom/os_insert/source/sam_on_objectstitch.py +253 -0
  143. libcom/os_insert/source/utils.py +41 -0
  144. libcom/os_insert/source/v1.yaml +108 -0
  145. libcom/painterly_image_harmonization/__init__.py +4 -0
  146. libcom/painterly_image_harmonization/painterly_image_harmonization.py +215 -0
  147. libcom/painterly_image_harmonization/source/PHDNet/__init__.py +0 -0
  148. libcom/painterly_image_harmonization/source/PHDNet/networks.py +202 -0
  149. libcom/painterly_image_harmonization/source/PHDNet/phdnet.py +47 -0
  150. libcom/painterly_image_harmonization/source/PHDNet/phdnet_config.py +18 -0
  151. libcom/painterly_image_harmonization/source/PHDiffusion/__init__.py +0 -0
  152. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/__init__.py +0 -0
  153. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/inference_base.py +300 -0
  154. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/lr_scheduler.py +98 -0
  155. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/__init__.py +0 -0
  156. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/autoencoder.py +211 -0
  157. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/__init__.py +3 -0
  158. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/ddim.py +336 -0
  159. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/ddpm.py +1400 -0
  160. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  161. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/dpm_solver.py +1217 -0
  162. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  163. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/plms.py +243 -0
  164. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/models/diffusion/scheduling_pndm.py +518 -0
  165. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/__init__.py +0 -0
  166. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/attention.py +378 -0
  167. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/__init__.py +0 -0
  168. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/model.py +857 -0
  169. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/openaimodel.py +823 -0
  170. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/diffusionmodules/util.py +270 -0
  171. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/distributions/__init__.py +0 -0
  172. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/distributions/distributions.py +92 -0
  173. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/ema.py +80 -0
  174. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/__init__.py +0 -0
  175. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/adapter.py +213 -0
  176. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/encoders/modules.py +82 -0
  177. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/__init__.py +1 -0
  178. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/api.py +269 -0
  179. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/__init__.py +0 -0
  180. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/api.py +175 -0
  181. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/__init__.py +0 -0
  182. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/base_model.py +16 -0
  183. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/blocks.py +342 -0
  184. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/dpt_depth.py +109 -0
  185. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/midas_net.py +76 -0
  186. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/midas_net_custom.py +128 -0
  187. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/transforms.py +234 -0
  188. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/midas/vit.py +491 -0
  189. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/midas/utils.py +189 -0
  190. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/model_edge.py +653 -0
  191. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/__init__.py +0 -0
  192. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/api.py +35 -0
  193. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/body.py +209 -0
  194. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/hand.py +75 -0
  195. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/model.py +178 -0
  196. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/openpose/util.py +203 -0
  197. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/extra_condition/utils.py +72 -0
  198. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/__init__.py +0 -0
  199. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/fusion/dual_encoder_fusion.py +93 -0
  200. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/__init__.py +2 -0
  201. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/bsrgan.py +730 -0
  202. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  203. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/modules/image_degradation/utils_image.py +916 -0
  204. libcom/painterly_image_harmonization/source/PHDiffusion/ldm/util.py +200 -0
  205. libcom/painterly_image_harmonization/source/PHDiffusion/stable_diffusion.yaml +75 -0
  206. libcom/painterly_image_harmonization/source/__init__.py +0 -0
  207. libcom/reflection_generation/__init__.py +3 -0
  208. libcom/reflection_generation/reflection_generation.py +174 -0
  209. libcom/reflection_generation/source/PostProcessModel.py +513 -0
  210. libcom/reflection_generation/source/__init__.py +0 -0
  211. libcom/reflection_generation/source/cldm/__init__.py +0 -0
  212. libcom/reflection_generation/source/cldm/base_network.py +166 -0
  213. libcom/reflection_generation/source/cldm/cldm.py +509 -0
  214. libcom/reflection_generation/source/cldm/ddim_hacked.py +317 -0
  215. libcom/reflection_generation/source/cldm/hack.py +111 -0
  216. libcom/reflection_generation/source/cldm/logger.py +127 -0
  217. libcom/reflection_generation/source/cldm/model.py +26 -0
  218. libcom/reflection_generation/source/cldm_v15.yaml +81 -0
  219. libcom/reflection_generation/source/cluster_feature.py +109 -0
  220. libcom/reflection_generation/source/ldm/__init__.py +0 -0
  221. libcom/reflection_generation/source/ldm/models/__init__.py +0 -0
  222. libcom/reflection_generation/source/ldm/models/autoencoder.py +219 -0
  223. libcom/reflection_generation/source/ldm/models/diffusion/__init__.py +0 -0
  224. libcom/reflection_generation/source/ldm/models/diffusion/attention_processor.py +188 -0
  225. libcom/reflection_generation/source/ldm/models/diffusion/ddim.py +340 -0
  226. libcom/reflection_generation/source/ldm/models/diffusion/ddpm.py +1892 -0
  227. libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  228. libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/dpm_solver.py +1154 -0
  229. libcom/reflection_generation/source/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  230. libcom/reflection_generation/source/ldm/models/diffusion/ip_adapter.py +57 -0
  231. libcom/reflection_generation/source/ldm/models/diffusion/plms.py +244 -0
  232. libcom/reflection_generation/source/ldm/models/diffusion/sampling_util.py +22 -0
  233. libcom/reflection_generation/source/ldm/modules/__init__.py +0 -0
  234. libcom/reflection_generation/source/ldm/modules/attention.py +341 -0
  235. libcom/reflection_generation/source/ldm/modules/diffusionmodules/__init__.py +0 -0
  236. libcom/reflection_generation/source/ldm/modules/diffusionmodules/model.py +852 -0
  237. libcom/reflection_generation/source/ldm/modules/diffusionmodules/openaimodel.py +786 -0
  238. libcom/reflection_generation/source/ldm/modules/diffusionmodules/upscaling.py +81 -0
  239. libcom/reflection_generation/source/ldm/modules/diffusionmodules/util.py +280 -0
  240. libcom/reflection_generation/source/ldm/modules/distributions/__init__.py +0 -0
  241. libcom/reflection_generation/source/ldm/modules/distributions/distributions.py +92 -0
  242. libcom/reflection_generation/source/ldm/modules/ema.py +80 -0
  243. libcom/reflection_generation/source/ldm/modules/encoders/__init__.py +0 -0
  244. libcom/reflection_generation/source/ldm/modules/encoders/modules.py +214 -0
  245. libcom/reflection_generation/source/ldm/modules/image_degradation/__init__.py +0 -0
  246. libcom/reflection_generation/source/ldm/modules/image_degradation/bsrgan.py +730 -0
  247. libcom/reflection_generation/source/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  248. libcom/reflection_generation/source/ldm/modules/image_degradation/utils_image.py +916 -0
  249. libcom/reflection_generation/source/ldm/util.py +198 -0
  250. libcom/reflection_generation/source/local_clip.py +749 -0
  251. libcom/reflection_generation/source/tool_add_control.py +49 -0
  252. libcom/shadow_generation/__init__.py +3 -0
  253. libcom/shadow_generation/shadow_generation.py +754 -0
  254. libcom/shadow_generation/source/PostProcessModel.py +511 -0
  255. libcom/shadow_generation/source/__init__.py +0 -0
  256. libcom/shadow_generation/source/attention_processor.py +234 -0
  257. libcom/shadow_generation/source/cldm/__init__.py +0 -0
  258. libcom/shadow_generation/source/cldm/base_network.py +166 -0
  259. libcom/shadow_generation/source/cldm/cldm.py +524 -0
  260. libcom/shadow_generation/source/cldm/model.py +28 -0
  261. libcom/shadow_generation/source/ldm/__init__.py +0 -0
  262. libcom/shadow_generation/source/ldm/models/__init__.py +0 -0
  263. libcom/shadow_generation/source/ldm/models/autoencoder.py +219 -0
  264. libcom/shadow_generation/source/ldm/models/diffusion/__init__.py +0 -0
  265. libcom/shadow_generation/source/ldm/models/diffusion/attention_processor.py +188 -0
  266. libcom/shadow_generation/source/ldm/models/diffusion/ddim.py +340 -0
  267. libcom/shadow_generation/source/ldm/models/diffusion/ddpm.py +1893 -0
  268. libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/__init__.py +1 -0
  269. libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/dpm_solver.py +1154 -0
  270. libcom/shadow_generation/source/ldm/models/diffusion/dpm_solver/sampler.py +87 -0
  271. libcom/shadow_generation/source/ldm/models/diffusion/ip_adapter.py +57 -0
  272. libcom/shadow_generation/source/ldm/models/diffusion/plms.py +244 -0
  273. libcom/shadow_generation/source/ldm/models/diffusion/sampling_util.py +22 -0
  274. libcom/shadow_generation/source/ldm/modules/__init__.py +0 -0
  275. libcom/shadow_generation/source/ldm/modules/attention.py +341 -0
  276. libcom/shadow_generation/source/ldm/modules/diffusionmodules/__init__.py +0 -0
  277. libcom/shadow_generation/source/ldm/modules/diffusionmodules/model.py +852 -0
  278. libcom/shadow_generation/source/ldm/modules/diffusionmodules/openaimodel.py +786 -0
  279. libcom/shadow_generation/source/ldm/modules/diffusionmodules/upscaling.py +81 -0
  280. libcom/shadow_generation/source/ldm/modules/diffusionmodules/util.py +270 -0
  281. libcom/shadow_generation/source/ldm/modules/distributions/__init__.py +0 -0
  282. libcom/shadow_generation/source/ldm/modules/distributions/distributions.py +92 -0
  283. libcom/shadow_generation/source/ldm/modules/ema.py +80 -0
  284. libcom/shadow_generation/source/ldm/modules/encoders/__init__.py +0 -0
  285. libcom/shadow_generation/source/ldm/modules/encoders/modules.py +214 -0
  286. libcom/shadow_generation/source/ldm/modules/image_degradation/__init__.py +2 -0
  287. libcom/shadow_generation/source/ldm/modules/image_degradation/bsrgan.py +730 -0
  288. libcom/shadow_generation/source/ldm/modules/image_degradation/bsrgan_light.py +651 -0
  289. libcom/shadow_generation/source/ldm/modules/image_degradation/utils_image.py +916 -0
  290. libcom/shadow_generation/source/ldm/modules/midas/__init__.py +0 -0
  291. libcom/shadow_generation/source/ldm/modules/midas/api.py +170 -0
  292. libcom/shadow_generation/source/ldm/modules/midas/midas/__init__.py +0 -0
  293. libcom/shadow_generation/source/ldm/modules/midas/midas/base_model.py +16 -0
  294. libcom/shadow_generation/source/ldm/modules/midas/midas/blocks.py +342 -0
  295. libcom/shadow_generation/source/ldm/modules/midas/midas/dpt_depth.py +109 -0
  296. libcom/shadow_generation/source/ldm/modules/midas/midas/midas_net.py +76 -0
  297. libcom/shadow_generation/source/ldm/modules/midas/midas/midas_net_custom.py +128 -0
  298. libcom/shadow_generation/source/ldm/modules/midas/midas/transforms.py +234 -0
  299. libcom/shadow_generation/source/ldm/modules/midas/midas/vit.py +491 -0
  300. libcom/shadow_generation/source/ldm/modules/midas/utils.py +189 -0
  301. libcom/shadow_generation/source/ldm/util.py +199 -0
  302. libcom/utils/__init__.py +0 -0
  303. libcom/utils/environment.py +12 -0
  304. libcom/utils/model_download.py +67 -0
  305. libcom/utils/process_image.py +108 -0
  306. libcom-0.1.8.dist-info/METADATA +157 -0
  307. libcom-0.1.8.dist-info/RECORD +310 -0
  308. libcom-0.1.8.dist-info/WHEEL +5 -0
  309. libcom-0.1.8.dist-info/licenses/LICENSE +201 -0
  310. 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)