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