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
@@ -0,0 +1,651 @@
1
+ # -*- coding: utf-8 -*-
2
+ import numpy as np
3
+ import cv2
4
+ import torch
5
+
6
+ from functools import partial
7
+ import random
8
+ from scipy import ndimage
9
+ import scipy
10
+ import scipy.stats as ss
11
+ from scipy.interpolate import interp2d
12
+ from scipy.linalg import orth
13
+ import albumentations
14
+
15
+ import utils_image as util
16
+
17
+ """
18
+ # --------------------------------------------
19
+ # Super-Resolution
20
+ # --------------------------------------------
21
+ #
22
+ # Kai Zhang (cskaizhang@gmail.com)
23
+ # https://github.com/cszn
24
+ # From 2019/03--2021/08
25
+ # --------------------------------------------
26
+ """
27
+
28
+ def modcrop_np(img, sf):
29
+ '''
30
+ Args:
31
+ img: numpy image, WxH or WxHxC
32
+ sf: scale factor
33
+ Return:
34
+ cropped image
35
+ '''
36
+ w, h = img.shape[:2]
37
+ im = np.copy(img)
38
+ return im[:w - w % sf, :h - h % sf, ...]
39
+
40
+
41
+ """
42
+ # --------------------------------------------
43
+ # anisotropic Gaussian kernels
44
+ # --------------------------------------------
45
+ """
46
+
47
+
48
+ def analytic_kernel(k):
49
+ """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
50
+ k_size = k.shape[0]
51
+ # Calculate the big kernels size
52
+ big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
53
+ # Loop over the small kernel to fill the big one
54
+ for r in range(k_size):
55
+ for c in range(k_size):
56
+ big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
57
+ # Crop the edges of the big kernel to ignore very small values and increase run time of SR
58
+ crop = k_size // 2
59
+ cropped_big_k = big_k[crop:-crop, crop:-crop]
60
+ # Normalize to 1
61
+ return cropped_big_k / cropped_big_k.sum()
62
+
63
+
64
+ def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
65
+ """ generate an anisotropic Gaussian kernel
66
+ Args:
67
+ ksize : e.g., 15, kernel size
68
+ theta : [0, pi], rotation angle range
69
+ l1 : [0.1,50], scaling of eigenvalues
70
+ l2 : [0.1,l1], scaling of eigenvalues
71
+ If l1 = l2, will get an isotropic Gaussian kernel.
72
+ Returns:
73
+ k : kernel
74
+ """
75
+
76
+ v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
77
+ V = np.array([[v[0], v[1]], [v[1], -v[0]]])
78
+ D = np.array([[l1, 0], [0, l2]])
79
+ Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
80
+ k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
81
+
82
+ return k
83
+
84
+
85
+ def gm_blur_kernel(mean, cov, size=15):
86
+ center = size / 2.0 + 0.5
87
+ k = np.zeros([size, size])
88
+ for y in range(size):
89
+ for x in range(size):
90
+ cy = y - center + 1
91
+ cx = x - center + 1
92
+ k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
93
+
94
+ k = k / np.sum(k)
95
+ return k
96
+
97
+
98
+ def shift_pixel(x, sf, upper_left=True):
99
+ """shift pixel for super-resolution with different scale factors
100
+ Args:
101
+ x: WxHxC or WxH
102
+ sf: scale factor
103
+ upper_left: shift direction
104
+ """
105
+ h, w = x.shape[:2]
106
+ shift = (sf - 1) * 0.5
107
+ xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
108
+ if upper_left:
109
+ x1 = xv + shift
110
+ y1 = yv + shift
111
+ else:
112
+ x1 = xv - shift
113
+ y1 = yv - shift
114
+
115
+ x1 = np.clip(x1, 0, w - 1)
116
+ y1 = np.clip(y1, 0, h - 1)
117
+
118
+ if x.ndim == 2:
119
+ x = interp2d(xv, yv, x)(x1, y1)
120
+ if x.ndim == 3:
121
+ for i in range(x.shape[-1]):
122
+ x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
123
+
124
+ return x
125
+
126
+
127
+ def blur(x, k):
128
+ '''
129
+ x: image, NxcxHxW
130
+ k: kernel, Nx1xhxw
131
+ '''
132
+ n, c = x.shape[:2]
133
+ p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
134
+ x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
135
+ k = k.repeat(1, c, 1, 1)
136
+ k = k.view(-1, 1, k.shape[2], k.shape[3])
137
+ x = x.view(1, -1, x.shape[2], x.shape[3])
138
+ x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c)
139
+ x = x.view(n, c, x.shape[2], x.shape[3])
140
+
141
+ return x
142
+
143
+
144
+ def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
145
+ """"
146
+ # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
147
+ # Kai Zhang
148
+ # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
149
+ # max_var = 2.5 * sf
150
+ """
151
+ # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
152
+ lambda_1 = min_var + np.random.rand() * (max_var - min_var)
153
+ lambda_2 = min_var + np.random.rand() * (max_var - min_var)
154
+ theta = np.random.rand() * np.pi # random theta
155
+ noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
156
+
157
+ # Set COV matrix using Lambdas and Theta
158
+ LAMBDA = np.diag([lambda_1, lambda_2])
159
+ Q = np.array([[np.cos(theta), -np.sin(theta)],
160
+ [np.sin(theta), np.cos(theta)]])
161
+ SIGMA = Q @ LAMBDA @ Q.T
162
+ INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
163
+
164
+ # Set expectation position (shifting kernel for aligned image)
165
+ MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
166
+ MU = MU[None, None, :, None]
167
+
168
+ # Create meshgrid for Gaussian
169
+ [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
170
+ Z = np.stack([X, Y], 2)[:, :, :, None]
171
+
172
+ # Calcualte Gaussian for every pixel of the kernel
173
+ ZZ = Z - MU
174
+ ZZ_t = ZZ.transpose(0, 1, 3, 2)
175
+ raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
176
+
177
+ # shift the kernel so it will be centered
178
+ # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
179
+
180
+ # Normalize the kernel and return
181
+ # kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
182
+ kernel = raw_kernel / np.sum(raw_kernel)
183
+ return kernel
184
+
185
+
186
+ def fspecial_gaussian(hsize, sigma):
187
+ hsize = [hsize, hsize]
188
+ siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
189
+ std = sigma
190
+ [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1))
191
+ arg = -(x * x + y * y) / (2 * std * std)
192
+ h = np.exp(arg)
193
+ h[h < scipy.finfo(float).eps * h.max()] = 0
194
+ sumh = h.sum()
195
+ if sumh != 0:
196
+ h = h / sumh
197
+ return h
198
+
199
+
200
+ def fspecial_laplacian(alpha):
201
+ alpha = max([0, min([alpha, 1])])
202
+ h1 = alpha / (alpha + 1)
203
+ h2 = (1 - alpha) / (alpha + 1)
204
+ h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
205
+ h = np.array(h)
206
+ return h
207
+
208
+
209
+ def fspecial(filter_type, *args, **kwargs):
210
+ '''
211
+ python code from:
212
+ https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
213
+ '''
214
+ if filter_type == 'gaussian':
215
+ return fspecial_gaussian(*args, **kwargs)
216
+ if filter_type == 'laplacian':
217
+ return fspecial_laplacian(*args, **kwargs)
218
+
219
+
220
+ """
221
+ # --------------------------------------------
222
+ # degradation models
223
+ # --------------------------------------------
224
+ """
225
+
226
+
227
+ def bicubic_degradation(x, sf=3):
228
+ '''
229
+ Args:
230
+ x: HxWxC image, [0, 1]
231
+ sf: down-scale factor
232
+ Return:
233
+ bicubicly downsampled LR image
234
+ '''
235
+ x = util.imresize_np(x, scale=1 / sf)
236
+ return x
237
+
238
+
239
+ def srmd_degradation(x, k, sf=3):
240
+ ''' blur + bicubic downsampling
241
+ Args:
242
+ x: HxWxC image, [0, 1]
243
+ k: hxw, double
244
+ sf: down-scale factor
245
+ Return:
246
+ downsampled LR image
247
+ Reference:
248
+ @inproceedings{zhang2018learning,
249
+ title={Learning a single convolutional super-resolution network for multiple degradations},
250
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
251
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
252
+ pages={3262--3271},
253
+ year={2018}
254
+ }
255
+ '''
256
+ x = ndimage.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror'
257
+ x = bicubic_degradation(x, sf=sf)
258
+ return x
259
+
260
+
261
+ def dpsr_degradation(x, k, sf=3):
262
+ ''' bicubic downsampling + blur
263
+ Args:
264
+ x: HxWxC image, [0, 1]
265
+ k: hxw, double
266
+ sf: down-scale factor
267
+ Return:
268
+ downsampled LR image
269
+ Reference:
270
+ @inproceedings{zhang2019deep,
271
+ title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
272
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
273
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
274
+ pages={1671--1681},
275
+ year={2019}
276
+ }
277
+ '''
278
+ x = bicubic_degradation(x, sf=sf)
279
+ x = ndimage.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
280
+ return x
281
+
282
+
283
+ def classical_degradation(x, k, sf=3):
284
+ ''' blur + downsampling
285
+ Args:
286
+ x: HxWxC image, [0, 1]/[0, 255]
287
+ k: hxw, double
288
+ sf: down-scale factor
289
+ Return:
290
+ downsampled LR image
291
+ '''
292
+ x = ndimage.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
293
+ # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
294
+ st = 0
295
+ return x[st::sf, st::sf, ...]
296
+
297
+
298
+ def add_sharpening(img, weight=0.5, radius=50, threshold=10):
299
+ """USM sharpening. borrowed from real-ESRGAN
300
+ Input image: I; Blurry image: B.
301
+ 1. K = I + weight * (I - B)
302
+ 2. Mask = 1 if abs(I - B) > threshold, else: 0
303
+ 3. Blur mask:
304
+ 4. Out = Mask * K + (1 - Mask) * I
305
+ Args:
306
+ img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
307
+ weight (float): Sharp weight. Default: 1.
308
+ radius (float): Kernel size of Gaussian blur. Default: 50.
309
+ threshold (int):
310
+ """
311
+ if radius % 2 == 0:
312
+ radius += 1
313
+ blur = cv2.GaussianBlur(img, (radius, radius), 0)
314
+ residual = img - blur
315
+ mask = np.abs(residual) * 255 > threshold
316
+ mask = mask.astype('float32')
317
+ soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
318
+
319
+ K = img + weight * residual
320
+ K = np.clip(K, 0, 1)
321
+ return soft_mask * K + (1 - soft_mask) * img
322
+
323
+
324
+ def add_blur(img, sf=4):
325
+ wd2 = 4.0 + sf
326
+ wd = 2.0 + 0.2 * sf
327
+
328
+ wd2 = wd2/4
329
+ wd = wd/4
330
+
331
+ if random.random() < 0.5:
332
+ l1 = wd2 * random.random()
333
+ l2 = wd2 * random.random()
334
+ k = anisotropic_Gaussian(ksize=random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2)
335
+ else:
336
+ k = fspecial('gaussian', random.randint(2, 4) + 3, wd * random.random())
337
+ img = ndimage.convolve(img, np.expand_dims(k, axis=2), mode='mirror')
338
+
339
+ return img
340
+
341
+
342
+ def add_resize(img, sf=4):
343
+ rnum = np.random.rand()
344
+ if rnum > 0.8: # up
345
+ sf1 = random.uniform(1, 2)
346
+ elif rnum < 0.7: # down
347
+ sf1 = random.uniform(0.5 / sf, 1)
348
+ else:
349
+ sf1 = 1.0
350
+ img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3]))
351
+ img = np.clip(img, 0.0, 1.0)
352
+
353
+ return img
354
+
355
+
356
+ # def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
357
+ # noise_level = random.randint(noise_level1, noise_level2)
358
+ # rnum = np.random.rand()
359
+ # if rnum > 0.6: # add color Gaussian noise
360
+ # img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
361
+ # elif rnum < 0.4: # add grayscale Gaussian noise
362
+ # img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
363
+ # else: # add noise
364
+ # L = noise_level2 / 255.
365
+ # D = np.diag(np.random.rand(3))
366
+ # U = orth(np.random.rand(3, 3))
367
+ # conv = np.dot(np.dot(np.transpose(U), D), U)
368
+ # img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
369
+ # img = np.clip(img, 0.0, 1.0)
370
+ # return img
371
+
372
+ def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
373
+ noise_level = random.randint(noise_level1, noise_level2)
374
+ rnum = np.random.rand()
375
+ if rnum > 0.6: # add color Gaussian noise
376
+ img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
377
+ elif rnum < 0.4: # add grayscale Gaussian noise
378
+ img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
379
+ else: # add noise
380
+ L = noise_level2 / 255.
381
+ D = np.diag(np.random.rand(3))
382
+ U = orth(np.random.rand(3, 3))
383
+ conv = np.dot(np.dot(np.transpose(U), D), U)
384
+ img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
385
+ img = np.clip(img, 0.0, 1.0)
386
+ return img
387
+
388
+
389
+ def add_speckle_noise(img, noise_level1=2, noise_level2=25):
390
+ noise_level = random.randint(noise_level1, noise_level2)
391
+ img = np.clip(img, 0.0, 1.0)
392
+ rnum = random.random()
393
+ if rnum > 0.6:
394
+ img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
395
+ elif rnum < 0.4:
396
+ img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
397
+ else:
398
+ L = noise_level2 / 255.
399
+ D = np.diag(np.random.rand(3))
400
+ U = orth(np.random.rand(3, 3))
401
+ conv = np.dot(np.dot(np.transpose(U), D), U)
402
+ img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
403
+ img = np.clip(img, 0.0, 1.0)
404
+ return img
405
+
406
+
407
+ def add_Poisson_noise(img):
408
+ img = np.clip((img * 255.0).round(), 0, 255) / 255.
409
+ vals = 10 ** (2 * random.random() + 2.0) # [2, 4]
410
+ if random.random() < 0.5:
411
+ img = np.random.poisson(img * vals).astype(np.float32) / vals
412
+ else:
413
+ img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
414
+ img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
415
+ noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray
416
+ img += noise_gray[:, :, np.newaxis]
417
+ img = np.clip(img, 0.0, 1.0)
418
+ return img
419
+
420
+
421
+ def add_JPEG_noise(img):
422
+ quality_factor = random.randint(80, 95)
423
+ img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
424
+ result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
425
+ img = cv2.imdecode(encimg, 1)
426
+ img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
427
+ return img
428
+
429
+
430
+ def random_crop(lq, hq, sf=4, lq_patchsize=64):
431
+ h, w = lq.shape[:2]
432
+ rnd_h = random.randint(0, h - lq_patchsize)
433
+ rnd_w = random.randint(0, w - lq_patchsize)
434
+ lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
435
+
436
+ rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
437
+ hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :]
438
+ return lq, hq
439
+
440
+
441
+ def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
442
+ """
443
+ This is the degradation model of BSRGAN from the paper
444
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
445
+ ----------
446
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
447
+ sf: scale factor
448
+ isp_model: camera ISP model
449
+ Returns
450
+ -------
451
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
452
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
453
+ """
454
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
455
+ sf_ori = sf
456
+
457
+ h1, w1 = img.shape[:2]
458
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
459
+ h, w = img.shape[:2]
460
+
461
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
462
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
463
+
464
+ hq = img.copy()
465
+
466
+ if sf == 4 and random.random() < scale2_prob: # downsample1
467
+ if np.random.rand() < 0.5:
468
+ img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])),
469
+ interpolation=random.choice([1, 2, 3]))
470
+ else:
471
+ img = util.imresize_np(img, 1 / 2, True)
472
+ img = np.clip(img, 0.0, 1.0)
473
+ sf = 2
474
+
475
+ shuffle_order = random.sample(range(7), 7)
476
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
477
+ if idx1 > idx2: # keep downsample3 last
478
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
479
+
480
+ for i in shuffle_order:
481
+
482
+ if i == 0:
483
+ img = add_blur(img, sf=sf)
484
+
485
+ elif i == 1:
486
+ img = add_blur(img, sf=sf)
487
+
488
+ elif i == 2:
489
+ a, b = img.shape[1], img.shape[0]
490
+ # downsample2
491
+ if random.random() < 0.75:
492
+ sf1 = random.uniform(1, 2 * sf)
493
+ img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])),
494
+ interpolation=random.choice([1, 2, 3]))
495
+ else:
496
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
497
+ k_shifted = shift_pixel(k, sf)
498
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
499
+ img = ndimage.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror')
500
+ img = img[0::sf, 0::sf, ...] # nearest downsampling
501
+ img = np.clip(img, 0.0, 1.0)
502
+
503
+ elif i == 3:
504
+ # downsample3
505
+ img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
506
+ img = np.clip(img, 0.0, 1.0)
507
+
508
+ elif i == 4:
509
+ # add Gaussian noise
510
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=8)
511
+
512
+ elif i == 5:
513
+ # add JPEG noise
514
+ if random.random() < jpeg_prob:
515
+ img = add_JPEG_noise(img)
516
+
517
+ elif i == 6:
518
+ # add processed camera sensor noise
519
+ if random.random() < isp_prob and isp_model is not None:
520
+ with torch.no_grad():
521
+ img, hq = isp_model.forward(img.copy(), hq)
522
+
523
+ # add final JPEG compression noise
524
+ img = add_JPEG_noise(img)
525
+
526
+ # random crop
527
+ img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
528
+
529
+ return img, hq
530
+
531
+
532
+ # todo no isp_model?
533
+ def degradation_bsrgan_variant(image, sf=4, isp_model=None, up=False):
534
+ """
535
+ This is the degradation model of BSRGAN from the paper
536
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
537
+ ----------
538
+ sf: scale factor
539
+ isp_model: camera ISP model
540
+ Returns
541
+ -------
542
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
543
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
544
+ """
545
+ image = util.uint2single(image)
546
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
547
+ sf_ori = sf
548
+
549
+ h1, w1 = image.shape[:2]
550
+ image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
551
+ h, w = image.shape[:2]
552
+
553
+ hq = image.copy()
554
+
555
+ if sf == 4 and random.random() < scale2_prob: # downsample1
556
+ if np.random.rand() < 0.5:
557
+ image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
558
+ interpolation=random.choice([1, 2, 3]))
559
+ else:
560
+ image = util.imresize_np(image, 1 / 2, True)
561
+ image = np.clip(image, 0.0, 1.0)
562
+ sf = 2
563
+
564
+ shuffle_order = random.sample(range(7), 7)
565
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
566
+ if idx1 > idx2: # keep downsample3 last
567
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
568
+
569
+ for i in shuffle_order:
570
+
571
+ if i == 0:
572
+ image = add_blur(image, sf=sf)
573
+
574
+ # elif i == 1:
575
+ # image = add_blur(image, sf=sf)
576
+
577
+ if i == 0:
578
+ pass
579
+
580
+ elif i == 2:
581
+ a, b = image.shape[1], image.shape[0]
582
+ # downsample2
583
+ if random.random() < 0.8:
584
+ sf1 = random.uniform(1, 2 * sf)
585
+ image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])),
586
+ interpolation=random.choice([1, 2, 3]))
587
+ else:
588
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
589
+ k_shifted = shift_pixel(k, sf)
590
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
591
+ image = ndimage.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror')
592
+ image = image[0::sf, 0::sf, ...] # nearest downsampling
593
+
594
+ image = np.clip(image, 0.0, 1.0)
595
+
596
+ elif i == 3:
597
+ # downsample3
598
+ image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
599
+ image = np.clip(image, 0.0, 1.0)
600
+
601
+ elif i == 4:
602
+ # add Gaussian noise
603
+ image = add_Gaussian_noise(image, noise_level1=1, noise_level2=2)
604
+
605
+ elif i == 5:
606
+ # add JPEG noise
607
+ if random.random() < jpeg_prob:
608
+ image = add_JPEG_noise(image)
609
+ #
610
+ # elif i == 6:
611
+ # # add processed camera sensor noise
612
+ # if random.random() < isp_prob and isp_model is not None:
613
+ # with torch.no_grad():
614
+ # img, hq = isp_model.forward(img.copy(), hq)
615
+
616
+ # add final JPEG compression noise
617
+ image = add_JPEG_noise(image)
618
+ image = util.single2uint(image)
619
+ if up:
620
+ image = cv2.resize(image, (w1, h1), interpolation=cv2.INTER_CUBIC) # todo: random, as above? want to condition on it then
621
+ example = {"image": image}
622
+ return example
623
+
624
+
625
+
626
+
627
+ if __name__ == '__main__':
628
+ print("hey")
629
+ img = util.imread_uint('utils/test.png', 3)
630
+ img = img[:448, :448]
631
+ h = img.shape[0] // 4
632
+ print("resizing to", h)
633
+ sf = 4
634
+ deg_fn = partial(degradation_bsrgan_variant, sf=sf)
635
+ for i in range(20):
636
+ print(i)
637
+ img_hq = img
638
+ img_lq = deg_fn(img)["image"]
639
+ img_hq, img_lq = util.uint2single(img_hq), util.uint2single(img_lq)
640
+ print(img_lq)
641
+ img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img_hq)["image"]
642
+ print(img_lq.shape)
643
+ print("bicubic", img_lq_bicubic.shape)
644
+ print(img_hq.shape)
645
+ lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
646
+ interpolation=0)
647
+ lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic),
648
+ (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
649
+ interpolation=0)
650
+ img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1)
651
+ util.imsave(img_concat, str(i) + '.png')