CSIGv3 AdcSR train scripts + A100 runbook

NOTES.md

@@ -0,0 +1,107 @@
+# 实施与 S0 冒烟要点（给下一执行者/服务器环境）
+
+## 已实现（本地已验证）
+- 工程脚手架 + 官方 clone(official/) + git 已提交(5a1afb7)
+- tools/*.py 全部 CPU 可跑: clean_images(阈值/去重/可复制到 clean) / tag_categories(7类启发式) /
+  build_manifest / split_val(70组隔离) / build_patches(>=20k patch) / download_batch(编排)
+  已用合成小图冒烟: 模糊/全黑/过曝/纯色正确剔除, 打标与 patch 构建正常。
+- src/*.py 与 scripts/*.py 语法编译通过(py_compile)。
+
+## 必须上 A100 用真实权重验证的点(S0, 1 迭代/配置)
+1. scripts/setup_env.sh + scripts/download_weights.sh(Plan A) 后运行 scripts/env_check.py。
+2. scripts/run_smoke.sh: train_lora 与 train_stage2 各 1 迭代; 重点核对:
+   - Net+halfDecoder 全链装配 forward 形状: LR[1,3,128,128] -> RGB[1,3,512,512]
+     (若 shape 报错, 以 official/test.py 官方装配为唯一基准修正 src/common.build_net/assemble_full_student)
+   - teacher z0 前向与判别器 dtype/shape 匹配(bf16 + autocast)
+   - loss 有限值且可保存 checkpoint、可续训
+3. 基线复现: 用官方 net_params_200.pkl + official/test.py 在 RealSR 复现论文量级指标;
+   并在 A100 记录 OSEDiff 512 fp16 时延(基准分母)。
+4. GDPO 教师: 先 python scripts/probe_gdpo.py <权重路径> 确认键格式; 若格式不兼容,
+   按 probe 输出在 src/common.load_gdpo_teacher 补映射, 否则回退 --teacher osediff(默认已支持)。
+
+## 执行顺序(每阶段 gate 见方案 §6)
+1. tools 数据构建 B1-B4 -> clean -> tag -> patches>=20k -> manifest -> split_val(70)
+2. pack_upload.ps1 上传代码+数据; 服务器下载权重
+3. S0 smoke -> S1(LoRA) -> eval_val 选 best -> S2(GDPO蒸馏) -> eval_val 选 best -> S3
+4. 每次选点: 用 src/eval_val.py 输出 proxy, 把最优权重软链/复制为
+   weight/s1/net_params_BEST.pkl (run_stage2 引用) / weight/s2/net_params_BEST.pkl (run_stage3 引用)
+5. 速度: src/latency_test.py; 出图: src/inference_4k.py; 导出: src/export_jit.py; 提交包检查: src/check_submission.py
+
+## 已知待办/说明
+- 训练脚本仅在服务器有 torch/权重时可运行; 本地 3050 4G 只做代码与数据开发。
+- run_stage1/2/3.sh 里的 net_params_BEST.pkl 为占位名, 按实际 eval 结果更新。
+- inference_4k/export_jit 采用“512窗->128->官方4x链->512”的 x1 包装语义;
+  若官方评测机 runner 口径不同, 以官方样例 runner 实测为准调整 export_jit.SR512。
+- 类别打标为启发式(文件名关键词+亮度/绿色占比), 抽样人工复核后再上传训练。
+
+## 试点批次进度 (B3 NKUSR8K 首 100 张, 2026-09-05)
+- raw 100 张(~1.9GB) -> clean 90 张(10 张 aHash 重复剔除; 8K 图质量均达标) -> 已删 raw 释放空间
+- 分类: city_view 41 / night_building 26 / plant 23 (skyscraper/shop_street/residential 为 0, 需 4KLSDB/RealSR 等源补充)
+- patches: 2998 张 512x512 (city_view 1366 / night_building 866 / plant 766), manifest: data/manifest_train.json
+- 人工复核图: data/review/NKUSR8K/*.jpg (每类 20 张拼图)
+- 注意: synthetic patch 的 val 划分须按"源图"分组(文件名前缀)而非单 patch, 防同源泄露; 最终 val 主要来自 real pairs。
+
+## 数据集构建进度 (2026-09-05 晚 更新, 本批次完成)
+- NKUSR8K 全量 1000 张: 下载 900(跳过试点100) -> 清洗 kept 639 (dup 259/过曝 2) -> 已删 raw(17.9GB 释放)
+- 合并 clean_all/NKUSR8K = 729 张(试点90 + 全量639, 4096 长边 jpg, 文件名前缀 NKUSR8K_)
+- 启发式分类: night_building 231 / plant 255 / city_view 243 (skyscraper/shop/residential 启发式分不出, 无碍训练)
+- 源级 val 划分: 70 源留作 val (split_val_sources.json), train 源 659
+- patches: data/patches 23,999 (city 7866/night 7684/plant 8449) + data/patches_val 560 (val 源)
+- manifest_train.json: synthetic 23,999 + RealSR 训练对 400; manifest_val.json: synthetic 560 + RealSR 测试对 99
+- RealSR 对清单: clean_pairs 配对命名 Canon_001_LR4.png <-> Canon_001_HR.png (build_manifest 已按 scene key 匹配)
+- 人工复核拼图: data/review/NKUSR8K_all/{city_view,night_building,plant}.jpg
+
+## 待办
+- [ ] DRealSR 训练集: 需人工(百度网盘 osiy / Google Drive 文件夹), 见 data/SOURCES.md; 到手后走 clean_pairs 入 real_pairs
+- [ ] 4KLSDB: 默认跳过(稀疏类抽取不划算, 结论见 SOURCES.md); 若需补量可用 tools 的 RangeReader 思路实现 extract_klsdb_rg.py
+- [ ] 打包上传: scripts/pack_upload.ps1 (需密码), 服务器 setup_env.sh + download_weights.sh + S0 smoke -> S1/S2/S3
+- [ ] 本地清理: 确认后删除 data/clean/NKUSR8K* / data/clean_all(已并入 patches, 仅留存备查)
+
+## 2026-09-05 深夜: CLIP/4KLSDB/稀疏类抓取 进展
+- 修正: NKUSR8K 无夜景(CLIP 确认 night_building=0); tag_categories 夜景仅允许文件名关键词, 不再按暗度推断
+- CLIP(openai/clip-vit-base-patch32, CPU, transformers5.x 需 .pooler_output) 打标 NKUSR8K 729:
+  skyscraper 204 / plant 354 / residential 63 / other_building 41 / storefront 22 / mall 21 / city_view 5, reject 19
+- 分类升级为 8 类: night_building / skyscraper / storefront(店铺门牌) / mall(商场外拍) / plant /
+  residential / city_view / other_building (tools: clip_filter.py, tag_categories.py)
+- 4KLSDB 抽取管线落地(全部本地可跑):
+  klsdb_index.py(206 分片 id->shard/rg 索引, ~27min) + klsdb_extract.py(rank 按 caption 评分排序 + extract 按 row-group 拉 hr)
+  第一批: 18 row-group, 11.6GB 流量 -> 1060 张 4K 原始 -> clean 802 张(data/clean/4KLSDB_b1)
+- 稀疏类原图 URL 抓取(fetch_klsdb_urls.py, 成功率 ~96%): night_building 575 / mall 355 /
+  storefront 501 / residential 495 / skyscraper 486 -> 后台逐类清洗中(data/clean/4KLSDB_url/<cls>)
+- DRealSR: 用户给的 GDrive 文件夹(1tP5m4k1...)只含 Test_x2/x3/x4.zip(测试集); 训练集在百度网盘 osiy(需登录)
+- SSDLite: 本地 Windows torchvision CPU wheel ABI 不兼容(2.7/2.6+cpu 均报 torchvision::nms), 已放弃本地;
+  方案: 本地用 CLIP; SSDLite person/物过滤代码留到 A100(Linux+torchvision 正常) 或后续换环境再启用
+- 待办: URL 清洗完成 -> CLIP 复核(去 people/indoor/other) -> 全池 pHash 去重 -> 合并 clean_all(带源前缀)
+  -> 8 类平衡报表 -> sample_review 拼图人工复核 -> build_patches(>=? 目标 5000-7000 HR 源后 50k+ patch)
+  -> 再跑 4KLSDB b2/b3 补量到总池 5000-7000
+
+## 2026-09-06 凌晨: LIU4K v2 引入 + 策略更新
+- 4KLSDB 政策: 建筑/风景/花草均可入池(增强鲁棒), 夜景适量即可(不堆量)
+- LIU4K v2 (GDrive 1FtVQtY2t_ecuy_gzJqZ-CatqrJBAdq_d): 4 组多段zip(Animal/Building/Street/Mountain),
+  只取 Building/Street/Mountain; Building = 400 张高分辨率 PNG(3992x2242~6208x4139, Pexels 建筑图, 10.2GB)
+  解压用 NVIDIA App 7-Zip 22.01(自带 7z.exe); 旧 7za 9.20 不支持 spanned zip
+- DRealSR: 放弃训练(用户确认); SSDLite: 取消(用户确认); 过滤以 CLIP 为主
+- URL 稀疏类清洗后: night_building 409 / storefront 446 / mall 248 / residential 339 / skyscraper 331 (共 1773)
+- 4KLSDB b1 CLIP 后 accept 604 (plant172/other_building171/residential124/city_view66/storefront35/skyscraper29/mall6/night1), reject 198
+- NKUSR8K CLIP accept 710 (skyscraper204/plant354/residential63/other_building41/storefront22/mall21/city_view5), reject 19
+- 新工具: merge_reports.py(多源合并/剔 reject/去重) + download_gdrive_parts.py(分卷断点下载)
+
+## 2026-09-06 凌晨2: 四源合并池(CLIP accept + URL文件夹标签 + reject剔除)
+- 池: logs/cat_pool_all.json, 共 3072 张清洗后 HR 原图
+  night_building 332 / skyscraper 575 / storefront 445 / mall 252 / plant 540 / residential 526 /
+  city_view 102 / other_building 300
+- 人工复核拼图: data/review/pool_all/{8类}.jpg (每类 20 张)
+- LIU4K Building: 400 PNG(10.2GB) -> clean 320 -> jpg(clean_all/LIU4K) -> CLIP accept 235(reject 85 室内/人像)
+- 剩余待补: city_view(102, 目标>=400), mall(252); 总池 3072 -> 目标 5000-7000
+- 下一批: 4KLSDB b2/b3(城景/商场/其他建筑加量) + LIU4K Street/Mountain(城市街道/山地风景)
+
+## 2026-09-06 凌晨: 收尾流水线启动(finalize_all.ps1 后台自动链)
+- 补量URL抓取: plant 790+383 / city_view 398+240 / other_building 554+259 / storefront 129 / mall 160 (raw)
+  (skyscraper/residential 候选枯竭未新增; fetch 去重由 jsonl 断点保证)
+- 4KLSDB b2 事故: klsdb_extract 未跳过已用 row-group -> b2 与 b1 完全重复(802/802), 已删 b2 目录;
+  已修复 extract 记录 used_rgs.json, 后续 b3 不再重复
+- clean_url_vol.ps1 正清洗 5 类 raw(plant/city_view/other_building/storefront/mall)
+- finalize_all.ps1(等清洗完): clip_url3(全URL目录) -> merge_reports(cat_pool_final, id去重)
+  -> split_pool(70源val) -> build_patches2(train 45k, val 560) -> build_manifest(train/val + RealSR对)
+- 产出: logs/cat_pool_final.json / cat_train.json / cat_val.json / data/patches / data/patches_val /
+  data/manifest_train.json / data/manifest_val.json / data/README.md

README.md

@@ -0,0 +1,36 @@
+# AdcSR 竞赛冲分工程（CSIG 2026 Camera学术之星 赛道二）
+
+以官方 Guaishou74851/AdcSR (CVPR 2025) 为基座的分阶段训练工程：
+S1 LoRA/QLoRA 像素域适配 -> S2 GDPO-SR 教师升级蒸馏 -> S3 域微调 -> 速度优化 -> 赛题提交包。
+
+## 目录
+- `official/` 官方仓库原样 clone（只读参照；训练脚本通过 sys.path 引用其 model/dataset/utils/ram）
+- `src/` 训练/蒸馏/导出/推理/评测
+- `tools/` 数据下载、清洗、打标、建 patch、清单、val 划分（纯 CPU）
+- `configs/` 各阶段 yaml
+- `scripts/` 服务器环境、权重、运行、打包上传、归档
+- `weight/ models/ logs/ data/` 运行产物（不入库）
+
+## 本地 vs 服务器
+- 本地（RTX 3050 4G / F 盘 ~44G）：代码开发、CPU 数据清洗、gt<=128 的 1 迭代冒烟。
+- aspire2a（单卡 A100 40G）：S1/S2/S3 全量训练、评测、导出、100 张 4K 出图。
+
+## 快速开始（服务器）
+```bash
+bash scripts/setup_env.sh              # conda + 依赖
+bash scripts/download_weights.sh       # Plan A: 服务器拉权重(AdcSR 5件套+SD2.1+GDPO)
+python scripts/env_check.py
+bash scripts/run_smoke.sh              # S0 冒烟（每个训练配置 1 迭代）
+bash scripts/run_stage1.sh             # S1 LoRA
+bash scripts/run_stage2.sh             # S2 GDPO 教师蒸馏
+bash scripts/run_stage3.sh             # S3 域微调
+python src/eval_val.py --weights <pkl> # val 评测 + 代理分
+python src/export_jit.py --weights <pkl> --out model_dir
+python src/inference_4k.py --lr_dir <testLR> --weights <pkl> --out output_dir
+python src/check_submission.py --zip <file.zip>
+```
+
+## 训练/评测口径（与官方一致）
+- 输入 LR 128x128 (scale=4, gt 512) -> 输出 512x512 RGB；官方 runner 用 512x512 包装/分块。
+- val 代理分 proxy = 0.25*CLIPIQA + 0.2*(MUSIQ/100) + 0.15*MANIQA + 0.15*(1-NIQE/10) + 0.25*(1-LPIPS)
+- 保真护栏：PSNR 不低于官方基线 0.5dB，SSIM 不低于 0.01。

RUNBOOK_A100.md

@@ -0,0 +1,51 @@
+# CSIGv3 AdcSR 训练脚本包 (A100 runbook)
+
+## 0) 数据
+```bash
+hf download XenderYang/CSIGv2_SR-Train --repo-type dataset --local-dir ./data
+cd data && for f in *.tar; do tar -xf "$f"; done && cd ..
+# 重新生成清单(manifest 内为 Windows 绝对路径, 不可直接用)
+python tools/build_manifest.py --hr_dirs data/patches \
+  --real_lr data/real_pairs/RealSR/train_lr --real_hr data/real_pairs/RealSR/train_hr \
+  --out data/manifest_train.json
+python tools/build_manifest.py --hr_dirs data/patches_val \
+  --real_lr data/real_pairs/RealSR/test_lr --real_hr data/real_pairs/RealSR/test_hr \
+  --out data/manifest_val.json
+```
+
+## 1) 环境
+```bash
+conda create -n AdcSR python=3.10 -y && conda activate AdcSR
+pip install -r requirements.txt
+# 数据仍以本仓库 scripts/setup_env.sh 指引; 权重:
+bash scripts/download_weights.sh    # 官方 AdcSR 5件 + SD2.1(fp16) + GDPO(可选)
+python scripts/env_check.py          # GPU/驱动/关键 import/单迭代自检
+```
+
+## 2) S0 冒烟(每个配置 2 迭代, 必须先过)
+```bash
+bash scripts/run_smoke.sh
+# 或手动:
+python src/train_lora.py   --config configs/config_smoke.yml --manifest data/manifest_train.json --real_prob 0.0 --steps 2 --batch_size 1 --grad_accum 1 --init_net weight/net_params_200.pkl --out weight/smoke --log_dir logs/smoke_lora
+python src/train_stage2.py --config configs/config_smoke.yml --steps 2 --batch_size 1 --grad_accum 1 --teacher osediff --skip_ram --init_net weight/net_params_200.pkl --out weight/smoke --log_dir logs/smoke_s2
+# GDPO 教师冒烟(先 probe)
+python scripts/probe_gdpo.py weight/gdpo
+python src/train_stage2.py --config configs/config_smoke.yml --steps 2 --batch_size 1 --teacher gdpo --gdpo_dir weight/gdpo --init_net weight/net_params_200.pkl --out weight/smoke --log_dir logs/smoke_s2_gdpo
+```
+
+## 3) 正式训练
+```bash
+bash scripts/run_stage1.sh   # S1 LoRA 域适配 (20k step)
+# eval 选 best -> weight/s1/net_params_BEST.pkl
+bash scripts/run_stage2.sh gdpo   # S2 GDPO 蒸馏(回退 osediff) (30k step)
+bash scripts/run_stage3.sh   # S3 域微调 (20k step)
+```
+
+## 4) 验证/导出/提交
+见 src/eval_val.py / export_jit.py / inference_4k.py / check_submission.py
+(赛题验证 3 对: data/val_official; 打分测试集: 见工程 README)
+
+## 注意
+- 训练脚本依赖 torchvision(官方 bsr/degradations.py import)。A100(Linux+torchvision)正常;
+  本机 Windows 沙箱若 torchvision 报 `operator torchvision::nms does not exist`, 请在自己终端(非沙箱)运行。
+- 数据清单中的真实退化对仅 RealSR 400/99; 其余母本在 data/clean_all, data/clean。

configs/config_base.yml

@@ -0,0 +1,35 @@
+# base degradation config (override dataroot_gt/gt_size/iter_num per stage)
+# Real-ESRGAN style degradation (shared by all stages)
+scale: 4
+resize_prob: [0.2, 0.7, 0.1]
+resize_range: [0.3, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2.0]
+gray_noise_prob: 0.4
+jpeg_range: [60, 95]
+second_blur_prob: 0.5
+resize_prob2: [0.3, 0.4, 0.3]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 12]
+poisson_scale_range2: [0.05, 1.0]
+gray_noise_prob2: 0.4
+jpeg_range2: [60, 100]
+blur_kernel_size: 21
+kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob: 0.1
+blur_sigma: [0.2, 1.5]
+betag_range: [0.5, 2.0]
+betap_range: [1, 1.5]
+blur_kernel_size2: 11
+kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob2: 0.1
+blur_sigma2: [0.2, 1.0]
+betag_range2: [0.5, 2.0]
+betap_range2: [1, 1.5]
+final_sinc_prob: 0.8
+use_hflip: True
+use_rot: False

configs/config_s1_lora.yml

@@ -0,0 +1,37 @@
+dataroot_gt: data/patches
+gt_size: 512
+iter_num: 1000
+# Real-ESRGAN style degradation (shared by all stages)
+scale: 4
+resize_prob: [0.2, 0.7, 0.1]
+resize_range: [0.3, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2.0]
+gray_noise_prob: 0.4
+jpeg_range: [60, 95]
+second_blur_prob: 0.5
+resize_prob2: [0.3, 0.4, 0.3]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 12]
+poisson_scale_range2: [0.05, 1.0]
+gray_noise_prob2: 0.4
+jpeg_range2: [60, 100]
+blur_kernel_size: 21
+kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob: 0.1
+blur_sigma: [0.2, 1.5]
+betag_range: [0.5, 2.0]
+betap_range: [1, 1.5]
+blur_kernel_size2: 11
+kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob2: 0.1
+blur_sigma2: [0.2, 1.0]
+betag_range2: [0.5, 2.0]
+betap_range2: [1, 1.5]
+final_sinc_prob: 0.8
+use_hflip: True
+use_rot: False

configs/config_s2_distill.yml

@@ -0,0 +1,37 @@
+dataroot_gt: data/patches
+gt_size: 512
+iter_num: 1000
+# Real-ESRGAN style degradation (shared by all stages)
+scale: 4
+resize_prob: [0.2, 0.7, 0.1]
+resize_range: [0.3, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2.0]
+gray_noise_prob: 0.4
+jpeg_range: [60, 95]
+second_blur_prob: 0.5
+resize_prob2: [0.3, 0.4, 0.3]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 12]
+poisson_scale_range2: [0.05, 1.0]
+gray_noise_prob2: 0.4
+jpeg_range2: [60, 100]
+blur_kernel_size: 21
+kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob: 0.1
+blur_sigma: [0.2, 1.5]
+betag_range: [0.5, 2.0]
+betap_range: [1, 1.5]
+blur_kernel_size2: 11
+kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob2: 0.1
+blur_sigma2: [0.2, 1.0]
+betag_range2: [0.5, 2.0]
+betap_range2: [1, 1.5]
+final_sinc_prob: 0.8
+use_hflip: True
+use_rot: False

configs/config_s3_ft.yml

@@ -0,0 +1,37 @@
+dataroot_gt: data/patches
+gt_size: 512
+iter_num: 1000
+# Real-ESRGAN style degradation (shared by all stages)
+scale: 4
+resize_prob: [0.2, 0.7, 0.1]
+resize_range: [0.3, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2.0]
+gray_noise_prob: 0.4
+jpeg_range: [60, 95]
+second_blur_prob: 0.5
+resize_prob2: [0.3, 0.4, 0.3]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 12]
+poisson_scale_range2: [0.05, 1.0]
+gray_noise_prob2: 0.4
+jpeg_range2: [60, 100]
+blur_kernel_size: 21
+kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob: 0.1
+blur_sigma: [0.2, 1.5]
+betag_range: [0.5, 2.0]
+betap_range: [1, 1.5]
+blur_kernel_size2: 11
+kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob2: 0.1
+blur_sigma2: [0.2, 1.0]
+betag_range2: [0.5, 2.0]
+betap_range2: [1, 1.5]
+final_sinc_prob: 0.8
+use_hflip: True
+use_rot: False

configs/config_smoke.yml

@@ -0,0 +1,37 @@
+dataroot_gt: data/patches
+gt_size: 128
+iter_num: 2
+# Real-ESRGAN style degradation (shared by all stages)
+scale: 4
+resize_prob: [0.2, 0.7, 0.1]
+resize_range: [0.3, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2.0]
+gray_noise_prob: 0.4
+jpeg_range: [60, 95]
+second_blur_prob: 0.5
+resize_prob2: [0.3, 0.4, 0.3]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 12]
+poisson_scale_range2: [0.05, 1.0]
+gray_noise_prob2: 0.4
+jpeg_range2: [60, 100]
+blur_kernel_size: 21
+kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob: 0.1
+blur_sigma: [0.2, 1.5]
+betag_range: [0.5, 2.0]
+betap_range: [1, 1.5]
+blur_kernel_size2: 11
+kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob2: 0.1
+blur_sigma2: [0.2, 1.0]
+betag_range2: [0.5, 2.0]
+betap_range2: [1, 1.5]
+final_sinc_prob: 0.8
+use_hflip: True
+use_rot: False

official/bsr/degradations.py

@@ -0,0 +1,764 @@
+import cv2
+import math
+import numpy as np
+import random
+import torch
+from scipy import special
+from scipy.stats import multivariate_normal
+from torchvision.transforms._functional_tensor import rgb_to_grayscale
+
+# -------------------------------------------------------------------- #
+# --------------------------- blur kernels --------------------------- #
+# -------------------------------------------------------------------- #
+
+
+# --------------------------- util functions --------------------------- #
+def sigma_matrix2(sig_x, sig_y, theta):
+    """Calculate the rotated sigma matrix (two dimensional matrix).
+
+    Args:
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+
+    Returns:
+        ndarray: Rotated sigma matrix.
+    """
+    d_matrix = np.array([[sig_x**2, 0], [0, sig_y**2]])
+    u_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
+    return np.dot(u_matrix, np.dot(d_matrix, u_matrix.T))
+
+
+def mesh_grid(kernel_size):
+    """Generate the mesh grid, centering at zero.
+
+    Args:
+        kernel_size (int):
+
+    Returns:
+        xy (ndarray): with the shape (kernel_size, kernel_size, 2)
+        xx (ndarray): with the shape (kernel_size, kernel_size)
+        yy (ndarray): with the shape (kernel_size, kernel_size)
+    """
+    ax = np.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.)
+    xx, yy = np.meshgrid(ax, ax)
+    xy = np.hstack((xx.reshape((kernel_size * kernel_size, 1)), yy.reshape(kernel_size * kernel_size,
+                                                                           1))).reshape(kernel_size, kernel_size, 2)
+    return xy, xx, yy
+
+
+def pdf2(sigma_matrix, grid):
+    """Calculate PDF of the bivariate Gaussian distribution.
+
+    Args:
+        sigma_matrix (ndarray): with the shape (2, 2)
+        grid (ndarray): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size.
+
+    Returns:
+        kernel (ndarrray): un-normalized kernel.
+    """
+    inverse_sigma = np.linalg.inv(sigma_matrix)
+    kernel = np.exp(-0.5 * np.sum(np.dot(grid, inverse_sigma) * grid, 2))
+    return kernel
+
+
+def cdf2(d_matrix, grid):
+    """Calculate the CDF of the standard bivariate Gaussian distribution.
+        Used in skewed Gaussian distribution.
+
+    Args:
+        d_matrix (ndarrasy): skew matrix.
+        grid (ndarray): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size.
+
+    Returns:
+        cdf (ndarray): skewed cdf.
+    """
+    rv = multivariate_normal([0, 0], [[1, 0], [0, 1]])
+    grid = np.dot(grid, d_matrix)
+    cdf = rv.cdf(grid)
+    return cdf
+
+
+def bivariate_Gaussian(kernel_size, sig_x, sig_y, theta, grid=None, isotropic=True):
+    """Generate a bivariate isotropic or anisotropic Gaussian kernel.
+
+    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
+
+    Args:
+        kernel_size (int):
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+        grid (ndarray, optional): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size. Default: None
+        isotropic (bool):
+
+    Returns:
+        kernel (ndarray): normalized kernel.
+    """
+    if grid is None:
+        grid, _, _ = mesh_grid(kernel_size)
+    if isotropic:
+        sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
+    else:
+        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
+    kernel = pdf2(sigma_matrix, grid)
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def bivariate_generalized_Gaussian(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
+    """Generate a bivariate generalized Gaussian kernel.
+
+    ``Paper: Parameter Estimation For Multivariate Generalized Gaussian Distributions``
+
+    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
+
+    Args:
+        kernel_size (int):
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+        beta (float): shape parameter, beta = 1 is the normal distribution.
+        grid (ndarray, optional): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size. Default: None
+
+    Returns:
+        kernel (ndarray): normalized kernel.
+    """
+    if grid is None:
+        grid, _, _ = mesh_grid(kernel_size)
+    if isotropic:
+        sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
+    else:
+        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
+    inverse_sigma = np.linalg.inv(sigma_matrix)
+    kernel = np.exp(-0.5 * np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta))
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def bivariate_plateau(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
+    """Generate a plateau-like anisotropic kernel.
+
+    1 / (1+x^(beta))
+
+    Reference: https://stats.stackexchange.com/questions/203629/is-there-a-plateau-shaped-distribution
+
+    In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
+
+    Args:
+        kernel_size (int):
+        sig_x (float):
+        sig_y (float):
+        theta (float): Radian measurement.
+        beta (float): shape parameter, beta = 1 is the normal distribution.
+        grid (ndarray, optional): generated by :func:`mesh_grid`,
+            with the shape (K, K, 2), K is the kernel size. Default: None
+
+    Returns:
+        kernel (ndarray): normalized kernel.
+    """
+    if grid is None:
+        grid, _, _ = mesh_grid(kernel_size)
+    if isotropic:
+        sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
+    else:
+        sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
+    inverse_sigma = np.linalg.inv(sigma_matrix)
+    kernel = np.reciprocal(np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta) + 1)
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def random_bivariate_Gaussian(kernel_size,
+                              sigma_x_range,
+                              sigma_y_range,
+                              rotation_range,
+                              noise_range=None,
+                              isotropic=True):
+    """Randomly generate bivariate isotropic or anisotropic Gaussian kernels.
+
+    In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
+
+    Args:
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi, math.pi]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
+    sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
+    if isotropic is False:
+        assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
+        assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
+        sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
+        rotation = np.random.uniform(rotation_range[0], rotation_range[1])
+    else:
+        sigma_y = sigma_x
+        rotation = 0
+
+    kernel = bivariate_Gaussian(kernel_size, sigma_x, sigma_y, rotation, isotropic=isotropic)
+
+    # add multiplicative noise
+    if noise_range is not None:
+        assert noise_range[0] < noise_range[1], 'Wrong noise range.'
+        noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
+        kernel = kernel * noise
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def random_bivariate_generalized_Gaussian(kernel_size,
+                                          sigma_x_range,
+                                          sigma_y_range,
+                                          rotation_range,
+                                          beta_range,
+                                          noise_range=None,
+                                          isotropic=True):
+    """Randomly generate bivariate generalized Gaussian kernels.
+
+    In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
+
+    Args:
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi, math.pi]
+        beta_range (tuple): [0.5, 8]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
+    sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
+    if isotropic is False:
+        assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
+        assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
+        sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
+        rotation = np.random.uniform(rotation_range[0], rotation_range[1])
+    else:
+        sigma_y = sigma_x
+        rotation = 0
+
+    # assume beta_range[0] < 1 < beta_range[1]
+    if np.random.uniform() < 0.5:
+        beta = np.random.uniform(beta_range[0], 1)
+    else:
+        beta = np.random.uniform(1, beta_range[1])
+
+    kernel = bivariate_generalized_Gaussian(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
+
+    # add multiplicative noise
+    if noise_range is not None:
+        assert noise_range[0] < noise_range[1], 'Wrong noise range.'
+        noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
+        kernel = kernel * noise
+    kernel = kernel / np.sum(kernel)
+    return kernel
+
+
+def random_bivariate_plateau(kernel_size,
+                             sigma_x_range,
+                             sigma_y_range,
+                             rotation_range,
+                             beta_range,
+                             noise_range=None,
+                             isotropic=True):
+    """Randomly generate bivariate plateau kernels.
+
+    In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
+
+    Args:
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi/2, math.pi/2]
+        beta_range (tuple): [1, 4]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
+    sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
+    if isotropic is False:
+        assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
+        assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
+        sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
+        rotation = np.random.uniform(rotation_range[0], rotation_range[1])
+    else:
+        sigma_y = sigma_x
+        rotation = 0
+
+    # TODO: this may be not proper
+    if np.random.uniform() < 0.5:
+        beta = np.random.uniform(beta_range[0], 1)
+    else:
+        beta = np.random.uniform(1, beta_range[1])
+
+    kernel = bivariate_plateau(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
+    # add multiplicative noise
+    if noise_range is not None:
+        assert noise_range[0] < noise_range[1], 'Wrong noise range.'
+        noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
+        kernel = kernel * noise
+    kernel = kernel / np.sum(kernel)
+
+    return kernel
+
+
+def random_mixed_kernels(kernel_list,
+                         kernel_prob,
+                         kernel_size=21,
+                         sigma_x_range=(0.6, 5),
+                         sigma_y_range=(0.6, 5),
+                         rotation_range=(-math.pi, math.pi),
+                         betag_range=(0.5, 8),
+                         betap_range=(0.5, 8),
+                         noise_range=None):
+    """Randomly generate mixed kernels.
+
+    Args:
+        kernel_list (tuple): a list name of kernel types,
+            support ['iso', 'aniso', 'skew', 'generalized', 'plateau_iso',
+            'plateau_aniso']
+        kernel_prob (tuple): corresponding kernel probability for each
+            kernel type
+        kernel_size (int):
+        sigma_x_range (tuple): [0.6, 5]
+        sigma_y_range (tuple): [0.6, 5]
+        rotation range (tuple): [-math.pi, math.pi]
+        beta_range (tuple): [0.5, 8]
+        noise_range(tuple, optional): multiplicative kernel noise,
+            [0.75, 1.25]. Default: None
+
+    Returns:
+        kernel (ndarray):
+    """
+    kernel_type = random.choices(kernel_list, kernel_prob)[0]
+    if kernel_type == 'iso':
+        kernel = random_bivariate_Gaussian(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=True)
+    elif kernel_type == 'aniso':
+        kernel = random_bivariate_Gaussian(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=False)
+    elif kernel_type == 'generalized_iso':
+        kernel = random_bivariate_generalized_Gaussian(
+            kernel_size,
+            sigma_x_range,
+            sigma_y_range,
+            rotation_range,
+            betag_range,
+            noise_range=noise_range,
+            isotropic=True)
+    elif kernel_type == 'generalized_aniso':
+        kernel = random_bivariate_generalized_Gaussian(
+            kernel_size,
+            sigma_x_range,
+            sigma_y_range,
+            rotation_range,
+            betag_range,
+            noise_range=noise_range,
+            isotropic=False)
+    elif kernel_type == 'plateau_iso':
+        kernel = random_bivariate_plateau(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=True)
+    elif kernel_type == 'plateau_aniso':
+        kernel = random_bivariate_plateau(
+            kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=False)
+    return kernel
+
+
+np.seterr(divide='ignore', invalid='ignore')
+
+
+def circular_lowpass_kernel(cutoff, kernel_size, pad_to=0):
+    """2D sinc filter
+
+    Reference: https://dsp.stackexchange.com/questions/58301/2-d-circularly-symmetric-low-pass-filter
+
+    Args:
+        cutoff (float): cutoff frequency in radians (pi is max)
+        kernel_size (int): horizontal and vertical size, must be odd.
+        pad_to (int): pad kernel size to desired size, must be odd or zero.
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    kernel = np.fromfunction(
+        lambda x, y: cutoff * special.j1(cutoff * np.sqrt(
+            (x - (kernel_size - 1) / 2)**2 + (y - (kernel_size - 1) / 2)**2)) / (2 * np.pi * np.sqrt(
+                (x - (kernel_size - 1) / 2)**2 + (y - (kernel_size - 1) / 2)**2)), [kernel_size, kernel_size])
+    kernel[(kernel_size - 1) // 2, (kernel_size - 1) // 2] = cutoff**2 / (4 * np.pi)
+    kernel = kernel / np.sum(kernel)
+    if pad_to > kernel_size:
+        pad_size = (pad_to - kernel_size) // 2
+        kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
+    return kernel
+
+
+# ------------------------------------------------------------- #
+# --------------------------- noise --------------------------- #
+# ------------------------------------------------------------- #
+
+# ----------------------- Gaussian Noise ----------------------- #
+
+
+def generate_gaussian_noise(img, sigma=10, gray_noise=False):
+    """Generate Gaussian noise.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        sigma (float): Noise scale (measured in range 255). Default: 10.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    if gray_noise:
+        noise = np.float32(np.random.randn(*(img.shape[0:2]))) * sigma / 255.
+        noise = np.expand_dims(noise, axis=2).repeat(3, axis=2)
+    else:
+        noise = np.float32(np.random.randn(*(img.shape))) * sigma / 255.
+    return noise
+
+
+def add_gaussian_noise(img, sigma=10, clip=True, rounds=False, gray_noise=False):
+    """Add Gaussian noise.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        sigma (float): Noise scale (measured in range 255). Default: 10.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    noise = generate_gaussian_noise(img, sigma, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def generate_gaussian_noise_pt(img, sigma=10, gray_noise=0):
+    """Add Gaussian noise (PyTorch version).
+
+    Args:
+        img (Tensor): Shape (b, c, h, w), range[0, 1], float32.
+        scale (float | Tensor): Noise scale. Default: 1.0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    b, _, h, w = img.size()
+    if not isinstance(sigma, (float, int)):
+        sigma = sigma.view(img.size(0), 1, 1, 1)
+    if isinstance(gray_noise, (float, int)):
+        cal_gray_noise = gray_noise > 0
+    else:
+        gray_noise = gray_noise.view(b, 1, 1, 1)
+        cal_gray_noise = torch.sum(gray_noise) > 0
+
+    if cal_gray_noise:
+        noise_gray = torch.randn(*img.size()[2:4], dtype=img.dtype, device=img.device) * sigma / 255.
+        noise_gray = noise_gray.view(b, 1, h, w)
+
+    # always calculate color noise
+    noise = torch.randn(*img.size(), dtype=img.dtype, device=img.device) * sigma / 255.
+
+    if cal_gray_noise:
+        noise = noise * (1 - gray_noise) + noise_gray * gray_noise
+    return noise
+
+
+def add_gaussian_noise_pt(img, sigma=10, gray_noise=0, clip=True, rounds=False):
+    """Add Gaussian noise (PyTorch version).
+
+    Args:
+        img (Tensor): Shape (b, c, h, w), range[0, 1], float32.
+        scale (float | Tensor): Noise scale. Default: 1.0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    noise = generate_gaussian_noise_pt(img, sigma, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ----------------------- Random Gaussian Noise ----------------------- #
+def random_generate_gaussian_noise(img, sigma_range=(0, 10), gray_prob=0):
+    sigma = np.random.uniform(sigma_range[0], sigma_range[1])
+    if np.random.uniform() < gray_prob:
+        gray_noise = True
+    else:
+        gray_noise = False
+    return generate_gaussian_noise(img, sigma, gray_noise)
+
+
+def random_add_gaussian_noise(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_gaussian_noise(img, sigma_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def random_generate_gaussian_noise_pt(img, sigma_range=(0, 10), gray_prob=0):
+    sigma = torch.rand(
+        img.size(0), dtype=img.dtype, device=img.device) * (sigma_range[1] - sigma_range[0]) + sigma_range[0]
+    gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
+    gray_noise = (gray_noise < gray_prob).float()
+    return generate_gaussian_noise_pt(img, sigma, gray_noise)
+
+
+def random_add_gaussian_noise_pt(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_gaussian_noise_pt(img, sigma_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ----------------------- Poisson (Shot) Noise ----------------------- #
+
+
+def generate_poisson_noise(img, scale=1.0, gray_noise=False):
+    """Generate poisson noise.
+
+    Reference: https://github.com/scikit-image/scikit-image/blob/main/skimage/util/noise.py#L37-L219
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        scale (float): Noise scale. Default: 1.0.
+        gray_noise (bool): Whether generate gray noise. Default: False.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    if gray_noise:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    # round and clip image for counting vals correctly
+    img = np.clip((img * 255.0).round(), 0, 255) / 255.
+    vals = len(np.unique(img))
+    vals = 2**np.ceil(np.log2(vals))
+    out = np.float32(np.random.poisson(img * vals) / float(vals))
+    noise = out - img
+    if gray_noise:
+        noise = np.repeat(noise[:, :, np.newaxis], 3, axis=2)
+    return noise * scale
+
+
+def add_poisson_noise(img, scale=1.0, clip=True, rounds=False, gray_noise=False):
+    """Add poisson noise.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        scale (float): Noise scale. Default: 1.0.
+        gray_noise (bool): Whether generate gray noise. Default: False.
+
+    Returns:
+        (Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    noise = generate_poisson_noise(img, scale, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def generate_poisson_noise_pt(img, scale=1.0, gray_noise=0):
+    """Generate a batch of poisson noise (PyTorch version)
+
+    Args:
+        img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
+        scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
+            Default: 1.0.
+        gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
+            0 for False, 1 for True. Default: 0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    b, _, h, w = img.size()
+    if isinstance(gray_noise, (float, int)):
+        cal_gray_noise = gray_noise > 0
+    else:
+        gray_noise = gray_noise.view(b, 1, 1, 1)
+        cal_gray_noise = torch.sum(gray_noise) > 0
+    if cal_gray_noise:
+        img_gray = rgb_to_grayscale(img, num_output_channels=1)
+        # round and clip image for counting vals correctly
+        img_gray = torch.clamp((img_gray * 255.0).round(), 0, 255) / 255.
+        # use for-loop to get the unique values for each sample
+        vals_list = [len(torch.unique(img_gray[i, :, :, :])) for i in range(b)]
+        vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]
+        vals = img_gray.new_tensor(vals_list).view(b, 1, 1, 1)
+        out = torch.poisson(img_gray * vals) / vals
+        noise_gray = out - img_gray
+        noise_gray = noise_gray.expand(b, 3, h, w)
+
+    # always calculate color noise
+    # round and clip image for counting vals correctly
+    img = torch.clamp((img * 255.0).round(), 0, 255) / 255.
+    # use for-loop to get the unique values for each sample
+    vals_list = [len(torch.unique(img[i, :, :, :])) for i in range(b)]
+    vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]
+    vals = img.new_tensor(vals_list).view(b, 1, 1, 1)
+    out = torch.poisson(img * vals) / vals
+    noise = out - img
+    if cal_gray_noise:
+        noise = noise * (1 - gray_noise) + noise_gray * gray_noise
+    if not isinstance(scale, (float, int)):
+        scale = scale.view(b, 1, 1, 1)
+    return noise * scale
+
+
+def add_poisson_noise_pt(img, scale=1.0, clip=True, rounds=False, gray_noise=0):
+    """Add poisson noise to a batch of images (PyTorch version).
+
+    Args:
+        img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
+        scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
+            Default: 1.0.
+        gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
+            0 for False, 1 for True. Default: 0.
+
+    Returns:
+        (Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
+            float32.
+    """
+    noise = generate_poisson_noise_pt(img, scale, gray_noise)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ----------------------- Random Poisson (Shot) Noise ----------------------- #
+
+
+def random_generate_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0):
+    scale = np.random.uniform(scale_range[0], scale_range[1])
+    if np.random.uniform() < gray_prob:
+        gray_noise = True
+    else:
+        gray_noise = False
+    return generate_poisson_noise(img, scale, gray_noise)
+
+
+def random_add_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_poisson_noise(img, scale_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = np.clip((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = np.clip(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+def random_generate_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0):
+    scale = torch.rand(
+        img.size(0), dtype=img.dtype, device=img.device) * (scale_range[1] - scale_range[0]) + scale_range[0]
+    gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
+    gray_noise = (gray_noise < gray_prob).float()
+    return generate_poisson_noise_pt(img, scale, gray_noise)
+
+
+def random_add_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
+    noise = random_generate_poisson_noise_pt(img, scale_range, gray_prob)
+    out = img + noise
+    if clip and rounds:
+        out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+    elif clip:
+        out = torch.clamp(out, 0, 1)
+    elif rounds:
+        out = (out * 255.0).round() / 255.
+    return out
+
+
+# ------------------------------------------------------------------------ #
+# --------------------------- JPEG compression --------------------------- #
+# ------------------------------------------------------------------------ #
+
+
+def add_jpg_compression(img, quality=90):
+    """Add JPG compression artifacts.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        quality (float): JPG compression quality. 0 for lowest quality, 100 for
+            best quality. Default: 90.
+
+    Returns:
+        (Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    img = np.clip(img, 0, 1)
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
+    _, encimg = cv2.imencode('.jpg', img * 255., encode_param)
+    img = np.float32(cv2.imdecode(encimg, 1)) / 255.
+    return img
+
+
+def random_add_jpg_compression(img, quality_range=(90, 100)):
+    """Randomly add JPG compression artifacts.
+
+    Args:
+        img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
+        quality_range (tuple[float] | list[float]): JPG compression quality
+            range. 0 for lowest quality, 100 for best quality.
+            Default: (90, 100).
+
+    Returns:
+        (Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],
+            float32.
+    """
+    quality = np.random.uniform(quality_range[0], quality_range[1])
+    return add_jpg_compression(img, quality)

official/bsr/transforms.py

@@ -0,0 +1,179 @@
+import cv2
+import random
+import torch
+
+
+def mod_crop(img, scale):
+    """Mod crop images, used during testing.
+
+    Args:
+        img (ndarray): Input image.
+        scale (int): Scale factor.
+
+    Returns:
+        ndarray: Result image.
+    """
+    img = img.copy()
+    if img.ndim in (2, 3):
+        h, w = img.shape[0], img.shape[1]
+        h_remainder, w_remainder = h % scale, w % scale
+        img = img[:h - h_remainder, :w - w_remainder, ...]
+    else:
+        raise ValueError(f'Wrong img ndim: {img.ndim}.')
+    return img
+
+
+def paired_random_crop(img_gts, img_lqs, gt_patch_size, scale, gt_path=None):
+    """Paired random crop. Support Numpy array and Tensor inputs.
+
+    It crops lists of lq and gt images with corresponding locations.
+
+    Args:
+        img_gts (list[ndarray] | ndarray | list[Tensor] | Tensor): GT images. Note that all images
+            should have the same shape. If the input is an ndarray, it will
+            be transformed to a list containing itself.
+        img_lqs (list[ndarray] | ndarray): LQ images. Note that all images
+            should have the same shape. If the input is an ndarray, it will
+            be transformed to a list containing itself.
+        gt_patch_size (int): GT patch size.
+        scale (int): Scale factor.
+        gt_path (str): Path to ground-truth. Default: None.
+
+    Returns:
+        list[ndarray] | ndarray: GT images and LQ images. If returned results
+            only have one element, just return ndarray.
+    """
+
+    if not isinstance(img_gts, list):
+        img_gts = [img_gts]
+    if not isinstance(img_lqs, list):
+        img_lqs = [img_lqs]
+
+    # determine input type: Numpy array or Tensor
+    input_type = 'Tensor' if torch.is_tensor(img_gts[0]) else 'Numpy'
+
+    if input_type == 'Tensor':
+        h_lq, w_lq = img_lqs[0].size()[-2:]
+        h_gt, w_gt = img_gts[0].size()[-2:]
+    else:
+        h_lq, w_lq = img_lqs[0].shape[0:2]
+        h_gt, w_gt = img_gts[0].shape[0:2]
+    lq_patch_size = gt_patch_size // scale
+
+    if h_gt != h_lq * scale or w_gt != w_lq * scale:
+        raise ValueError(f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',
+                         f'multiplication of LQ ({h_lq}, {w_lq}).')
+    if h_lq < lq_patch_size or w_lq < lq_patch_size:
+        raise ValueError(f'LQ ({h_lq}, {w_lq}) is smaller than patch size '
+                         f'({lq_patch_size}, {lq_patch_size}). '
+                         f'Please remove {gt_path}.')
+
+    # randomly choose top and left coordinates for lq patch
+    top = random.randint(0, h_lq - lq_patch_size)
+    left = random.randint(0, w_lq - lq_patch_size)
+
+    # crop lq patch
+    if input_type == 'Tensor':
+        img_lqs = [v[:, :, top:top + lq_patch_size, left:left + lq_patch_size] for v in img_lqs]
+    else:
+        img_lqs = [v[top:top + lq_patch_size, left:left + lq_patch_size, ...] for v in img_lqs]
+
+    # crop corresponding gt patch
+    top_gt, left_gt = int(top * scale), int(left * scale)
+    if input_type == 'Tensor':
+        img_gts = [v[:, :, top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size] for v in img_gts]
+    else:
+        img_gts = [v[top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size, ...] for v in img_gts]
+    if len(img_gts) == 1:
+        img_gts = img_gts[0]
+    if len(img_lqs) == 1:
+        img_lqs = img_lqs[0]
+    return img_gts, img_lqs
+
+
+def augment(imgs, hflip=True, rotation=True, flows=None, return_status=False):
+    """Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees).
+
+    We use vertical flip and transpose for rotation implementation.
+    All the images in the list use the same augmentation.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Images to be augmented. If the input
+            is an ndarray, it will be transformed to a list.
+        hflip (bool): Horizontal flip. Default: True.
+        rotation (bool): Ratotation. Default: True.
+        flows (list[ndarray]: Flows to be augmented. If the input is an
+            ndarray, it will be transformed to a list.
+            Dimension is (h, w, 2). Default: None.
+        return_status (bool): Return the status of flip and rotation.
+            Default: False.
+
+    Returns:
+        list[ndarray] | ndarray: Augmented images and flows. If returned
+            results only have one element, just return ndarray.
+
+    """
+    hflip = hflip and random.random() < 0.5
+    vflip = rotation and random.random() < 0.5
+    rot90 = rotation and random.random() < 0.5
+
+    def _augment(img):
+        if hflip:  # horizontal
+            cv2.flip(img, 1, img)
+        if vflip:  # vertical
+            cv2.flip(img, 0, img)
+        if rot90:
+            img = img.transpose(1, 0, 2)
+        return img
+
+    def _augment_flow(flow):
+        if hflip:  # horizontal
+            cv2.flip(flow, 1, flow)
+            flow[:, :, 0] *= -1
+        if vflip:  # vertical
+            cv2.flip(flow, 0, flow)
+            flow[:, :, 1] *= -1
+        if rot90:
+            flow = flow.transpose(1, 0, 2)
+            flow = flow[:, :, [1, 0]]
+        return flow
+
+    if not isinstance(imgs, list):
+        imgs = [imgs]
+    imgs = [_augment(img) for img in imgs]
+    if len(imgs) == 1:
+        imgs = imgs[0]
+
+    if flows is not None:
+        if not isinstance(flows, list):
+            flows = [flows]
+        flows = [_augment_flow(flow) for flow in flows]
+        if len(flows) == 1:
+            flows = flows[0]
+        return imgs, flows
+    else:
+        if return_status:
+            return imgs, (hflip, vflip, rot90)
+        else:
+            return imgs
+
+
+def img_rotate(img, angle, center=None, scale=1.0):
+    """Rotate image.
+
+    Args:
+        img (ndarray): Image to be rotated.
+        angle (float): Rotation angle in degrees. Positive values mean
+            counter-clockwise rotation.
+        center (tuple[int]): Rotation center. If the center is None,
+            initialize it as the center of the image. Default: None.
+        scale (float): Isotropic scale factor. Default: 1.0.
+    """
+    (h, w) = img.shape[:2]
+
+    if center is None:
+        center = (w // 2, h // 2)
+
+    matrix = cv2.getRotationMatrix2D(center, angle, scale)
+    rotated_img = cv2.warpAffine(img, matrix, (w, h))
+    return rotated_img

official/bsr/utils/__init__.py

@@ -0,0 +1,47 @@
+from .color_util import bgr2ycbcr, rgb2ycbcr, rgb2ycbcr_pt, ycbcr2bgr, ycbcr2rgb
+from .diffjpeg import DiffJPEG
+from .file_client import FileClient
+from .img_process_util import USMSharp, usm_sharp
+from .img_util import crop_border, imfrombytes, img2tensor, imwrite, tensor2img
+from .logger import AvgTimer, MessageLogger, get_env_info, get_root_logger, init_tb_logger, init_wandb_logger
+from .misc import check_resume, get_time_str, make_exp_dirs, mkdir_and_rename, scandir, set_random_seed, sizeof_fmt
+from .options import yaml_load
+
+__all__ = [
+    #  color_util.py
+    'bgr2ycbcr',
+    'rgb2ycbcr',
+    'rgb2ycbcr_pt',
+    'ycbcr2bgr',
+    'ycbcr2rgb',
+    # file_client.py
+    'FileClient',
+    # img_util.py
+    'img2tensor',
+    'tensor2img',
+    'imfrombytes',
+    'imwrite',
+    'crop_border',
+    # logger.py
+    'MessageLogger',
+    'AvgTimer',
+    'init_tb_logger',
+    'init_wandb_logger',
+    'get_root_logger',
+    'get_env_info',
+    # misc.py
+    'set_random_seed',
+    'get_time_str',
+    'mkdir_and_rename',
+    'make_exp_dirs',
+    'scandir',
+    'check_resume',
+    'sizeof_fmt',
+    # diffjpeg
+    'DiffJPEG',
+    # img_process_util
+    'USMSharp',
+    'usm_sharp',
+    # options
+    'yaml_load'
+]

official/bsr/utils/color_util.py

@@ -0,0 +1,208 @@
+import numpy as np
+import torch
+
+
+def rgb2ycbcr(img, y_only=False):
+    """Convert a RGB image to YCbCr image.
+
+    This function produces the same results as Matlab's `rgb2ycbcr` function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def bgr2ycbcr(img, y_only=False):
+    """Convert a BGR image to YCbCr image.
+
+    The bgr version of rgb2ycbcr.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def ycbcr2rgb(img):
+    """Convert a YCbCr image to RGB image.
+
+    This function produces the same results as Matlab's ycbcr2rgb function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> RGB`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        ndarray: The converted RGB image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+                              [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]  # noqa: E126
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def ycbcr2bgr(img):
+    """Convert a YCbCr image to BGR image.
+
+    The bgr version of ycbcr2rgb.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> BGR`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        ndarray: The converted BGR image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0.00791071, -0.00153632, 0],
+                              [0, -0.00318811, 0.00625893]]) * 255.0 + [-276.836, 135.576, -222.921]  # noqa: E126
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def _convert_input_type_range(img):
+    """Convert the type and range of the input image.
+
+    It converts the input image to np.float32 type and range of [0, 1].
+    It is mainly used for pre-processing the input image in colorspace
+    conversion functions such as rgb2ycbcr and ycbcr2rgb.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        (ndarray): The converted image with type of np.float32 and range of
+            [0, 1].
+    """
+    img_type = img.dtype
+    img = img.astype(np.float32)
+    if img_type == np.float32:
+        pass
+    elif img_type == np.uint8:
+        img /= 255.
+    else:
+        raise TypeError(f'The img type should be np.float32 or np.uint8, but got {img_type}')
+    return img
+
+
+def _convert_output_type_range(img, dst_type):
+    """Convert the type and range of the image according to dst_type.
+
+    It converts the image to desired type and range. If `dst_type` is np.uint8,
+    images will be converted to np.uint8 type with range [0, 255]. If
+    `dst_type` is np.float32, it converts the image to np.float32 type with
+    range [0, 1].
+    It is mainly used for post-processing images in colorspace conversion
+    functions such as rgb2ycbcr and ycbcr2rgb.
+
+    Args:
+        img (ndarray): The image to be converted with np.float32 type and
+            range [0, 255].
+        dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
+            converts the image to np.uint8 type with range [0, 255]. If
+            dst_type is np.float32, it converts the image to np.float32 type
+            with range [0, 1].
+
+    Returns:
+        (ndarray): The converted image with desired type and range.
+    """
+    if dst_type not in (np.uint8, np.float32):
+        raise TypeError(f'The dst_type should be np.float32 or np.uint8, but got {dst_type}')
+    if dst_type == np.uint8:
+        img = img.round()
+    else:
+        img /= 255.
+    return img.astype(dst_type)
+
+
+def rgb2ycbcr_pt(img, y_only=False):
+    """Convert RGB images to YCbCr images (PyTorch version).
+
+    It implements the ITU-R BT.601 conversion for standard-definition television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    Args:
+        img (Tensor): Images with shape (n, 3, h, w), the range [0, 1], float, RGB format.
+         y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        (Tensor): converted images with the shape (n, 3/1, h, w), the range [0, 1], float.
+    """
+    if y_only:
+        weight = torch.tensor([[65.481], [128.553], [24.966]]).to(img)
+        out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + 16.0
+    else:
+        weight = torch.tensor([[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]).to(img)
+        bias = torch.tensor([16, 128, 128]).view(1, 3, 1, 1).to(img)
+        out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + bias
+
+    out_img = out_img / 255.
+    return out_img

official/bsr/utils/diffjpeg.py

@@ -0,0 +1,515 @@
+"""
+Modified from https://github.com/mlomnitz/DiffJPEG
+
+For images not divisible by 8
+https://dsp.stackexchange.com/questions/35339/jpeg-dct-padding/35343#35343
+"""
+import itertools
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+# ------------------------ utils ------------------------#
+y_table = np.array(
+    [[16, 11, 10, 16, 24, 40, 51, 61], [12, 12, 14, 19, 26, 58, 60, 55], [14, 13, 16, 24, 40, 57, 69, 56],
+     [14, 17, 22, 29, 51, 87, 80, 62], [18, 22, 37, 56, 68, 109, 103, 77], [24, 35, 55, 64, 81, 104, 113, 92],
+     [49, 64, 78, 87, 103, 121, 120, 101], [72, 92, 95, 98, 112, 100, 103, 99]],
+    dtype=np.float32).T
+y_table = nn.Parameter(torch.from_numpy(y_table))
+c_table = np.empty((8, 8), dtype=np.float32)
+c_table.fill(99)
+c_table[:4, :4] = np.array([[17, 18, 24, 47], [18, 21, 26, 66], [24, 26, 56, 99], [47, 66, 99, 99]]).T
+c_table = nn.Parameter(torch.from_numpy(c_table))
+
+
+def diff_round(x):
+    """ Differentiable rounding function
+    """
+    return torch.round(x) + (x - torch.round(x))**3
+
+
+def quality_to_factor(quality):
+    """ Calculate factor corresponding to quality
+
+    Args:
+        quality(float): Quality for jpeg compression.
+
+    Returns:
+        float: Compression factor.
+    """
+    if quality < 50:
+        quality = 5000. / quality
+    else:
+        quality = 200. - quality * 2
+    return quality / 100.
+
+
+# ------------------------ compression ------------------------#
+class RGB2YCbCrJpeg(nn.Module):
+    """ Converts RGB image to YCbCr
+    """
+
+    def __init__(self):
+        super(RGB2YCbCrJpeg, self).__init__()
+        matrix = np.array([[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], [0.5, -0.418688, -0.081312]],
+                          dtype=np.float32).T
+        self.shift = nn.Parameter(torch.tensor([0., 128., 128.]))
+        self.matrix = nn.Parameter(torch.from_numpy(matrix))
+
+    def forward(self, image):
+        """
+        Args:
+            image(Tensor): batch x 3 x height x width
+
+        Returns:
+            Tensor: batch x height x width x 3
+        """
+        image = image.permute(0, 2, 3, 1)
+        result = torch.tensordot(image, self.matrix, dims=1) + self.shift
+        return result.view(image.shape)
+
+
+class ChromaSubsampling(nn.Module):
+    """ Chroma subsampling on CbCr channels
+    """
+
+    def __init__(self):
+        super(ChromaSubsampling, self).__init__()
+
+    def forward(self, image):
+        """
+        Args:
+            image(tensor): batch x height x width x 3
+
+        Returns:
+            y(tensor): batch x height x width
+            cb(tensor): batch x height/2 x width/2
+            cr(tensor): batch x height/2 x width/2
+        """
+        image_2 = image.permute(0, 3, 1, 2).clone()
+        cb = F.avg_pool2d(image_2[:, 1, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False)
+        cr = F.avg_pool2d(image_2[:, 2, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False)
+        cb = cb.permute(0, 2, 3, 1)
+        cr = cr.permute(0, 2, 3, 1)
+        return image[:, :, :, 0], cb.squeeze(3), cr.squeeze(3)
+
+
+class BlockSplitting(nn.Module):
+    """ Splitting image into patches
+    """
+
+    def __init__(self):
+        super(BlockSplitting, self).__init__()
+        self.k = 8
+
+    def forward(self, image):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor:  batch x h*w/64 x h x w
+        """
+        height, _ = image.shape[1:3]
+        batch_size = image.shape[0]
+        image_reshaped = image.view(batch_size, height // self.k, self.k, -1, self.k)
+        image_transposed = image_reshaped.permute(0, 1, 3, 2, 4)
+        return image_transposed.contiguous().view(batch_size, -1, self.k, self.k)
+
+
+class DCT8x8(nn.Module):
+    """ Discrete Cosine Transformation
+    """
+
+    def __init__(self):
+        super(DCT8x8, self).__init__()
+        tensor = np.zeros((8, 8, 8, 8), dtype=np.float32)
+        for x, y, u, v in itertools.product(range(8), repeat=4):
+            tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos((2 * y + 1) * v * np.pi / 16)
+        alpha = np.array([1. / np.sqrt(2)] + [1] * 7)
+        self.tensor = nn.Parameter(torch.from_numpy(tensor).float())
+        self.scale = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha) * 0.25).float())
+
+    def forward(self, image):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        image = image - 128
+        result = self.scale * torch.tensordot(image, self.tensor, dims=2)
+        result.view(image.shape)
+        return result
+
+
+class YQuantize(nn.Module):
+    """ JPEG Quantization for Y channel
+
+    Args:
+        rounding(function): rounding function to use
+    """
+
+    def __init__(self, rounding):
+        super(YQuantize, self).__init__()
+        self.rounding = rounding
+        self.y_table = y_table
+
+    def forward(self, image, factor=1):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        if isinstance(factor, (int, float)):
+            image = image.float() / (self.y_table * factor)
+        else:
+            b = factor.size(0)
+            table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
+            image = image.float() / table
+        image = self.rounding(image)
+        return image
+
+
+class CQuantize(nn.Module):
+    """ JPEG Quantization for CbCr channels
+
+    Args:
+        rounding(function): rounding function to use
+    """
+
+    def __init__(self, rounding):
+        super(CQuantize, self).__init__()
+        self.rounding = rounding
+        self.c_table = c_table
+
+    def forward(self, image, factor=1):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        if isinstance(factor, (int, float)):
+            image = image.float() / (self.c_table * factor)
+        else:
+            b = factor.size(0)
+            table = self.c_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
+            image = image.float() / table
+        image = self.rounding(image)
+        return image
+
+
+class CompressJpeg(nn.Module):
+    """Full JPEG compression algorithm
+
+    Args:
+        rounding(function): rounding function to use
+    """
+
+    def __init__(self, rounding=torch.round):
+        super(CompressJpeg, self).__init__()
+        self.l1 = nn.Sequential(RGB2YCbCrJpeg(), ChromaSubsampling())
+        self.l2 = nn.Sequential(BlockSplitting(), DCT8x8())
+        self.c_quantize = CQuantize(rounding=rounding)
+        self.y_quantize = YQuantize(rounding=rounding)
+
+    def forward(self, image, factor=1):
+        """
+        Args:
+            image(tensor): batch x 3 x height x width
+
+        Returns:
+            dict(tensor): Compressed tensor with batch x h*w/64 x 8 x 8.
+        """
+        y, cb, cr = self.l1(image * 255)
+        components = {'y': y, 'cb': cb, 'cr': cr}
+        for k in components.keys():
+            comp = self.l2(components[k])
+            if k in ('cb', 'cr'):
+                comp = self.c_quantize(comp, factor=factor)
+            else:
+                comp = self.y_quantize(comp, factor=factor)
+
+            components[k] = comp
+
+        return components['y'], components['cb'], components['cr']
+
+
+# ------------------------ decompression ------------------------#
+
+
+class YDequantize(nn.Module):
+    """Dequantize Y channel
+    """
+
+    def __init__(self):
+        super(YDequantize, self).__init__()
+        self.y_table = y_table
+
+    def forward(self, image, factor=1):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        if isinstance(factor, (int, float)):
+            out = image * (self.y_table * factor)
+        else:
+            b = factor.size(0)
+            table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
+            out = image * table
+        return out
+
+
+class CDequantize(nn.Module):
+    """Dequantize CbCr channel
+    """
+
+    def __init__(self):
+        super(CDequantize, self).__init__()
+        self.c_table = c_table
+
+    def forward(self, image, factor=1):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        if isinstance(factor, (int, float)):
+            out = image * (self.c_table * factor)
+        else:
+            b = factor.size(0)
+            table = self.c_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
+            out = image * table
+        return out
+
+
+class iDCT8x8(nn.Module):
+    """Inverse discrete Cosine Transformation
+    """
+
+    def __init__(self):
+        super(iDCT8x8, self).__init__()
+        alpha = np.array([1. / np.sqrt(2)] + [1] * 7)
+        self.alpha = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha)).float())
+        tensor = np.zeros((8, 8, 8, 8), dtype=np.float32)
+        for x, y, u, v in itertools.product(range(8), repeat=4):
+            tensor[x, y, u, v] = np.cos((2 * u + 1) * x * np.pi / 16) * np.cos((2 * v + 1) * y * np.pi / 16)
+        self.tensor = nn.Parameter(torch.from_numpy(tensor).float())
+
+    def forward(self, image):
+        """
+        Args:
+            image(tensor): batch x height x width
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        image = image * self.alpha
+        result = 0.25 * torch.tensordot(image, self.tensor, dims=2) + 128
+        result.view(image.shape)
+        return result
+
+
+class BlockMerging(nn.Module):
+    """Merge patches into image
+    """
+
+    def __init__(self):
+        super(BlockMerging, self).__init__()
+
+    def forward(self, patches, height, width):
+        """
+        Args:
+            patches(tensor) batch x height*width/64, height x width
+            height(int)
+            width(int)
+
+        Returns:
+            Tensor: batch x height x width
+        """
+        k = 8
+        batch_size = patches.shape[0]
+        image_reshaped = patches.view(batch_size, height // k, width // k, k, k)
+        image_transposed = image_reshaped.permute(0, 1, 3, 2, 4)
+        return image_transposed.contiguous().view(batch_size, height, width)
+
+
+class ChromaUpsampling(nn.Module):
+    """Upsample chroma layers
+    """
+
+    def __init__(self):
+        super(ChromaUpsampling, self).__init__()
+
+    def forward(self, y, cb, cr):
+        """
+        Args:
+            y(tensor): y channel image
+            cb(tensor): cb channel
+            cr(tensor): cr channel
+
+        Returns:
+            Tensor: batch x height x width x 3
+        """
+
+        def repeat(x, k=2):
+            height, width = x.shape[1:3]
+            x = x.unsqueeze(-1)
+            x = x.repeat(1, 1, k, k)
+            x = x.view(-1, height * k, width * k)
+            return x
+
+        cb = repeat(cb)
+        cr = repeat(cr)
+        return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3)
+
+
+class YCbCr2RGBJpeg(nn.Module):
+    """Converts YCbCr image to RGB JPEG
+    """
+
+    def __init__(self):
+        super(YCbCr2RGBJpeg, self).__init__()
+
+        matrix = np.array([[1., 0., 1.402], [1, -0.344136, -0.714136], [1, 1.772, 0]], dtype=np.float32).T
+        self.shift = nn.Parameter(torch.tensor([0, -128., -128.]))
+        self.matrix = nn.Parameter(torch.from_numpy(matrix))
+
+    def forward(self, image):
+        """
+        Args:
+            image(tensor): batch x height x width x 3
+
+        Returns:
+            Tensor: batch x 3 x height x width
+        """
+        result = torch.tensordot(image + self.shift, self.matrix, dims=1)
+        return result.view(image.shape).permute(0, 3, 1, 2)
+
+
+class DeCompressJpeg(nn.Module):
+    """Full JPEG decompression algorithm
+
+    Args:
+        rounding(function): rounding function to use
+    """
+
+    def __init__(self, rounding=torch.round):
+        super(DeCompressJpeg, self).__init__()
+        self.c_dequantize = CDequantize()
+        self.y_dequantize = YDequantize()
+        self.idct = iDCT8x8()
+        self.merging = BlockMerging()
+        self.chroma = ChromaUpsampling()
+        self.colors = YCbCr2RGBJpeg()
+
+    def forward(self, y, cb, cr, imgh, imgw, factor=1):
+        """
+        Args:
+            compressed(dict(tensor)): batch x h*w/64 x 8 x 8
+            imgh(int)
+            imgw(int)
+            factor(float)
+
+        Returns:
+            Tensor: batch x 3 x height x width
+        """
+        components = {'y': y, 'cb': cb, 'cr': cr}
+        for k in components.keys():
+            if k in ('cb', 'cr'):
+                comp = self.c_dequantize(components[k], factor=factor)
+                height, width = int(imgh / 2), int(imgw / 2)
+            else:
+                comp = self.y_dequantize(components[k], factor=factor)
+                height, width = imgh, imgw
+            comp = self.idct(comp)
+            components[k] = self.merging(comp, height, width)
+            #
+        image = self.chroma(components['y'], components['cb'], components['cr'])
+        image = self.colors(image)
+
+        image = torch.min(255 * torch.ones_like(image), torch.max(torch.zeros_like(image), image))
+        return image / 255
+
+
+# ------------------------ main DiffJPEG ------------------------ #
+
+
+class DiffJPEG(nn.Module):
+    """This JPEG algorithm result is slightly different from cv2.
+    DiffJPEG supports batch processing.
+
+    Args:
+        differentiable(bool): If True, uses custom differentiable rounding function, if False, uses standard torch.round
+    """
+
+    def __init__(self, differentiable=True):
+        super(DiffJPEG, self).__init__()
+        if differentiable:
+            rounding = diff_round
+        else:
+            rounding = torch.round
+
+        self.compress = CompressJpeg(rounding=rounding)
+        self.decompress = DeCompressJpeg(rounding=rounding)
+
+    def forward(self, x, quality):
+        """
+        Args:
+            x (Tensor): Input image, bchw, rgb, [0, 1]
+            quality(float): Quality factor for jpeg compression scheme.
+        """
+        factor = quality
+        if isinstance(factor, (int, float)):
+            factor = quality_to_factor(factor)
+        else:
+            for i in range(factor.size(0)):
+                factor[i] = quality_to_factor(factor[i])
+        h, w = x.size()[-2:]
+        h_pad, w_pad = 0, 0
+        # why should use 16
+        if h % 16 != 0:
+            h_pad = 16 - h % 16
+        if w % 16 != 0:
+            w_pad = 16 - w % 16
+        x = F.pad(x, (0, w_pad, 0, h_pad), mode='constant', value=0)
+
+        y, cb, cr = self.compress(x, factor=factor)
+        recovered = self.decompress(y, cb, cr, (h + h_pad), (w + w_pad), factor=factor)
+        recovered = recovered[:, :, 0:h, 0:w]
+        return recovered
+
+
+if __name__ == '__main__':
+    import cv2
+
+    from bsr.utils import img2tensor, tensor2img
+
+    img_gt = cv2.imread('test.png') / 255.
+
+    # -------------- cv2 -------------- #
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 20]
+    _, encimg = cv2.imencode('.jpg', img_gt * 255., encode_param)
+    img_lq = np.float32(cv2.imdecode(encimg, 1))
+    cv2.imwrite('cv2_JPEG_20.png', img_lq)
+
+    # -------------- DiffJPEG -------------- #
+    jpeger = DiffJPEG(differentiable=False).cuda()
+    img_gt = img2tensor(img_gt)
+    img_gt = torch.stack([img_gt, img_gt]).cuda()
+    quality = img_gt.new_tensor([20, 40])
+    out = jpeger(img_gt, quality=quality)
+
+    cv2.imwrite('pt_JPEG_20.png', tensor2img(out[0]))
+    cv2.imwrite('pt_JPEG_40.png', tensor2img(out[1]))

official/bsr/utils/dist_util.py

@@ -0,0 +1,82 @@
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E501
+import functools
+import os
+import subprocess
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+
+def init_dist(launcher, backend='nccl', **kwargs):
+    if mp.get_start_method(allow_none=True) is None:
+        mp.set_start_method('spawn')
+    if launcher == 'pytorch':
+        _init_dist_pytorch(backend, **kwargs)
+    elif launcher == 'slurm':
+        _init_dist_slurm(backend, **kwargs)
+    else:
+        raise ValueError(f'Invalid launcher type: {launcher}')
+
+
+def _init_dist_pytorch(backend, **kwargs):
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+
+def _init_dist_slurm(backend, port=None):
+    """Initialize slurm distributed training environment.
+
+    If argument ``port`` is not specified, then the master port will be system
+    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
+    environment variable, then a default port ``29500`` will be used.
+
+    Args:
+        backend (str): Backend of torch.distributed.
+        port (int, optional): Master port. Defaults to None.
+    """
+    proc_id = int(os.environ['SLURM_PROCID'])
+    ntasks = int(os.environ['SLURM_NTASKS'])
+    node_list = os.environ['SLURM_NODELIST']
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(proc_id % num_gpus)
+    addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')
+    # specify master port
+    if port is not None:
+        os.environ['MASTER_PORT'] = str(port)
+    elif 'MASTER_PORT' in os.environ:
+        pass  # use MASTER_PORT in the environment variable
+    else:
+        # 29500 is torch.distributed default port
+        os.environ['MASTER_PORT'] = '29500'
+    os.environ['MASTER_ADDR'] = addr
+    os.environ['WORLD_SIZE'] = str(ntasks)
+    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
+    os.environ['RANK'] = str(proc_id)
+    dist.init_process_group(backend=backend)
+
+
+def get_dist_info():
+    if dist.is_available():
+        initialized = dist.is_initialized()
+    else:
+        initialized = False
+    if initialized:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    else:
+        rank = 0
+        world_size = 1
+    return rank, world_size
+
+
+def master_only(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        rank, _ = get_dist_info()
+        if rank == 0:
+            return func(*args, **kwargs)
+
+    return wrapper

official/bsr/utils/download_util.py

@@ -0,0 +1,98 @@
+import math
+import os
+import requests
+from torch.hub import download_url_to_file, get_dir
+from tqdm import tqdm
+from urllib.parse import urlparse
+
+from .misc import sizeof_fmt
+
+
+def download_file_from_google_drive(file_id, save_path):
+    """Download files from google drive.
+
+    Reference: https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive
+
+    Args:
+        file_id (str): File id.
+        save_path (str): Save path.
+    """
+
+    session = requests.Session()
+    URL = 'https://docs.google.com/uc?export=download'
+    params = {'id': file_id}
+
+    response = session.get(URL, params=params, stream=True)
+    token = get_confirm_token(response)
+    if token:
+        params['confirm'] = token
+        response = session.get(URL, params=params, stream=True)
+
+    # get file size
+    response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
+    if 'Content-Range' in response_file_size.headers:
+        file_size = int(response_file_size.headers['Content-Range'].split('/')[1])
+    else:
+        file_size = None
+
+    save_response_content(response, save_path, file_size)
+
+
+def get_confirm_token(response):
+    for key, value in response.cookies.items():
+        if key.startswith('download_warning'):
+            return value
+    return None
+
+
+def save_response_content(response, destination, file_size=None, chunk_size=32768):
+    if file_size is not None:
+        pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
+
+        readable_file_size = sizeof_fmt(file_size)
+    else:
+        pbar = None
+
+    with open(destination, 'wb') as f:
+        downloaded_size = 0
+        for chunk in response.iter_content(chunk_size):
+            downloaded_size += chunk_size
+            if pbar is not None:
+                pbar.update(1)
+                pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')
+            if chunk:  # filter out keep-alive new chunks
+                f.write(chunk)
+        if pbar is not None:
+            pbar.close()
+
+
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Load file form http url, will download models if necessary.
+
+    Reference: https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+
+    Args:
+        url (str): URL to be downloaded.
+        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
+            Default: None.
+        progress (bool): Whether to show the download progress. Default: True.
+        file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
+
+    Returns:
+        str: The path to the downloaded file.
+    """
+    if model_dir is None:  # use the pytorch hub_dir
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.abspath(os.path.join(model_dir, filename))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file

official/bsr/utils/file_client.py

@@ -0,0 +1,167 @@
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py  # noqa: E501
+from abc import ABCMeta, abstractmethod
+
+
+class BaseStorageBackend(metaclass=ABCMeta):
+    """Abstract class of storage backends.
+
+    All backends need to implement two apis: ``get()`` and ``get_text()``.
+    ``get()`` reads the file as a byte stream and ``get_text()`` reads the file
+    as texts.
+    """
+
+    @abstractmethod
+    def get(self, filepath):
+        pass
+
+    @abstractmethod
+    def get_text(self, filepath):
+        pass
+
+
+class MemcachedBackend(BaseStorageBackend):
+    """Memcached storage backend.
+
+    Attributes:
+        server_list_cfg (str): Config file for memcached server list.
+        client_cfg (str): Config file for memcached client.
+        sys_path (str | None): Additional path to be appended to `sys.path`.
+            Default: None.
+    """
+
+    def __init__(self, server_list_cfg, client_cfg, sys_path=None):
+        if sys_path is not None:
+            import sys
+            sys.path.append(sys_path)
+        try:
+            import mc
+        except ImportError:
+            raise ImportError('Please install memcached to enable MemcachedBackend.')
+
+        self.server_list_cfg = server_list_cfg
+        self.client_cfg = client_cfg
+        self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg)
+        # mc.pyvector servers as a point which points to a memory cache
+        self._mc_buffer = mc.pyvector()
+
+    def get(self, filepath):
+        filepath = str(filepath)
+        import mc
+        self._client.Get(filepath, self._mc_buffer)
+        value_buf = mc.ConvertBuffer(self._mc_buffer)
+        return value_buf
+
+    def get_text(self, filepath):
+        raise NotImplementedError
+
+
+class HardDiskBackend(BaseStorageBackend):
+    """Raw hard disks storage backend."""
+
+    def get(self, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'rb') as f:
+            value_buf = f.read()
+        return value_buf
+
+    def get_text(self, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'r') as f:
+            value_buf = f.read()
+        return value_buf
+
+
+class LmdbBackend(BaseStorageBackend):
+    """Lmdb storage backend.
+
+    Args:
+        db_paths (str | list[str]): Lmdb database paths.
+        client_keys (str | list[str]): Lmdb client keys. Default: 'default'.
+        readonly (bool, optional): Lmdb environment parameter. If True,
+            disallow any write operations. Default: True.
+        lock (bool, optional): Lmdb environment parameter. If False, when
+            concurrent access occurs, do not lock the database. Default: False.
+        readahead (bool, optional): Lmdb environment parameter. If False,
+            disable the OS filesystem readahead mechanism, which may improve
+            random read performance when a database is larger than RAM.
+            Default: False.
+
+    Attributes:
+        db_paths (list): Lmdb database path.
+        _client (list): A list of several lmdb envs.
+    """
+
+    def __init__(self, db_paths, client_keys='default', readonly=True, lock=False, readahead=False, **kwargs):
+        try:
+            import lmdb
+        except ImportError:
+            raise ImportError('Please install lmdb to enable LmdbBackend.')
+
+        if isinstance(client_keys, str):
+            client_keys = [client_keys]
+
+        if isinstance(db_paths, list):
+            self.db_paths = [str(v) for v in db_paths]
+        elif isinstance(db_paths, str):
+            self.db_paths = [str(db_paths)]
+        assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '
+                                                        f'but received {len(client_keys)} and {len(self.db_paths)}.')
+
+        self._client = {}
+        for client, path in zip(client_keys, self.db_paths):
+            self._client[client] = lmdb.open(path, readonly=readonly, lock=lock, readahead=readahead, **kwargs)
+
+    def get(self, filepath, client_key):
+        """Get values according to the filepath from one lmdb named client_key.
+
+        Args:
+            filepath (str | obj:`Path`): Here, filepath is the lmdb key.
+            client_key (str): Used for distinguishing different lmdb envs.
+        """
+        filepath = str(filepath)
+        assert client_key in self._client, (f'client_key {client_key} is not in lmdb clients.')
+        client = self._client[client_key]
+        with client.begin(write=False) as txn:
+            value_buf = txn.get(filepath.encode('ascii'))
+        return value_buf
+
+    def get_text(self, filepath):
+        raise NotImplementedError
+
+
+class FileClient(object):
+    """A general file client to access files in different backend.
+
+    The client loads a file or text in a specified backend from its path
+    and return it as a binary file. it can also register other backend
+    accessor with a given name and backend class.
+
+    Attributes:
+        backend (str): The storage backend type. Options are "disk",
+            "memcached" and "lmdb".
+        client (:obj:`BaseStorageBackend`): The backend object.
+    """
+
+    _backends = {
+        'disk': HardDiskBackend,
+        'memcached': MemcachedBackend,
+        'lmdb': LmdbBackend,
+    }
+
+    def __init__(self, backend='disk', **kwargs):
+        if backend not in self._backends:
+            raise ValueError(f'Backend {backend} is not supported. Currently supported ones'
+                             f' are {list(self._backends.keys())}')
+        self.backend = backend
+        self.client = self._backends[backend](**kwargs)
+
+    def get(self, filepath, client_key='default'):
+        # client_key is used only for lmdb, where different fileclients have
+        # different lmdb environments.
+        if self.backend == 'lmdb':
+            return self.client.get(filepath, client_key)
+        else:
+            return self.client.get(filepath)
+
+    def get_text(self, filepath):
+        return self.client.get_text(filepath)

official/bsr/utils/flow_util.py

@@ -0,0 +1,170 @@
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/video/optflow.py  # noqa: E501
+import cv2
+import numpy as np
+import os
+
+
+def flowread(flow_path, quantize=False, concat_axis=0, *args, **kwargs):
+    """Read an optical flow map.
+
+    Args:
+        flow_path (ndarray or str): Flow path.
+        quantize (bool): whether to read quantized pair, if set to True,
+            remaining args will be passed to :func:`dequantize_flow`.
+        concat_axis (int): The axis that dx and dy are concatenated,
+            can be either 0 or 1. Ignored if quantize is False.
+
+    Returns:
+        ndarray: Optical flow represented as a (h, w, 2) numpy array
+    """
+    if quantize:
+        assert concat_axis in [0, 1]
+        cat_flow = cv2.imread(flow_path, cv2.IMREAD_UNCHANGED)
+        if cat_flow.ndim != 2:
+            raise IOError(f'{flow_path} is not a valid quantized flow file, its dimension is {cat_flow.ndim}.')
+        assert cat_flow.shape[concat_axis] % 2 == 0
+        dx, dy = np.split(cat_flow, 2, axis=concat_axis)
+        flow = dequantize_flow(dx, dy, *args, **kwargs)
+    else:
+        with open(flow_path, 'rb') as f:
+            try:
+                header = f.read(4).decode('utf-8')
+            except Exception:
+                raise IOError(f'Invalid flow file: {flow_path}')
+            else:
+                if header != 'PIEH':
+                    raise IOError(f'Invalid flow file: {flow_path}, header does not contain PIEH')
+
+            w = np.fromfile(f, np.int32, 1).squeeze()
+            h = np.fromfile(f, np.int32, 1).squeeze()
+            flow = np.fromfile(f, np.float32, w * h * 2).reshape((h, w, 2))
+
+    return flow.astype(np.float32)
+
+
+def flowwrite(flow, filename, quantize=False, concat_axis=0, *args, **kwargs):
+    """Write optical flow to file.
+
+    If the flow is not quantized, it will be saved as a .flo file losslessly,
+    otherwise a jpeg image which is lossy but of much smaller size. (dx and dy
+    will be concatenated horizontally into a single image if quantize is True.)
+
+    Args:
+        flow (ndarray): (h, w, 2) array of optical flow.
+        filename (str): Output filepath.
+        quantize (bool): Whether to quantize the flow and save it to 2 jpeg
+            images. If set to True, remaining args will be passed to
+            :func:`quantize_flow`.
+        concat_axis (int): The axis that dx and dy are concatenated,
+            can be either 0 or 1. Ignored if quantize is False.
+    """
+    if not quantize:
+        with open(filename, 'wb') as f:
+            f.write('PIEH'.encode('utf-8'))
+            np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f)
+            flow = flow.astype(np.float32)
+            flow.tofile(f)
+            f.flush()
+    else:
+        assert concat_axis in [0, 1]
+        dx, dy = quantize_flow(flow, *args, **kwargs)
+        dxdy = np.concatenate((dx, dy), axis=concat_axis)
+        os.makedirs(os.path.dirname(filename), exist_ok=True)
+        cv2.imwrite(filename, dxdy)
+
+
+def quantize_flow(flow, max_val=0.02, norm=True):
+    """Quantize flow to [0, 255].
+
+    After this step, the size of flow will be much smaller, and can be
+    dumped as jpeg images.
+
+    Args:
+        flow (ndarray): (h, w, 2) array of optical flow.
+        max_val (float): Maximum value of flow, values beyond
+                        [-max_val, max_val] will be truncated.
+        norm (bool): Whether to divide flow values by image width/height.
+
+    Returns:
+        tuple[ndarray]: Quantized dx and dy.
+    """
+    h, w, _ = flow.shape
+    dx = flow[..., 0]
+    dy = flow[..., 1]
+    if norm:
+        dx = dx / w  # avoid inplace operations
+        dy = dy / h
+    # use 255 levels instead of 256 to make sure 0 is 0 after dequantization.
+    flow_comps = [quantize(d, -max_val, max_val, 255, np.uint8) for d in [dx, dy]]
+    return tuple(flow_comps)
+
+
+def dequantize_flow(dx, dy, max_val=0.02, denorm=True):
+    """Recover from quantized flow.
+
+    Args:
+        dx (ndarray): Quantized dx.
+        dy (ndarray): Quantized dy.
+        max_val (float): Maximum value used when quantizing.
+        denorm (bool): Whether to multiply flow values with width/height.
+
+    Returns:
+        ndarray: Dequantized flow.
+    """
+    assert dx.shape == dy.shape
+    assert dx.ndim == 2 or (dx.ndim == 3 and dx.shape[-1] == 1)
+
+    dx, dy = [dequantize(d, -max_val, max_val, 255) for d in [dx, dy]]
+
+    if denorm:
+        dx *= dx.shape[1]
+        dy *= dx.shape[0]
+    flow = np.dstack((dx, dy))
+    return flow
+
+
+def quantize(arr, min_val, max_val, levels, dtype=np.int64):
+    """Quantize an array of (-inf, inf) to [0, levels-1].
+
+    Args:
+        arr (ndarray): Input array.
+        min_val (scalar): Minimum value to be clipped.
+        max_val (scalar): Maximum value to be clipped.
+        levels (int): Quantization levels.
+        dtype (np.type): The type of the quantized array.
+
+    Returns:
+        tuple: Quantized array.
+    """
+    if not (isinstance(levels, int) and levels > 1):
+        raise ValueError(f'levels must be a positive integer, but got {levels}')
+    if min_val >= max_val:
+        raise ValueError(f'min_val ({min_val}) must be smaller than max_val ({max_val})')
+
+    arr = np.clip(arr, min_val, max_val) - min_val
+    quantized_arr = np.minimum(np.floor(levels * arr / (max_val - min_val)).astype(dtype), levels - 1)
+
+    return quantized_arr
+
+
+def dequantize(arr, min_val, max_val, levels, dtype=np.float64):
+    """Dequantize an array.
+
+    Args:
+        arr (ndarray): Input array.
+        min_val (scalar): Minimum value to be clipped.
+        max_val (scalar): Maximum value to be clipped.
+        levels (int): Quantization levels.
+        dtype (np.type): The type of the dequantized array.
+
+    Returns:
+        tuple: Dequantized array.
+    """
+    if not (isinstance(levels, int) and levels > 1):
+        raise ValueError(f'levels must be a positive integer, but got {levels}')
+    if min_val >= max_val:
+        raise ValueError(f'min_val ({min_val}) must be smaller than max_val ({max_val})')
+
+    dequantized_arr = (arr + 0.5).astype(dtype) * (max_val - min_val) / levels + min_val
+
+    return dequantized_arr

official/bsr/utils/img_process_util.py

@@ -0,0 +1,83 @@
+import cv2
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+
+def filter2D(img, kernel):
+    """PyTorch version of cv2.filter2D
+
+    Args:
+        img (Tensor): (b, c, h, w)
+        kernel (Tensor): (b, k, k)
+    """
+    k = kernel.size(-1)
+    b, c, h, w = img.size()
+    if k % 2 == 1:
+        img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect')
+    else:
+        raise ValueError('Wrong kernel size')
+
+    ph, pw = img.size()[-2:]
+
+    if kernel.size(0) == 1:
+        # apply the same kernel to all batch images
+        img = img.view(b * c, 1, ph, pw)
+        kernel = kernel.view(1, 1, k, k)
+        return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
+    else:
+        img = img.view(1, b * c, ph, pw)
+        kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
+        return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
+
+
+def usm_sharp(img, weight=0.5, radius=50, threshold=10):
+    """USM sharpening.
+
+    Input image: I; Blurry image: B.
+    1. sharp = I + weight * (I - B)
+    2. Mask = 1 if abs(I - B) > threshold, else: 0
+    3. Blur mask:
+    4. Out = Mask * sharp + (1 - Mask) * I
+
+
+    Args:
+        img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
+        weight (float): Sharp weight. Default: 1.
+        radius (float): Kernel size of Gaussian blur. Default: 50.
+        threshold (int):
+    """
+    if radius % 2 == 0:
+        radius += 1
+    blur = cv2.GaussianBlur(img, (radius, radius), 0)
+    residual = img - blur
+    mask = np.abs(residual) * 255 > threshold
+    mask = mask.astype('float32')
+    soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
+
+    sharp = img + weight * residual
+    sharp = np.clip(sharp, 0, 1)
+    return soft_mask * sharp + (1 - soft_mask) * img
+
+
+class USMSharp(torch.nn.Module):
+
+    def __init__(self, radius=50, sigma=0):
+        super(USMSharp, self).__init__()
+        if radius % 2 == 0:
+            radius += 1
+        self.radius = radius
+        kernel = cv2.getGaussianKernel(radius, sigma)
+        kernel = torch.FloatTensor(np.dot(kernel, kernel.transpose())).unsqueeze_(0)
+        self.register_buffer('kernel', kernel)
+
+    def forward(self, img, weight=0.5, threshold=10):
+        blur = filter2D(img, self.kernel)
+        residual = img - blur
+
+        mask = torch.abs(residual) * 255 > threshold
+        mask = mask.float()
+        soft_mask = filter2D(mask, self.kernel)
+        sharp = img + weight * residual
+        sharp = torch.clip(sharp, 0, 1)
+        return soft_mask * sharp + (1 - soft_mask) * img

official/bsr/utils/img_util.py

@@ -0,0 +1,172 @@
+import cv2
+import math
+import numpy as np
+import os
+import torch
+from torchvision.utils import make_grid
+
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)
+
+
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    """Convert torch Tensors into image numpy arrays.
+
+    After clamping to [min, max], values will be normalized to [0, 1].
+
+    Args:
+        tensor (Tensor or list[Tensor]): Accept shapes:
+            1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
+            2) 3D Tensor of shape (3/1 x H x W);
+            3) 2D Tensor of shape (H x W).
+            Tensor channel should be in RGB order.
+        rgb2bgr (bool): Whether to change rgb to bgr.
+        out_type (numpy type): output types. If ``np.uint8``, transform outputs
+            to uint8 type with range [0, 255]; otherwise, float type with
+            range [0, 1]. Default: ``np.uint8``.
+        min_max (tuple[int]): min and max values for clamp.
+
+    Returns:
+        (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
+        shape (H x W). The channel order is BGR.
+    """
+    if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
+
+    if torch.is_tensor(tensor):
+        tensor = [tensor]
+    result = []
+    for _tensor in tensor:
+        _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
+        _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
+
+        n_dim = _tensor.dim()
+        if n_dim == 4:
+            img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if rgb2bgr:
+                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 3:
+            img_np = _tensor.numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if img_np.shape[2] == 1:  # gray image
+                img_np = np.squeeze(img_np, axis=2)
+            else:
+                if rgb2bgr:
+                    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 2:
+            img_np = _tensor.numpy()
+        else:
+            raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
+        if out_type == np.uint8:
+            # Unlike MATLAB, numpy.unit8() WILL NOT round by default.
+            img_np = (img_np * 255.0).round()
+        img_np = img_np.astype(out_type)
+        result.append(img_np)
+    if len(result) == 1:
+        result = result[0]
+    return result
+
+
+def tensor2img_fast(tensor, rgb2bgr=True, min_max=(0, 1)):
+    """This implementation is slightly faster than tensor2img.
+    It now only supports torch tensor with shape (1, c, h, w).
+
+    Args:
+        tensor (Tensor): Now only support torch tensor with (1, c, h, w).
+        rgb2bgr (bool): Whether to change rgb to bgr. Default: True.
+        min_max (tuple[int]): min and max values for clamp.
+    """
+    output = tensor.squeeze(0).detach().clamp_(*min_max).permute(1, 2, 0)
+    output = (output - min_max[0]) / (min_max[1] - min_max[0]) * 255
+    output = output.type(torch.uint8).cpu().numpy()
+    if rgb2bgr:
+        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+    return output
+
+
+def imfrombytes(content, flag='color', float32=False):
+    """Read an image from bytes.
+
+    Args:
+        content (bytes): Image bytes got from files or other streams.
+        flag (str): Flags specifying the color type of a loaded image,
+            candidates are `color`, `grayscale` and `unchanged`.
+        float32 (bool): Whether to change to float32., If True, will also norm
+            to [0, 1]. Default: False.
+
+    Returns:
+        ndarray: Loaded image array.
+    """
+    img_np = np.frombuffer(content, np.uint8)
+    imread_flags = {'color': cv2.IMREAD_COLOR, 'grayscale': cv2.IMREAD_GRAYSCALE, 'unchanged': cv2.IMREAD_UNCHANGED}
+    img = cv2.imdecode(img_np, imread_flags[flag])
+    if float32:
+        img = img.astype(np.float32) / 255.
+    return img
+
+
+def imwrite(img, file_path, params=None, auto_mkdir=True):
+    """Write image to file.
+
+    Args:
+        img (ndarray): Image array to be written.
+        file_path (str): Image file path.
+        params (None or list): Same as opencv's :func:`imwrite` interface.
+        auto_mkdir (bool): If the parent folder of `file_path` does not exist,
+            whether to create it automatically.
+
+    Returns:
+        bool: Successful or not.
+    """
+    if auto_mkdir:
+        dir_name = os.path.abspath(os.path.dirname(file_path))
+        os.makedirs(dir_name, exist_ok=True)
+    ok = cv2.imwrite(file_path, img, params)
+    if not ok:
+        raise IOError('Failed in writing images.')
+
+
+def crop_border(imgs, crop_border):
+    """Crop borders of images.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Images with shape (h, w, c).
+        crop_border (int): Crop border for each end of height and weight.
+
+    Returns:
+        list[ndarray]: Cropped images.
+    """
+    if crop_border == 0:
+        return imgs
+    else:
+        if isinstance(imgs, list):
+            return [v[crop_border:-crop_border, crop_border:-crop_border, ...] for v in imgs]
+        else:
+            return imgs[crop_border:-crop_border, crop_border:-crop_border, ...]

official/bsr/utils/lmdb_util.py

@@ -0,0 +1,199 @@
+import cv2
+import lmdb
+import sys
+from multiprocessing import Pool
+from os import path as osp
+from tqdm import tqdm
+
+
+def make_lmdb_from_imgs(data_path,
+                        lmdb_path,
+                        img_path_list,
+                        keys,
+                        batch=5000,
+                        compress_level=1,
+                        multiprocessing_read=False,
+                        n_thread=40,
+                        map_size=None):
+    """Make lmdb from images.
+
+    Contents of lmdb. The file structure is:
+
+    ::
+
+        example.lmdb
+        ├── data.mdb
+        ├── lock.mdb
+        ├── meta_info.txt
+
+    The data.mdb and lock.mdb are standard lmdb files and you can refer to
+    https://lmdb.readthedocs.io/en/release/ for more details.
+
+    The meta_info.txt is a specified txt file to record the meta information
+    of our datasets. It will be automatically created when preparing
+    datasets by our provided dataset tools.
+    Each line in the txt file records 1)image name (with extension),
+    2)image shape, and 3)compression level, separated by a white space.
+
+    For example, the meta information could be:
+    `000_00000000.png (720,1280,3) 1`, which means:
+    1) image name (with extension): 000_00000000.png;
+    2) image shape: (720,1280,3);
+    3) compression level: 1
+
+    We use the image name without extension as the lmdb key.
+
+    If `multiprocessing_read` is True, it will read all the images to memory
+    using multiprocessing. Thus, your server needs to have enough memory.
+
+    Args:
+        data_path (str): Data path for reading images.
+        lmdb_path (str): Lmdb save path.
+        img_path_list (str): Image path list.
+        keys (str): Used for lmdb keys.
+        batch (int): After processing batch images, lmdb commits.
+            Default: 5000.
+        compress_level (int): Compress level when encoding images. Default: 1.
+        multiprocessing_read (bool): Whether use multiprocessing to read all
+            the images to memory. Default: False.
+        n_thread (int): For multiprocessing.
+        map_size (int | None): Map size for lmdb env. If None, use the
+            estimated size from images. Default: None
+    """
+
+    assert len(img_path_list) == len(keys), ('img_path_list and keys should have the same length, '
+                                             f'but got {len(img_path_list)} and {len(keys)}')
+    print(f'Create lmdb for {data_path}, save to {lmdb_path}...')
+    print(f'Totoal images: {len(img_path_list)}')
+    if not lmdb_path.endswith('.lmdb'):
+        raise ValueError("lmdb_path must end with '.lmdb'.")
+    if osp.exists(lmdb_path):
+        print(f'Folder {lmdb_path} already exists. Exit.')
+        sys.exit(1)
+
+    if multiprocessing_read:
+        # read all the images to memory (multiprocessing)
+        dataset = {}  # use dict to keep the order for multiprocessing
+        shapes = {}
+        print(f'Read images with multiprocessing, #thread: {n_thread} ...')
+        pbar = tqdm(total=len(img_path_list), unit='image')
+
+        def callback(arg):
+            """get the image data and update pbar."""
+            key, dataset[key], shapes[key] = arg
+            pbar.update(1)
+            pbar.set_description(f'Read {key}')
+
+        pool = Pool(n_thread)
+        for path, key in zip(img_path_list, keys):
+            pool.apply_async(read_img_worker, args=(osp.join(data_path, path), key, compress_level), callback=callback)
+        pool.close()
+        pool.join()
+        pbar.close()
+        print(f'Finish reading {len(img_path_list)} images.')
+
+    # create lmdb environment
+    if map_size is None:
+        # obtain data size for one image
+        img = cv2.imread(osp.join(data_path, img_path_list[0]), cv2.IMREAD_UNCHANGED)
+        _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
+        data_size_per_img = img_byte.nbytes
+        print('Data size per image is: ', data_size_per_img)
+        data_size = data_size_per_img * len(img_path_list)
+        map_size = data_size * 10
+
+    env = lmdb.open(lmdb_path, map_size=map_size)
+
+    # write data to lmdb
+    pbar = tqdm(total=len(img_path_list), unit='chunk')
+    txn = env.begin(write=True)
+    txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
+    for idx, (path, key) in enumerate(zip(img_path_list, keys)):
+        pbar.update(1)
+        pbar.set_description(f'Write {key}')
+        key_byte = key.encode('ascii')
+        if multiprocessing_read:
+            img_byte = dataset[key]
+            h, w, c = shapes[key]
+        else:
+            _, img_byte, img_shape = read_img_worker(osp.join(data_path, path), key, compress_level)
+            h, w, c = img_shape
+
+        txn.put(key_byte, img_byte)
+        # write meta information
+        txt_file.write(f'{key}.png ({h},{w},{c}) {compress_level}\n')
+        if idx % batch == 0:
+            txn.commit()
+            txn = env.begin(write=True)
+    pbar.close()
+    txn.commit()
+    env.close()
+    txt_file.close()
+    print('\nFinish writing lmdb.')
+
+
+def read_img_worker(path, key, compress_level):
+    """Read image worker.
+
+    Args:
+        path (str): Image path.
+        key (str): Image key.
+        compress_level (int): Compress level when encoding images.
+
+    Returns:
+        str: Image key.
+        byte: Image byte.
+        tuple[int]: Image shape.
+    """
+
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+    if img.ndim == 2:
+        h, w = img.shape
+        c = 1
+    else:
+        h, w, c = img.shape
+    _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
+    return (key, img_byte, (h, w, c))
+
+
+class LmdbMaker():
+    """LMDB Maker.
+
+    Args:
+        lmdb_path (str): Lmdb save path.
+        map_size (int): Map size for lmdb env. Default: 1024 ** 4, 1TB.
+        batch (int): After processing batch images, lmdb commits.
+            Default: 5000.
+        compress_level (int): Compress level when encoding images. Default: 1.
+    """
+
+    def __init__(self, lmdb_path, map_size=1024**4, batch=5000, compress_level=1):
+        if not lmdb_path.endswith('.lmdb'):
+            raise ValueError("lmdb_path must end with '.lmdb'.")
+        if osp.exists(lmdb_path):
+            print(f'Folder {lmdb_path} already exists. Exit.')
+            sys.exit(1)
+
+        self.lmdb_path = lmdb_path
+        self.batch = batch
+        self.compress_level = compress_level
+        self.env = lmdb.open(lmdb_path, map_size=map_size)
+        self.txn = self.env.begin(write=True)
+        self.txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
+        self.counter = 0
+
+    def put(self, img_byte, key, img_shape):
+        self.counter += 1
+        key_byte = key.encode('ascii')
+        self.txn.put(key_byte, img_byte)
+        # write meta information
+        h, w, c = img_shape
+        self.txt_file.write(f'{key}.png ({h},{w},{c}) {self.compress_level}\n')
+        if self.counter % self.batch == 0:
+            self.txn.commit()
+            self.txn = self.env.begin(write=True)
+
+    def close(self):
+        self.txn.commit()
+        self.env.close()
+        self.txt_file.close()

official/bsr/utils/logger.py

@@ -0,0 +1,213 @@
+import datetime
+import logging
+import time
+
+from .dist_util import get_dist_info, master_only
+
+initialized_logger = {}
+
+
+class AvgTimer():
+
+    def __init__(self, window=200):
+        self.window = window  # average window
+        self.current_time = 0
+        self.total_time = 0
+        self.count = 0
+        self.avg_time = 0
+        self.start()
+
+    def start(self):
+        self.start_time = self.tic = time.time()
+
+    def record(self):
+        self.count += 1
+        self.toc = time.time()
+        self.current_time = self.toc - self.tic
+        self.total_time += self.current_time
+        # calculate average time
+        self.avg_time = self.total_time / self.count
+
+        # reset
+        if self.count > self.window:
+            self.count = 0
+            self.total_time = 0
+
+        self.tic = time.time()
+
+    def get_current_time(self):
+        return self.current_time
+
+    def get_avg_time(self):
+        return self.avg_time
+
+
+class MessageLogger():
+    """Message logger for printing.
+
+    Args:
+        opt (dict): Config. It contains the following keys:
+            name (str): Exp name.
+            logger (dict): Contains 'print_freq' (str) for logger interval.
+            train (dict): Contains 'total_iter' (int) for total iters.
+            use_tb_logger (bool): Use tensorboard logger.
+        start_iter (int): Start iter. Default: 1.
+        tb_logger (obj:`tb_logger`): Tensorboard logger. Default： None.
+    """
+
+    def __init__(self, opt, start_iter=1, tb_logger=None):
+        self.exp_name = opt['name']
+        self.interval = opt['logger']['print_freq']
+        self.start_iter = start_iter
+        self.max_iters = opt['train']['total_iter']
+        self.use_tb_logger = opt['logger']['use_tb_logger']
+        self.tb_logger = tb_logger
+        self.start_time = time.time()
+        self.logger = get_root_logger()
+
+    def reset_start_time(self):
+        self.start_time = time.time()
+
+    @master_only
+    def __call__(self, log_vars):
+        """Format logging message.
+
+        Args:
+            log_vars (dict): It contains the following keys:
+                epoch (int): Epoch number.
+                iter (int): Current iter.
+                lrs (list): List for learning rates.
+
+                time (float): Iter time.
+                data_time (float): Data time for each iter.
+        """
+        # epoch, iter, learning rates
+        epoch = log_vars.pop('epoch')
+        current_iter = log_vars.pop('iter')
+        lrs = log_vars.pop('lrs')
+
+        message = (f'[{self.exp_name[:5]}..][epoch:{epoch:3d}, iter:{current_iter:8,d}, lr:(')
+        for v in lrs:
+            message += f'{v:.3e},'
+        message += ')] '
+
+        # time and estimated time
+        if 'time' in log_vars.keys():
+            iter_time = log_vars.pop('time')
+            data_time = log_vars.pop('data_time')
+
+            total_time = time.time() - self.start_time
+            time_sec_avg = total_time / (current_iter - self.start_iter + 1)
+            eta_sec = time_sec_avg * (self.max_iters - current_iter - 1)
+            eta_str = str(datetime.timedelta(seconds=int(eta_sec)))
+            message += f'[eta: {eta_str}, '
+            message += f'time (data): {iter_time:.3f} ({data_time:.3f})] '
+
+        # other items, especially losses
+        for k, v in log_vars.items():
+            message += f'{k}: {v:.4e} '
+            # tensorboard logger
+            if self.use_tb_logger and 'debug' not in self.exp_name:
+                if k.startswith('l_'):
+                    self.tb_logger.add_scalar(f'losses/{k}', v, current_iter)
+                else:
+                    self.tb_logger.add_scalar(k, v, current_iter)
+        self.logger.info(message)
+
+
+@master_only
+def init_tb_logger(log_dir):
+    from torch.utils.tensorboard import SummaryWriter
+    tb_logger = SummaryWriter(log_dir=log_dir)
+    return tb_logger
+
+
+@master_only
+def init_wandb_logger(opt):
+    """We now only use wandb to sync tensorboard log."""
+    import wandb
+    logger = get_root_logger()
+
+    project = opt['logger']['wandb']['project']
+    resume_id = opt['logger']['wandb'].get('resume_id')
+    if resume_id:
+        wandb_id = resume_id
+        resume = 'allow'
+        logger.warning(f'Resume wandb logger with id={wandb_id}.')
+    else:
+        wandb_id = wandb.util.generate_id()
+        resume = 'never'
+
+    wandb.init(id=wandb_id, resume=resume, name=opt['name'], config=opt, project=project, sync_tensorboard=True)
+
+    logger.info(f'Use wandb logger with id={wandb_id}; project={project}.')
+
+
+def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
+    """Get the root logger.
+
+    The logger will be initialized if it has not been initialized. By default a
+    StreamHandler will be added. If `log_file` is specified, a FileHandler will
+    also be added.
+
+    Args:
+        logger_name (str): root logger name. Default: 'basicsr'.
+        log_file (str | None): The log filename. If specified, a FileHandler
+            will be added to the root logger.
+        log_level (int): The root logger level. Note that only the process of
+            rank 0 is affected, while other processes will set the level to
+            "Error" and be silent most of the time.
+
+    Returns:
+        logging.Logger: The root logger.
+    """
+    logger = logging.getLogger(logger_name)
+    # if the logger has been initialized, just return it
+    if logger_name in initialized_logger:
+        return logger
+
+    format_str = '%(asctime)s %(levelname)s: %(message)s'
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(logging.Formatter(format_str))
+    logger.addHandler(stream_handler)
+    logger.propagate = False
+    rank, _ = get_dist_info()
+    if rank != 0:
+        logger.setLevel('ERROR')
+    elif log_file is not None:
+        logger.setLevel(log_level)
+        # add file handler
+        file_handler = logging.FileHandler(log_file, 'w')
+        file_handler.setFormatter(logging.Formatter(format_str))
+        file_handler.setLevel(log_level)
+        logger.addHandler(file_handler)
+    initialized_logger[logger_name] = True
+    return logger
+
+
+def get_env_info():
+    """Get environment information.
+
+    Currently, only log the software version.
+    """
+    import torch
+    import torchvision
+
+    from basicsr.version import __version__
+    msg = r"""
+                ____                _       _____  ____
+               / __ ) ____ _ _____ (_)_____/ ___/ / __ \
+              / __  |/ __ `// ___// // ___/\__ \ / /_/ /
+             / /_/ // /_/ /(__  )/ // /__ ___/ // _, _/
+            /_____/ \__,_//____//_/ \___//____//_/ |_|
+     ______                   __   __                 __      __
+    / ____/____   ____   ____/ /  / /   __  __ _____ / /__   / /
+   / / __ / __ \ / __ \ / __  /  / /   / / / // ___// //_/  / /
+  / /_/ // /_/ // /_/ // /_/ /  / /___/ /_/ // /__ / /<    /_/
+  \____/ \____/ \____/ \____/  /_____/\____/ \___//_/|_|  (_)
+    """
+    msg += ('\nVersion Information: '
+            f'\n\tBasicSR: {__version__}'
+            f'\n\tPyTorch: {torch.__version__}'
+            f'\n\tTorchVision: {torchvision.__version__}')
+    return msg

official/bsr/utils/matlab_functions.py

@@ -0,0 +1,178 @@
+import math
+import numpy as np
+import torch
+
+
+def cubic(x):
+    """cubic function used for calculate_weights_indices."""
+    absx = torch.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    return (1.5 * absx3 - 2.5 * absx2 + 1) * (
+        (absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * (((absx > 1) *
+                                                                                     (absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+    """Calculate weights and indices, used for imresize function.
+
+    Args:
+        in_length (int): Input length.
+        out_length (int): Output length.
+        scale (float): Scale factor.
+        kernel_width (int): Kernel width.
+        antialisaing (bool): Whether to apply anti-aliasing when downsampling.
+    """
+
+    if (scale < 1) and antialiasing:
+        # Use a modified kernel (larger kernel width) to simultaneously
+        # interpolate and antialias
+        kernel_width = kernel_width / scale
+
+    # Output-space coordinates
+    x = torch.linspace(1, out_length, out_length)
+
+    # Input-space coordinates. Calculate the inverse mapping such that 0.5
+    # in output space maps to 0.5 in input space, and 0.5 + scale in output
+    # space maps to 1.5 in input space.
+    u = x / scale + 0.5 * (1 - 1 / scale)
+
+    # What is the left-most pixel that can be involved in the computation?
+    left = torch.floor(u - kernel_width / 2)
+
+    # What is the maximum number of pixels that can be involved in the
+    # computation?  Note: it's OK to use an extra pixel here; if the
+    # corresponding weights are all zero, it will be eliminated at the end
+    # of this function.
+    p = math.ceil(kernel_width) + 2
+
+    # The indices of the input pixels involved in computing the k-th output
+    # pixel are in row k of the indices matrix.
+    indices = left.view(out_length, 1).expand(out_length, p) + torch.linspace(0, p - 1, p).view(1, p).expand(
+        out_length, p)
+
+    # The weights used to compute the k-th output pixel are in row k of the
+    # weights matrix.
+    distance_to_center = u.view(out_length, 1).expand(out_length, p) - indices
+
+    # apply cubic kernel
+    if (scale < 1) and antialiasing:
+        weights = scale * cubic(distance_to_center * scale)
+    else:
+        weights = cubic(distance_to_center)
+
+    # Normalize the weights matrix so that each row sums to 1.
+    weights_sum = torch.sum(weights, 1).view(out_length, 1)
+    weights = weights / weights_sum.expand(out_length, p)
+
+    # If a column in weights is all zero, get rid of it. only consider the
+    # first and last column.
+    weights_zero_tmp = torch.sum((weights == 0), 0)
+    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 1, p - 2)
+        weights = weights.narrow(1, 1, p - 2)
+    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 0, p - 2)
+        weights = weights.narrow(1, 0, p - 2)
+    weights = weights.contiguous()
+    indices = indices.contiguous()
+    sym_len_s = -indices.min() + 1
+    sym_len_e = indices.max() - in_length
+    indices = indices + sym_len_s - 1
+    return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+@torch.no_grad()
+def imresize(img, scale, antialiasing=True):
+    """imresize function same as MATLAB.
+
+    It now only supports bicubic.
+    The same scale applies for both height and width.
+
+    Args:
+        img (Tensor | Numpy array):
+            Tensor: Input image with shape (c, h, w), [0, 1] range.
+            Numpy: Input image with shape (h, w, c), [0, 1] range.
+        scale (float): Scale factor. The same scale applies for both height
+            and width.
+        antialisaing (bool): Whether to apply anti-aliasing when downsampling.
+            Default: True.
+
+    Returns:
+        Tensor: Output image with shape (c, h, w), [0, 1] range, w/o round.
+    """
+    squeeze_flag = False
+    if type(img).__module__ == np.__name__:  # numpy type
+        numpy_type = True
+        if img.ndim == 2:
+            img = img[:, :, None]
+            squeeze_flag = True
+        img = torch.from_numpy(img.transpose(2, 0, 1)).float()
+    else:
+        numpy_type = False
+        if img.ndim == 2:
+            img = img.unsqueeze(0)
+            squeeze_flag = True
+
+    in_c, in_h, in_w = img.size()
+    out_h, out_w = math.ceil(in_h * scale), math.ceil(in_w * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # get weights and indices
+    weights_h, indices_h, sym_len_hs, sym_len_he = calculate_weights_indices(in_h, out_h, scale, kernel, kernel_width,
+                                                                             antialiasing)
+    weights_w, indices_w, sym_len_ws, sym_len_we = calculate_weights_indices(in_w, out_w, scale, kernel, kernel_width,
+                                                                             antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_c, in_h + sym_len_hs + sym_len_he, in_w)
+    img_aug.narrow(1, sym_len_hs, in_h).copy_(img)
+
+    sym_patch = img[:, :sym_len_hs, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, 0, sym_len_hs).copy_(sym_patch_inv)
+
+    sym_patch = img[:, -sym_len_he:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, sym_len_hs + in_h, sym_len_he).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(in_c, out_h, in_w)
+    kernel_width = weights_h.size(1)
+    for i in range(out_h):
+        idx = int(indices_h[i][0])
+        for j in range(in_c):
+            out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_h[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(in_c, out_h, in_w + sym_len_ws + sym_len_we)
+    out_1_aug.narrow(2, sym_len_ws, in_w).copy_(out_1)
+
+    sym_patch = out_1[:, :, :sym_len_ws]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, 0, sym_len_ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, :, -sym_len_we:]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, sym_len_ws + in_w, sym_len_we).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(in_c, out_h, out_w)
+    kernel_width = weights_w.size(1)
+    for i in range(out_w):
+        idx = int(indices_w[i][0])
+        for j in range(in_c):
+            out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_w[i])
+
+    if squeeze_flag:
+        out_2 = out_2.squeeze(0)
+    if numpy_type:
+        out_2 = out_2.numpy()
+        if not squeeze_flag:
+            out_2 = out_2.transpose(1, 2, 0)
+
+    return out_2

official/bsr/utils/misc.py

@@ -0,0 +1,141 @@
+import numpy as np
+import os
+import random
+import time
+import torch
+from os import path as osp
+
+from .dist_util import master_only
+
+
+def set_random_seed(seed):
+    """Set random seeds."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def get_time_str():
+    return time.strftime('%Y%m%d_%H%M%S', time.localtime())
+
+
+def mkdir_and_rename(path):
+    """mkdirs. If path exists, rename it with timestamp and create a new one.
+
+    Args:
+        path (str): Folder path.
+    """
+    if osp.exists(path):
+        new_name = path + '_archived_' + get_time_str()
+        print(f'Path already exists. Rename it to {new_name}', flush=True)
+        os.rename(path, new_name)
+    os.makedirs(path, exist_ok=True)
+
+
+@master_only
+def make_exp_dirs(opt):
+    """Make dirs for experiments."""
+    path_opt = opt['path'].copy()
+    if opt['is_train']:
+        mkdir_and_rename(path_opt.pop('experiments_root'))
+    else:
+        mkdir_and_rename(path_opt.pop('results_root'))
+    for key, path in path_opt.items():
+        if ('strict_load' in key) or ('pretrain_network' in key) or ('resume' in key) or ('param_key' in key):
+            continue
+        else:
+            os.makedirs(path, exist_ok=True)
+
+
+def scandir(dir_path, suffix=None, recursive=False, full_path=False):
+    """Scan a directory to find the interested files.
+
+    Args:
+        dir_path (str): Path of the directory.
+        suffix (str | tuple(str), optional): File suffix that we are
+            interested in. Default: None.
+        recursive (bool, optional): If set to True, recursively scan the
+            directory. Default: False.
+        full_path (bool, optional): If set to True, include the dir_path.
+            Default: False.
+
+    Returns:
+        A generator for all the interested files with relative paths.
+    """
+
+    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
+
+    root = dir_path
+
+    def _scandir(dir_path, suffix, recursive):
+        for entry in os.scandir(dir_path):
+            if not entry.name.startswith('.') and entry.is_file():
+                if full_path:
+                    return_path = entry.path
+                else:
+                    return_path = osp.relpath(entry.path, root)
+
+                if suffix is None:
+                    yield return_path
+                elif return_path.endswith(suffix):
+                    yield return_path
+            else:
+                if recursive:
+                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
+                else:
+                    continue
+
+    return _scandir(dir_path, suffix=suffix, recursive=recursive)
+
+
+def check_resume(opt, resume_iter):
+    """Check resume states and pretrain_network paths.
+
+    Args:
+        opt (dict): Options.
+        resume_iter (int): Resume iteration.
+    """
+    if opt['path']['resume_state']:
+        # get all the networks
+        networks = [key for key in opt.keys() if key.startswith('network_')]
+        flag_pretrain = False
+        for network in networks:
+            if opt['path'].get(f'pretrain_{network}') is not None:
+                flag_pretrain = True
+        if flag_pretrain:
+            print('pretrain_network path will be ignored during resuming.')
+        # set pretrained model paths
+        for network in networks:
+            name = f'pretrain_{network}'
+            basename = network.replace('network_', '')
+            if opt['path'].get('ignore_resume_networks') is None or (network
+                                                                     not in opt['path']['ignore_resume_networks']):
+                opt['path'][name] = osp.join(opt['path']['models'], f'net_{basename}_{resume_iter}.pth')
+                print(f"Set {name} to {opt['path'][name]}")
+
+        # change param_key to params in resume
+        param_keys = [key for key in opt['path'].keys() if key.startswith('param_key')]
+        for param_key in param_keys:
+            if opt['path'][param_key] == 'params_ema':
+                opt['path'][param_key] = 'params'
+                print(f'Set {param_key} to params')
+
+
+def sizeof_fmt(size, suffix='B'):
+    """Get human readable file size.
+
+    Args:
+        size (int): File size.
+        suffix (str): Suffix. Default: 'B'.
+
+    Return:
+        str: Formatted file size.
+    """
+    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
+        if abs(size) < 1024.0:
+            return f'{size:3.1f} {unit}{suffix}'
+        size /= 1024.0
+    return f'{size:3.1f} Y{suffix}'

official/bsr/utils/options.py

@@ -0,0 +1,218 @@
+import argparse
+import os
+import random
+import torch
+import yaml
+from collections import OrderedDict
+from os import path as osp
+
+from bsr.utils import set_random_seed
+from bsr.utils.dist_util import get_dist_info, init_dist, master_only
+
+
+def ordered_yaml():
+    """Support OrderedDict for yaml.
+
+    Returns:
+        tuple: yaml Loader and Dumper.
+    """
+    try:
+        from yaml import CDumper as Dumper
+        from yaml import CLoader as Loader
+    except ImportError:
+        from yaml import Dumper, Loader
+
+    _mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG
+
+    def dict_representer(dumper, data):
+        return dumper.represent_dict(data.items())
+
+    def dict_constructor(loader, node):
+        return OrderedDict(loader.construct_pairs(node))
+
+    Dumper.add_representer(OrderedDict, dict_representer)
+    Loader.add_constructor(_mapping_tag, dict_constructor)
+    return Loader, Dumper
+
+
+def yaml_load(f):
+    """Load yaml file or string.
+
+    Args:
+        f (str): File path or a python string.
+
+    Returns:
+        dict: Loaded dict.
+    """
+    if os.path.isfile(f):
+        with open(f, 'r') as f:
+            return yaml.load(f, Loader=ordered_yaml()[0])
+    else:
+        return yaml.load(f, Loader=ordered_yaml()[0])
+
+
+def dict2str(opt, indent_level=1):
+    """dict to string for printing options.
+
+    Args:
+        opt (dict): Option dict.
+        indent_level (int): Indent level. Default: 1.
+
+    Return:
+        (str): Option string for printing.
+    """
+    msg = '\n'
+    for k, v in opt.items():
+        if isinstance(v, dict):
+            msg += ' ' * (indent_level * 2) + k + ':['
+            msg += dict2str(v, indent_level + 1)
+            msg += ' ' * (indent_level * 2) + ']\n'
+        else:
+            msg += ' ' * (indent_level * 2) + k + ': ' + str(v) + '\n'
+    return msg
+
+
+def _postprocess_yml_value(value):
+    # None
+    if value == '~' or value.lower() == 'none':
+        return None
+    # bool
+    if value.lower() == 'true':
+        return True
+    elif value.lower() == 'false':
+        return False
+    # !!float number
+    if value.startswith('!!float'):
+        return float(value.replace('!!float', ''))
+    # number
+    if value.isdigit():
+        return int(value)
+    elif value.replace('.', '', 1).isdigit() and value.count('.') < 2:
+        return float(value)
+    # list
+    if value.startswith('['):
+        return eval(value)
+    # str
+    return value
+
+
+def parse_options(root_path, is_train=True):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-opt', type=str, required=True, help='Path to option YAML file.')
+    parser.add_argument('--launcher', choices=['none', 'pytorch', 'slurm'], default='none', help='job launcher')
+    parser.add_argument('--auto_resume', action='store_true')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--local_rank', type=int, default=0)
+    parser.add_argument(
+        '--force_yml', nargs='+', default=None, help='Force to update yml files. Examples: train:ema_decay=0.999')
+    args = parser.parse_args()
+
+    # parse yml to dict
+    opt = yaml_load(args.opt)
+
+    # distributed settings
+    if args.launcher == 'none':
+        opt['dist'] = False
+        print('Disable distributed.', flush=True)
+    else:
+        opt['dist'] = True
+        if args.launcher == 'slurm' and 'dist_params' in opt:
+            init_dist(args.launcher, **opt['dist_params'])
+        else:
+            init_dist(args.launcher)
+    opt['rank'], opt['world_size'] = get_dist_info()
+
+    # random seed
+    seed = opt.get('manual_seed')
+    if seed is None:
+        seed = random.randint(1, 10000)
+        opt['manual_seed'] = seed
+    set_random_seed(seed + opt['rank'])
+
+    # force to update yml options
+    if args.force_yml is not None:
+        for entry in args.force_yml:
+            # now do not support creating new keys
+            keys, value = entry.split('=')
+            keys, value = keys.strip(), value.strip()
+            value = _postprocess_yml_value(value)
+            eval_str = 'opt'
+            for key in keys.split(':'):
+                eval_str += f'["{key}"]'
+            eval_str += '=value'
+            # using exec function
+            exec(eval_str)
+
+    opt['auto_resume'] = args.auto_resume
+    opt['is_train'] = is_train
+
+    # debug setting
+    if args.debug and not opt['name'].startswith('debug'):
+        opt['name'] = 'debug_' + opt['name']
+
+    if opt['num_gpu'] == 'auto':
+        opt['num_gpu'] = torch.cuda.device_count()
+
+    # datasets
+    for phase, dataset in opt['datasets'].items():
+        # for multiple datasets, e.g., val_1, val_2; test_1, test_2
+        phase = phase.split('_')[0]
+        dataset['phase'] = phase
+        if 'scale' in opt:
+            dataset['scale'] = opt['scale']
+        if dataset.get('dataroot_gt') is not None:
+            dataset['dataroot_gt'] = osp.expanduser(dataset['dataroot_gt'])
+        if dataset.get('dataroot_lq') is not None:
+            dataset['dataroot_lq'] = osp.expanduser(dataset['dataroot_lq'])
+
+    # paths
+    for key, val in opt['path'].items():
+        if (val is not None) and ('resume_state' in key or 'pretrain_network' in key):
+            opt['path'][key] = osp.expanduser(val)
+
+    if is_train:
+        experiments_root = opt['path'].get('experiments_root')
+        if experiments_root is None:
+            experiments_root = osp.join(root_path, 'experiments')
+        experiments_root = osp.join(experiments_root, opt['name'])
+
+        opt['path']['experiments_root'] = experiments_root
+        opt['path']['models'] = osp.join(experiments_root, 'models')
+        opt['path']['training_states'] = osp.join(experiments_root, 'training_states')
+        opt['path']['log'] = experiments_root
+        opt['path']['visualization'] = osp.join(experiments_root, 'visualization')
+
+        # change some options for debug mode
+        if 'debug' in opt['name']:
+            if 'val' in opt:
+                opt['val']['val_freq'] = 8
+            opt['logger']['print_freq'] = 1
+            opt['logger']['save_checkpoint_freq'] = 8
+    else:  # test
+        results_root = opt['path'].get('results_root')
+        if results_root is None:
+            results_root = osp.join(root_path, 'results')
+        results_root = osp.join(results_root, opt['name'])
+
+        opt['path']['results_root'] = results_root
+        opt['path']['log'] = results_root
+        opt['path']['visualization'] = osp.join(results_root, 'visualization')
+
+    return opt, args
+
+
+@master_only
+def copy_opt_file(opt_file, experiments_root):
+    # copy the yml file to the experiment root
+    import sys
+    import time
+    from shutil import copyfile
+    cmd = ' '.join(sys.argv)
+    filename = osp.join(experiments_root, osp.basename(opt_file))
+    copyfile(opt_file, filename)
+
+    with open(filename, 'r+') as f:
+        lines = f.readlines()
+        lines.insert(0, f'# GENERATE TIME: {time.asctime()}\n# CMD:\n# {cmd}\n\n')
+        f.seek(0)
+        f.writelines(lines)

official/bsr/utils/plot_util.py

@@ -0,0 +1,83 @@
+import re
+
+
+def read_data_from_tensorboard(log_path, tag):
+    """Get raw data (steps and values) from tensorboard events.
+
+    Args:
+        log_path (str): Path to the tensorboard log.
+        tag (str): tag to be read.
+    """
+    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
+
+    # tensorboard event
+    event_acc = EventAccumulator(log_path)
+    event_acc.Reload()
+    scalar_list = event_acc.Tags()['scalars']
+    print('tag list: ', scalar_list)
+    steps = [int(s.step) for s in event_acc.Scalars(tag)]
+    values = [s.value for s in event_acc.Scalars(tag)]
+    return steps, values
+
+
+def read_data_from_txt_2v(path, pattern, step_one=False):
+    """Read data from txt with 2 returned values (usually [step, value]).
+
+    Args:
+        path (str): path to the txt file.
+        pattern (str): re (regular expression) pattern.
+        step_one (bool): add 1 to steps. Default: False.
+    """
+    with open(path) as f:
+        lines = f.readlines()
+    lines = [line.strip() for line in lines]
+    steps = []
+    values = []
+
+    pattern = re.compile(pattern)
+    for line in lines:
+        match = pattern.match(line)
+        if match:
+            steps.append(int(match.group(1)))
+            values.append(float(match.group(2)))
+    if step_one:
+        steps = [v + 1 for v in steps]
+    return steps, values
+
+
+def read_data_from_txt_1v(path, pattern):
+    """Read data from txt with 1 returned values.
+
+    Args:
+        path (str): path to the txt file.
+        pattern (str): re (regular expression) pattern.
+    """
+    with open(path) as f:
+        lines = f.readlines()
+    lines = [line.strip() for line in lines]
+    data = []
+
+    pattern = re.compile(pattern)
+    for line in lines:
+        match = pattern.match(line)
+        if match:
+            data.append(float(match.group(1)))
+    return data
+
+
+def smooth_data(values, smooth_weight):
+    """ Smooth data using 1st-order IIR low-pass filter (what tensorflow does).
+
+    Reference: https://github.com/tensorflow/tensorboard/blob/f801ebf1f9fbfe2baee1ddd65714d0bccc640fb1/tensorboard/plugins/scalar/vz_line_chart/vz-line-chart.ts#L704  # noqa: E501
+
+    Args:
+        values (list): A list of values to be smoothed.
+        smooth_weight (float): Smooth weight.
+    """
+    values_sm = []
+    last_sm_value = values[0]
+    for value in values:
+        value_sm = last_sm_value * smooth_weight + (1 - smooth_weight) * value
+        values_sm.append(value_sm)
+        last_sm_value = value_sm
+    return values_sm

official/bsr/utils/registry.py

@@ -0,0 +1,88 @@
+# Modified from: https://github.com/facebookresearch/fvcore/blob/master/fvcore/common/registry.py  # noqa: E501
+
+
+class Registry():
+    """
+    The registry that provides name -> object mapping, to support third-party
+    users' custom modules.
+
+    To create a registry (e.g. a backbone registry):
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY = Registry('BACKBONE')
+
+    To register an object:
+
+    .. code-block:: python
+
+        @BACKBONE_REGISTRY.register()
+        class MyBackbone():
+            ...
+
+    Or:
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY.register(MyBackbone)
+    """
+
+    def __init__(self, name):
+        """
+        Args:
+            name (str): the name of this registry
+        """
+        self._name = name
+        self._obj_map = {}
+
+    def _do_register(self, name, obj, suffix=None):
+        if isinstance(suffix, str):
+            name = name + '_' + suffix
+
+        assert (name not in self._obj_map), (f"An object named '{name}' was already registered "
+                                             f"in '{self._name}' registry!")
+        self._obj_map[name] = obj
+
+    def register(self, obj=None, suffix=None):
+        """
+        Register the given object under the the name `obj.__name__`.
+        Can be used as either a decorator or not.
+        See docstring of this class for usage.
+        """
+        if obj is None:
+            # used as a decorator
+            def deco(func_or_class):
+                name = func_or_class.__name__
+                self._do_register(name, func_or_class, suffix)
+                return func_or_class
+
+            return deco
+
+        # used as a function call
+        name = obj.__name__
+        self._do_register(name, obj, suffix)
+
+    def get(self, name, suffix='basicsr'):
+        ret = self._obj_map.get(name)
+        if ret is None:
+            ret = self._obj_map.get(name + '_' + suffix)
+            print(f'Name {name} is not found, use name: {name}_{suffix}!')
+        if ret is None:
+            raise KeyError(f"No object named '{name}' found in '{self._name}' registry!")
+        return ret
+
+    def __contains__(self, name):
+        return name in self._obj_map
+
+    def __iter__(self):
+        return iter(self._obj_map.items())
+
+    def keys(self):
+        return self._obj_map.keys()
+
+
+DATASET_REGISTRY = Registry('dataset')
+ARCH_REGISTRY = Registry('arch')
+MODEL_REGISTRY = Registry('model')
+LOSS_REGISTRY = Registry('loss')
+METRIC_REGISTRY = Registry('metric')

official/dataset.py

@@ -0,0 +1,280 @@
+import torch, random, cv2, os, math, glob
+import torch.nn.functional as F
+import numpy as np
+from bsr.degradations import circular_lowpass_kernel, random_mixed_kernels, random_add_gaussian_noise_pt, random_add_poisson_noise_pt
+from bsr.transforms import augment, paired_random_crop
+from bsr.utils import FileClient, imfrombytes, img2tensor, DiffJPEG
+from bsr.utils.img_process_util import filter2D
+
+class RealESRGANDataset(torch.utils.data.Dataset):
+    def __init__(self, opt, bsz):
+        super(RealESRGANDataset, self).__init__()
+        self.opt = opt
+        self.file_client = FileClient("disk")
+        self.gt_folder = opt["dataroot_gt"]
+        self.len = bsz * opt["iter_num"]
+        self.paths = glob.glob(os.path.join(self.gt_folder, "**/*"), recursive=True)
+
+        # blur settings for the first degradation
+        self.blur_kernel_size = opt["blur_kernel_size"]
+        self.kernel_list = opt["kernel_list"]
+        self.kernel_prob = opt["kernel_prob"]  # a list for each kernel probability
+        self.blur_sigma = opt["blur_sigma"]
+        self.betag_range = opt["betag_range"]  # betag used in generalized Gaussian blur kernels
+        self.betap_range = opt["betap_range"]  # betap used in plateau blur kernels
+        self.sinc_prob = opt["sinc_prob"]  # the probability for sinc filters
+
+        # blur settings for the second degradation
+        self.blur_kernel_size2 = opt["blur_kernel_size2"]
+        self.kernel_list2 = opt["kernel_list2"]
+        self.kernel_prob2 = opt["kernel_prob2"]
+        self.blur_sigma2 = opt["blur_sigma2"]
+        self.betag_range2 = opt["betag_range2"]
+        self.betap_range2 = opt["betap_range2"]
+        self.sinc_prob2 = opt["sinc_prob2"]
+
+        # a final sinc filter
+        self.final_sinc_prob = opt["final_sinc_prob"]
+
+        self.kernel_range = [2 * v + 1 for v in range(3, 11)]  # kernel size ranges from 7 to 21
+        # TODO: kernel range is now hard-coded, should be in the configure file
+        self.pulse_tensor = torch.zeros(21, 21).float()  # convolving with pulse tensor brings no blurry effect
+        self.pulse_tensor[10, 10] = 1
+
+    def __getitem__(self, index):
+        index = random.randint(0, len(self.paths) - 1)
+        gt_path = self.paths[index]
+        img_gt = imfrombytes(self.file_client.get(gt_path, "gt"), float32=True)
+        img_gt = augment(img_gt, self.opt["use_hflip"], self.opt["use_rot"])
+        h, w = img_gt.shape[0:2]
+        crop_pad_size = self.opt.gt_size
+        if h < crop_pad_size or w < crop_pad_size:
+            pad_h = max(0, crop_pad_size - h)
+            pad_w = max(0, crop_pad_size - w)
+            img_gt = cv2.copyMakeBorder(img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
+        if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size:
+            h, w = img_gt.shape[0:2]
+            top = random.randint(0, h - crop_pad_size)
+            left = random.randint(0, w - crop_pad_size)
+            img_gt = img_gt[top:top + crop_pad_size, left:left + crop_pad_size, ...]
+
+        # ------------------------ Generate kernels (used in the first degradation) ------------------------ #
+        kernel_size = random.choice(self.kernel_range)
+        if np.random.uniform() < self.opt["sinc_prob"]:
+            # this sinc filter setting is for kernels ranging from [7, 21]
+            if kernel_size < 13:
+                omega_c = np.random.uniform(np.pi / 3, np.pi)
+            else:
+                omega_c = np.random.uniform(np.pi / 5, np.pi)
+            kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
+        else:
+            kernel = random_mixed_kernels(
+                self.kernel_list,
+                self.kernel_prob,
+                kernel_size,
+                self.blur_sigma,
+                self.blur_sigma, [-math.pi, math.pi],
+                self.betag_range,
+                self.betap_range,
+                noise_range=None)
+        # pad kernel
+        pad_size = (21 - kernel_size) // 2
+        kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
+
+        # ------------------------ Generate kernels (used in the second degradation) ------------------------ #
+        kernel_size = random.choice(self.kernel_range)
+        if np.random.uniform() < self.opt["sinc_prob2"]:
+            if kernel_size < 13:
+                omega_c = np.random.uniform(np.pi / 3, np.pi)
+            else:
+                omega_c = np.random.uniform(np.pi / 5, np.pi)
+            kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
+        else:
+            kernel2 = random_mixed_kernels(
+                self.kernel_list2,
+                self.kernel_prob2,
+                kernel_size,
+                self.blur_sigma2,
+                self.blur_sigma2, [-math.pi, math.pi],
+                self.betag_range2,
+                self.betap_range2,
+                noise_range=None)
+
+        # pad kernel
+        pad_size = (21 - kernel_size) // 2
+        kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
+
+        # ------------------------------------- the final sinc kernel ------------------------------------- #
+        if np.random.uniform() < self.opt["final_sinc_prob"]:
+            kernel_size = random.choice(self.kernel_range)
+            omega_c = np.random.uniform(np.pi / 3, np.pi)
+            sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
+            sinc_kernel = torch.FloatTensor(sinc_kernel)
+        else:
+            sinc_kernel = self.pulse_tensor
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
+        kernel = torch.FloatTensor(kernel)
+        kernel2 = torch.FloatTensor(kernel2)
+
+        return_d = {"gt": img_gt, "kernel1": kernel, "kernel2": kernel2, "sinc_kernel": sinc_kernel, "gt_path": gt_path}
+        return return_d
+
+    def __len__(self):
+        return self.len
+
+class RealESRGANDegrader:
+    def __init__(self, opt, device):
+        self.opt = opt
+        self.device = device
+        self.jpeger = DiffJPEG(differentiable=False).to(device)  # simulate JPEG compression artifacts
+        self.queue_size = 1200
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self):
+        """It is the training pair pool for increasing the diversity in a batch.
+
+        Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
+        batch could not have different resize scaling factors. Therefore, we employ this training pair pool
+        to increase the degradation diversity in a batch.
+        """
+        # initialize
+        b, c, h, w = self.lq.size()
+        if not hasattr(self, "queue_lr"):
+            assert self.queue_size % b == 0, f"queue size {self.queue_size} should be divisible by batch size {b}"
+            self.queue_lr = torch.zeros(self.queue_size, c, h, w).to(self.device)
+            _, c, h, w = self.gt.size()
+            self.queue_gt = torch.zeros(self.queue_size, c, h, w).to(self.device)
+            self.queue_ptr = 0
+        if self.queue_ptr == self.queue_size:  # the pool is full
+            # do dequeue and enqueue
+            # shuffle
+            idx = torch.randperm(self.queue_size)
+            self.queue_lr = self.queue_lr[idx]
+            self.queue_gt = self.queue_gt[idx]
+            # get first b samples
+            lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
+            gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
+            # update the queue
+            self.queue_lr[0:b, :, :, :] = self.lq.clone()
+            self.queue_gt[0:b, :, :, :] = self.gt.clone()
+
+            self.lq = lq_dequeue
+            self.gt = gt_dequeue
+        else:
+            # only do enqueue
+            self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
+            self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
+            self.queue_ptr = self.queue_ptr + b
+
+    @torch.no_grad()
+    def degrade(self, data):
+        """Accept data from dataloader, and then add two-order degradations to obtain LQ images.
+        """
+        # training data synthesis
+        self.gt = data["gt"].to(self.device)
+
+        self.kernel1 = data["kernel1"].to(self.device)
+        self.kernel2 = data["kernel2"].to(self.device)
+        self.sinc_kernel = data["sinc_kernel"].to(self.device)
+
+        ori_h, ori_w = self.gt.size()[2:4]
+
+        # ----------------------- The first degradation process ----------------------- #
+        # blur
+        out = filter2D(self.gt, self.kernel1)
+        # random resize
+        updown_type = random.choices(["up", "down", "keep"], self.opt["resize_prob"])[0]
+        if updown_type == "up":
+            scale = np.random.uniform(1, self.opt["resize_range"][1])
+        elif updown_type == "down":
+            scale = np.random.uniform(self.opt["resize_range"][0], 1)
+        else:
+            scale = 1
+        mode = random.choice(["area", "bilinear", "bicubic"])
+        out = F.interpolate(out, scale_factor=scale, mode=mode)
+        # add noise
+        gray_noise_prob = self.opt["gray_noise_prob"]
+        if np.random.uniform() < self.opt["gaussian_noise_prob"]:
+            out = random_add_gaussian_noise_pt(
+                out, sigma_range=self.opt["noise_range"], clip=True, rounds=False, gray_prob=gray_noise_prob)
+        else:
+            out = random_add_poisson_noise_pt(
+                out,
+                scale_range=self.opt["poisson_scale_range"],
+                gray_prob=gray_noise_prob,
+                clip=True,
+                rounds=False)
+        # JPEG compression
+        jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt["jpeg_range"])
+        out = torch.clamp(out, 0, 1)  # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
+        out = self.jpeger(out, quality=jpeg_p)
+
+        # ----------------------- The second degradation process ----------------------- #
+        # blur
+        if np.random.uniform() < self.opt["second_blur_prob"]:
+            out = filter2D(out, self.kernel2)
+        # random resize
+        updown_type = random.choices(["up", "down", "keep"], self.opt["resize_prob2"])[0]
+        if updown_type == "up":
+            scale = np.random.uniform(1, self.opt["resize_range2"][1])
+        elif updown_type == "down":
+            scale = np.random.uniform(self.opt["resize_range2"][0], 1)
+        else:
+            scale = 1
+        mode = random.choice(["area", "bilinear", "bicubic"])
+        out = F.interpolate(
+            out, size=(int(ori_h / self.opt["scale"] * scale), int(ori_w / self.opt["scale"] * scale)), mode=mode)
+        # add noise
+        gray_noise_prob = self.opt["gray_noise_prob2"]
+        if np.random.uniform() < self.opt["gaussian_noise_prob2"]:
+            out = random_add_gaussian_noise_pt(
+                out, sigma_range=self.opt["noise_range2"], clip=True, rounds=False, gray_prob=gray_noise_prob)
+        else:
+            out = random_add_poisson_noise_pt(
+                out,
+                scale_range=self.opt["poisson_scale_range2"],
+                gray_prob=gray_noise_prob,
+                clip=True,
+                rounds=False)
+
+        # JPEG compression + the final sinc filter
+        # We also need to resize images to desired sizes. We group [resize back + sinc filter] together
+        # as one operation.
+        # We consider two orders:
+        #   1. [resize back + sinc filter] + JPEG compression
+        #   2. JPEG compression + [resize back + sinc filter]
+        # Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
+        if np.random.uniform() < 0.5:
+            # resize back + the final sinc filter
+            mode = random.choice(["area", "bilinear", "bicubic"])
+            out = F.interpolate(out, size=(ori_h // self.opt["scale"], ori_w // self.opt["scale"]), mode=mode)
+            out = filter2D(out, self.sinc_kernel)
+            # JPEG compression
+            jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt["jpeg_range2"])
+            out = torch.clamp(out, 0, 1)
+            out = self.jpeger(out, quality=jpeg_p)
+        else:
+            # JPEG compression
+            jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt["jpeg_range2"])
+            out = torch.clamp(out, 0, 1)
+            out = self.jpeger(out, quality=jpeg_p)
+            # resize back + the final sinc filter
+            mode = random.choice(["area", "bilinear", "bicubic"])
+            out = F.interpolate(out, size=(ori_h // self.opt["scale"], ori_w // self.opt["scale"]), mode=mode)
+            out = filter2D(out, self.sinc_kernel)
+
+        # clamp and round
+        self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
+
+        # random crop
+        gt_size = self.opt["gt_size"]
+        self.gt, self.lq = paired_random_crop(self.gt, self.lq, gt_size, self.opt["scale"])
+
+        # training pair pool
+        self._dequeue_and_enqueue()
+        # sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
+        self.lq = self.lq.contiguous()  # for the warning: grad and param do not obey the gradient layout contract
+
+        return self.lq, self.gt

official/evaluate.py

@@ -0,0 +1,55 @@
+import torch, os, glob, pyiqa
+from argparse import ArgumentParser
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from torchvision import transforms
+
+parser = ArgumentParser()
+parser.add_argument("--HR_dir", type=str, default="testset/RealSR/HR")
+parser.add_argument("--SR_dir", type=str, default="result/RealSR")
+args = parser.parse_args()
+
+device = torch.device("cuda")
+
+psnr = pyiqa.create_metric("psnr", test_y_channel=True, color_space="ycbcr", device=device)
+ssim = pyiqa.create_metric("ssim", test_y_channel=True, color_space="ycbcr", device=device)
+lpips = pyiqa.create_metric("lpips", device=device)
+dists = pyiqa.create_metric("dists", device=device)
+fid = pyiqa.create_metric("fid", device=device)
+niqe = pyiqa.create_metric("niqe", device=device)
+maniqa = pyiqa.create_metric("maniqa-pipal", device=device)
+clipiqa = pyiqa.create_metric("clipiqa", device=device)
+musiq = pyiqa.create_metric("musiq", device=device)
+
+test_SR_paths = list(sorted(glob.glob(os.path.join(args.SR_dir, "*"))))
+test_HR_paths = list(sorted(glob.glob(os.path.join(args.HR_dir, "*"))))
+
+metrics = {"psnr": [], "ssim": [], "lpips": [], "dists": [], "niqe": [], "maniqa": [], "musiq": [], "clipiqa": []}
+
+for i, (SR_path, HR_path) in tqdm(enumerate(zip(test_SR_paths, test_HR_paths))):
+    SR = Image.open(SR_path).convert("RGB")
+    SR = transforms.ToTensor()(SR).to(device).unsqueeze(0)
+    HR = Image.open(HR_path).convert("RGB")
+    HR = transforms.ToTensor()(HR).to(device).unsqueeze(0)
+    metrics["psnr"].append(psnr(SR, HR).item())
+    metrics["ssim"].append(ssim(SR, HR).item())
+    metrics["lpips"].append(lpips(SR, HR).item())
+    metrics["dists"].append(dists(SR, HR).item())
+    metrics["niqe"].append(niqe(SR).item())
+    metrics["maniqa"].append(maniqa(SR).item())
+    metrics["clipiqa"].append(clipiqa(SR).item())
+    metrics["musiq"].append(musiq(SR).item())
+
+for k in metrics.keys():
+    metrics[k] = np.mean(metrics[k])
+
+metrics["fid"] = fid(args.SR_dir, args.HR_dir)
+
+for k, v in metrics.items():
+    if k == "niqe":
+        print(k, f"{v:.3g}")
+    elif k == "fid":
+        print(k, f"{v:.5g}")
+    else:
+        print(k, f"{v:.4g}")

official/forward.py

@@ -0,0 +1,67 @@
+import torch
+
+def MyUNet2DConditionModel_SD_forward(self, x):
+    global skip
+    x = self.conv_in(x)
+    skip = [x]
+    x = self.body(x)
+    return x
+
+def MyCrossAttnDownBlock2D_SD_forward(self, x):
+    for i in range(2):
+        x = self.resnets[i](x)
+        x = self.attentions[i](x)
+        skip.append(x)
+    if self.downsamplers is not None:
+        x = self.downsamplers[0](x)
+        skip.append(x)
+    return x
+
+def MyCrossAttnUpBlock2D_SD_forward(self, x):
+    for i in range(3):
+        x = self.resnets[i](torch.cat([x, skip.pop()], dim=1))
+        x = self.attentions[i](x)
+    if self.upsamplers is not None:
+        x = self.upsamplers[0](x)
+    return x
+
+def MyDownBlock2D_SD_forward(self, x):
+    for i in range(2):
+        x = self.resnets[i](x)
+        skip.append(x)
+    return x
+
+def MyUNetMidBlock2DCrossAttn_SD_forward(self, x):
+    x = self.resnets[0](x)
+    x = self.attentions[0](x)
+    x = self.resnets[1](x)
+    return x
+
+def MyUpBlock2D_SD_forward(self, x):
+    for i in range(3):
+        x = self.resnets[i](torch.cat([x, skip.pop()], dim=1))
+    x = self.upsamplers[0](x)
+    return x
+
+def MyResnetBlock2D_SD_forward(self, x_in):
+    x = self.norm1(x_in)
+    x = self.nonlinearity(x)
+    x = self.conv1(x)
+    x = self.norm2(x)
+    x = self.nonlinearity(x)
+    x = self.conv2(x)
+    if self.in_channels == self.out_channels:
+        return x + x_in
+    return x + self.conv_shortcut(x_in)
+
+def MyTransformer2DModel_SD_forward(self, x_in):
+    b, c, h, w = x_in.shape
+    x = self.norm(x_in)
+    x = x.permute(0, 2, 3, 1).reshape(b, h * w, c).contiguous()
+    x = self.proj_in(x)
+    for block in self.transformer_blocks:
+        x = x + block.attn1(block.norm1(x))
+        x = x + block.ff(block.norm3(x))
+    x = self.proj_out(x)
+    x = x.reshape(b, h, w, c).permute(0, 3, 1, 2).contiguous()
+    return x + x_in

official/model.py

@@ -0,0 +1,152 @@
+import torch, types, copy
+from torch import nn
+import torch.nn.functional as F
+from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D, \
+                                                  CrossAttnUpBlock2D, \
+                                                  DownBlock2D, \
+                                                  UpBlock2D, \
+                                                  UNetMidBlock2DCrossAttn
+from diffusers.models.resnet import ResnetBlock2D
+from diffusers.models.transformers.transformer_2d import Transformer2DModel
+from diffusers.models.attention import BasicTransformerBlock
+from diffusers.models.downsampling import Downsample2D
+from diffusers.models.upsampling import Upsample2D
+from forward import MyUNet2DConditionModel_SD_forward, \
+                    MyCrossAttnDownBlock2D_SD_forward, \
+                    MyDownBlock2D_SD_forward, \
+                    MyUNetMidBlock2DCrossAttn_SD_forward, \
+                    MyCrossAttnUpBlock2D_SD_forward, \
+                    MyUpBlock2D_SD_forward, \
+                    MyResnetBlock2D_SD_forward, \
+                    MyTransformer2DModel_SD_forward
+
+def find_parent(model, module_name):
+    components = module_name.split(".")
+    parent = model
+    for comp in components[:-1]:
+        parent = getattr(parent, comp)
+    return parent, components[-1]
+
+def halve_channels(model):
+    for name, module in model.named_modules():
+        if hasattr(module, "pruned"):
+            continue
+        if isinstance(module, nn.Conv2d):
+            in_channels = int(module.in_channels * 0.75)
+            out_channels = int(module.out_channels * 0.75)
+            new_conv = nn.Conv2d(in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=module.kernel_size,
+                                 stride=module.stride,
+                                 padding=module.padding,
+                                 dilation=module.dilation,
+                                 groups=module.groups,
+                                 bias=module.bias is not None)
+            with torch.no_grad():
+                new_conv.weight.copy_(module.weight[:out_channels, :in_channels])
+                if module.bias is not None:
+                    new_conv.bias.copy_(module.bias[:out_channels])
+            parent, last_name = find_parent(model, name)
+            setattr(parent, last_name, new_conv)
+            new_conv.pruned = True
+        elif isinstance(module, nn.Linear):
+            in_features = int(module.in_features * 0.75)
+            out_features = int(module.out_features * 0.75)
+            new_linear = nn.Linear(in_features=in_features,
+                                   out_features=out_features,
+                                   bias=module.bias is not None)
+            with torch.no_grad():
+                new_linear.weight.copy_(module.weight[:out_features, :in_features])
+                if module.bias is not None:
+                    new_linear.bias.copy_(module.bias[:out_features])
+            parent, last_name = find_parent(model, name)
+            setattr(parent, last_name, new_linear)
+            new_linear.pruned = True
+        elif isinstance(module, nn.GroupNorm):
+            num_channels = int(module.num_channels * 0.75)
+            for num_groups in [32, 24, 16, 12, 8, 6, 4, 2, 1]:
+                if num_channels % num_groups == 0:
+                    break
+            new_gn = nn.GroupNorm(num_groups=num_groups,
+                                  num_channels=num_channels,
+                                  eps=module.eps,
+                                  affine=module.affine)
+            with torch.no_grad():
+                new_gn.weight.copy_(module.weight[:num_channels])
+                new_gn.bias.copy_(module.bias[:num_channels])
+            parent, last_name = find_parent(model, name)
+            setattr(parent, last_name, new_gn)
+            new_gn.pruned = True
+        elif isinstance(module, nn.LayerNorm):
+            normalized_shape = int(module.normalized_shape[0] * 0.75)
+            new_ln = nn.LayerNorm(normalized_shape, 
+                                  eps=module.eps, 
+                                  elementwise_affine=module.elementwise_affine)
+            with torch.no_grad():
+                new_ln.weight.copy_(module.weight[:normalized_shape])
+                new_ln.bias.copy_(module.bias[:normalized_shape])
+            parent, last_name = find_parent(model, name)
+            setattr(parent, last_name, new_ln)
+            new_ln.pruned = True
+        elif isinstance(module, Downsample2D) or isinstance(module, Upsample2D):
+            module.channels = int(module.channels * 0.75)
+
+class Net(nn.Module):
+    def __init__(self, unet, decoder):
+        super().__init__()
+        del unet.time_embedding
+        new_conv_in = nn.Conv2d(16, 320, 3, padding=1)
+        new_conv_in.weight.data = unet.conv_in.weight.data.repeat(1, 4, 1, 1)
+        new_conv_in.bias.data = unet.conv_in.bias.data
+        unet.conv_in = new_conv_in
+        new_conv_out = nn.Conv2d(320, 342, 3, padding=1)
+        new_conv_out.weight.data = unet.conv_out.weight.data.repeat(86, 1, 1, 1)[:342]
+        new_conv_out.bias.data = unet.conv_out.bias.data.repeat(86,)[:342]
+        unet.conv_out = new_conv_out
+        def ResnetBlock2D_remove_time_emb_proj(module):
+            if isinstance(module, ResnetBlock2D):
+                del module.time_emb_proj
+        unet.apply(ResnetBlock2D_remove_time_emb_proj)
+        def BasicTransformerBlock_remove_cross_attn(module):
+            if isinstance(module, BasicTransformerBlock):
+                del module.attn2, module.norm2
+        unet.apply(BasicTransformerBlock_remove_cross_attn)
+        def set_inplace_to_true(module):
+            if isinstance(module, nn.Dropout) or isinstance(module, nn.SiLU):
+                module.inplace = True
+        unet.apply(set_inplace_to_true)
+        def replace_forward_methods(module):
+            if isinstance(module, CrossAttnDownBlock2D):
+                module.forward = types.MethodType(MyCrossAttnDownBlock2D_SD_forward, module)
+            elif isinstance(module, DownBlock2D):
+                module.forward = types.MethodType(MyDownBlock2D_SD_forward, module)
+            elif isinstance(module, UNetMidBlock2DCrossAttn):
+                module.forward = types.MethodType(MyUNetMidBlock2DCrossAttn_SD_forward, module)
+            elif isinstance(module, UpBlock2D):
+                module.forward = types.MethodType(MyUpBlock2D_SD_forward, module)
+            elif isinstance(module, CrossAttnUpBlock2D):
+                module.forward = types.MethodType(MyCrossAttnUpBlock2D_SD_forward, module)
+            elif isinstance(module, ResnetBlock2D):
+                module.forward = types.MethodType(MyResnetBlock2D_SD_forward, module)
+            elif isinstance(module, Transformer2DModel):
+                module.forward = types.MethodType(MyTransformer2DModel_SD_forward, module)
+        unet.apply(replace_forward_methods)
+        unet.forward = types.MethodType(MyUNet2DConditionModel_SD_forward, unet)
+        halve_channels(unet)
+        unet.body = nn.Sequential(
+            *unet.down_blocks,
+            unet.mid_block,
+            *unet.up_blocks,
+            unet.conv_norm_out,
+            unet.conv_act,
+            unet.conv_out,
+        )
+        del decoder.conv_in, decoder.up_blocks, decoder.conv_norm_out, decoder.conv_act, decoder.conv_out
+        self.body = nn.Sequential(
+            nn.PixelUnshuffle(2),
+            unet,
+            decoder.mid_block,
+        )
+    
+    def forward(self, x):
+        return self.body(x)

official/test.py

@@ -0,0 +1,70 @@
+import torch, os, glob, copy
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+from argparse import ArgumentParser
+from torchvision import transforms
+from model import Net
+
+parser = ArgumentParser()
+parser.add_argument("--epoch", type=int, default=200)
+parser.add_argument("--model_dir", type=str, default="weight")
+parser.add_argument("--LR_dir", type=str, default="testset/RealSR/LR")
+parser.add_argument("--HR_dir", type=str, default="testset/RealSR/HR")
+parser.add_argument("--SR_dir", type=str, default="result/RealSR")
+args = parser.parse_args()
+
+device = torch.device("cuda")
+
+from diffusers import StableDiffusionPipeline
+model_id = "stabilityai/stable-diffusion-2-1-base"
+pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)
+
+vae = pipe.vae
+tokenizer = pipe.tokenizer
+unet = pipe.unet
+noise_scheduler = pipe.scheduler
+text_encoder = pipe.text_encoder
+
+from diffusers.models.autoencoders.vae import Decoder 
+ckpt_halfdecoder = torch.load("./weight/pretrained/halfDecoder.ckpt", weights_only=False)
+decoder = Decoder(in_channels=4,
+            out_channels=3,
+            up_block_types=["UpDecoderBlock2D" for _ in range(4)],
+            block_out_channels=[64, 128, 256, 256],
+            layers_per_block=2, 
+            norm_num_groups=32, 
+            act_fn="silu", 
+            norm_type="group", 
+            mid_block_add_attention=True).to(device)
+decoder_ckpt = {}
+for k,v in ckpt_halfdecoder["state_dict"].items():
+    if "decoder" in k:
+        new_k = k.replace("decoder.", "")
+        decoder_ckpt[new_k] = v
+decoder.load_state_dict(decoder_ckpt, strict=True)
+
+model = torch.nn.DataParallel(Net(unet, copy.deepcopy(decoder)))
+model.load_state_dict(torch.load("./%s/net_params_%d.pkl" % (args.model_dir, args.epoch), weights_only=False))
+model = torch.nn.Sequential(
+    model.module,
+    *decoder.up_blocks,
+    decoder.conv_norm_out,
+    decoder.conv_act,
+    decoder.conv_out,
+).to(device)
+
+test_LR_paths = list(sorted(glob.glob(os.path.join(args.LR_dir, "*.png"))))
+test_HR_paths = list(sorted(glob.glob(os.path.join(args.HR_dir, "*.png"))))
+
+os.makedirs(args.SR_dir, exist_ok=True)
+
+with torch.no_grad():
+    for i, path in enumerate(test_LR_paths):
+        LR = Image.open(path).convert("RGB")
+        LR = transforms.ToTensor()(LR).to(device).unsqueeze(0) * 2 - 1
+        SR = model(LR)
+        SR = (SR - SR.mean(dim=[2,3],keepdim=True)) / SR.std(dim=[2,3],keepdim=True) \
+             * LR.std(dim=[2,3],keepdim=True) + LR.mean(dim=[2,3],keepdim=True)
+        SR = transforms.ToPILImage()((SR[0] / 2 + 0.5).clamp(0, 1).cpu())
+        SR.save(os.path.join(args.SR_dir, os.path.basename(path)))

official/train.py

@@ -0,0 +1,225 @@
+import torch, os, glob, random, copy
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+import numpy as np
+from argparse import ArgumentParser
+from time import time
+from tqdm import tqdm
+from omegaconf import OmegaConf
+from dataset import RealESRGANDataset, RealESRGANDegrader
+from model import Net
+from ram.models.ram_lora import ram
+from torchvision import transforms
+from utils import add_lora_to_unet
+
+dist.init_process_group(backend="nccl", init_method="env://")
+rank = dist.get_rank()
+world_size = dist.get_world_size()
+
+parser = ArgumentParser()
+parser.add_argument("--epoch", type=int, default=200)
+parser.add_argument("--batch_size", type=int, default=12)
+parser.add_argument("--learning_rate", type=float, default=1e-4)
+parser.add_argument("--model_dir", type=str, default="weight")
+parser.add_argument("--log_dir", type=str, default="log")
+parser.add_argument("--save_interval", type=int, default=10)
+
+args = parser.parse_args()
+
+# fixed seed for reproduction
+seed = rank
+random.seed(seed)
+np.random.seed(seed)
+torch.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+
+config = OmegaConf.load("config.yml")
+
+epoch = args.epoch
+learning_rate = args.learning_rate
+bsz = args.batch_size
+
+device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cuda.matmul.allow_tf32 = True
+
+if rank == 0:
+    print("batch size per gpu =", bsz)
+
+from diffusers import StableDiffusionPipeline
+model_id = "stabilityai/stable-diffusion-2-1-base"
+pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)
+
+vae = pipe.vae
+tokenizer = pipe.tokenizer
+unet = pipe.unet
+text_encoder = pipe.text_encoder
+
+unet_D = copy.deepcopy(unet)
+new_conv_in = torch.nn.Conv2d(256, 320, 3, padding=1).to(device)
+new_conv_in.weight.data = unet_D.conv_in.weight.data.repeat(1, 64, 1, 1) / 64
+new_conv_in.bias.data = unet_D.conv_in.bias.data
+unet_D.conv_in = new_conv_in
+unet_D = add_lora_to_unet(unet_D)
+unet_D.set_adapters(["default_encoder", "default_decoder", "default_others"])
+
+vae_teacher = copy.deepcopy(vae)
+unet_teacher = copy.deepcopy(unet)
+
+osediff = torch.load("./weight/pretrained/osediff.pkl", weights_only=False)
+vae_teacher.load_state_dict(osediff["vae"])
+unet_teacher.load_state_dict(osediff["unet"])
+
+from diffusers.models.autoencoders.vae import Decoder 
+ckpt_halfdecoder = torch.load("./weight/pretrained/halfDecoder.ckpt", weights_only=False)
+decoder = Decoder(in_channels=4,
+                  out_channels=3,
+                  up_block_types=["UpDecoderBlock2D" for _ in range(4)],
+                  block_out_channels=[64, 128, 256, 256],
+                  layers_per_block=2,
+                  norm_num_groups=32,
+                  act_fn="silu",
+                  norm_type="group",
+                  mid_block_add_attention=True).to(device)
+decoder_ckpt = {}
+for k, v in ckpt_halfdecoder["state_dict"].items():
+    if "decoder" in k:
+        new_k = k.replace("decoder.", "")
+        decoder_ckpt[new_k] = v
+decoder.load_state_dict(decoder_ckpt, strict=True)
+
+ram_transforms = transforms.Compose([
+    transforms.Resize((384, 384)),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+DAPE = ram(pretrained="./weight/pretrained/ram_swin_large_14m.pth",
+           pretrained_condition="./weight/pretrained/DAPE.pth",
+           image_size=384,
+           vit="swin_l").eval().to(device)
+
+vae.requires_grad_(False)
+unet.requires_grad_(False)
+text_encoder.requires_grad_(False)
+vae_teacher.requires_grad_(False)
+unet_teacher.requires_grad_(False)
+decoder.requires_grad_(False)
+DAPE.requires_grad_(False)
+
+model = DDP(Net(unet, copy.deepcopy(decoder)).to(device), device_ids=[rank])
+model_D = DDP(unet_D.to(device), device_ids=[rank])
+model.requires_grad_(True)
+model_D.requires_grad_(False)
+params_to_opt = []
+for n, p in model_D.named_parameters():
+    if "lora" in n or "conv_in" in n:
+        p.requires_grad = True
+        params_to_opt.append(p)
+
+if rank == 0:
+    param_cnt = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("#Param.", param_cnt/1e6, "M")
+
+dataset = RealESRGANDataset(config, bsz)
+degrader = RealESRGANDegrader(config, device)
+dataloader = DataLoader(dataset, batch_size=bsz, num_workers=8)
+optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+optimizer_D = torch.optim.Adam(params_to_opt, lr=1e-6)
+scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100,], gamma=0.5)
+scaler = torch.cuda.amp.GradScaler()
+
+model_dir = "./%s" % (args.model_dir,)
+log_path = "./%s/log.txt" % (args.log_dir,)
+os.makedirs(model_dir, exist_ok=True)
+os.makedirs(args.log_dir, exist_ok=True)
+
+print("start training...")
+timesteps = torch.tensor([999], device=device).long().expand(bsz,)
+alpha = pipe.scheduler.alphas_cumprod[999]
+for epoch_i in range(1, epoch + 1):
+    start_time = time()
+    loss_avg = 0.0
+    loss_distil_avg = 0.0
+    loss_adv_avg = 0.0
+    loss_D_avg = 0.0
+    iter_num = 0
+    dist.barrier()
+    for batch in tqdm(dataloader):
+        with torch.cuda.amp.autocast(enabled=True):
+            with torch.no_grad():
+                LR, HR = degrader.degrade(batch)
+                text_input = tokenizer(DAPE.generate_tag(ram_transforms(LR))[0],
+                                       max_length=tokenizer.model_max_length,
+                                       padding="max_length", return_tensors="pt").to(device)
+                encoder_hidden_states = text_encoder(text_input.input_ids, return_dict=False)[0]
+                LR, HR = LR * 2 - 1, HR * 2 - 1
+                LR_ = F.interpolate(LR, scale_factor=4, mode="bicubic")
+                LR_latents = vae_teacher.encode(LR_).latent_dist.mean * vae_teacher.config.scaling_factor
+                HR_latents = vae.encode(HR).latent_dist.mean
+                pred_teacher = unet_teacher(
+                    LR_latents,
+                    timesteps,
+                    encoder_hidden_states=encoder_hidden_states,
+                    return_dict=False,
+                )[0]
+                z0_teacher = (LR_latents-((1-alpha)**0.5)*pred_teacher)/(alpha**0.5)
+                z0_teacher = vae_teacher.post_quant_conv(z0_teacher / vae_teacher.config.scaling_factor)
+                z0_teacher = decoder.conv_in(z0_teacher)
+                z0_teacher = decoder.mid_block(z0_teacher)
+                z0_gt = vae.post_quant_conv(HR_latents)
+                z0_gt = decoder.conv_in(z0_gt)
+                z0_gt = decoder.mid_block(z0_gt)
+            z0_student = model(LR)
+            loss_distil = (z0_student - z0_teacher).abs().mean()
+            loss_adv = F.softplus(-model_D(
+                z0_student,
+                timesteps,
+                encoder_hidden_states=encoder_hidden_states,
+                return_dict=False,
+            )[0]).mean()
+            loss = loss_distil + loss_adv
+        optimizer.zero_grad(set_to_none=True)
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+        with torch.cuda.amp.autocast(enabled=True):
+            pred_real = model_D(
+                z0_gt.detach(),
+                timesteps,
+                encoder_hidden_states=encoder_hidden_states,
+                return_dict=False,
+            )[0]
+            pred_fake = model_D(
+                z0_student.detach(),
+                timesteps,
+                encoder_hidden_states=encoder_hidden_states,
+                return_dict=False,
+            )[0]
+            loss_D = F.softplus(pred_fake).mean() + F.softplus(-pred_real).mean()
+        optimizer_D.zero_grad(set_to_none=True)
+        scaler.scale(loss_D).backward()
+        scaler.step(optimizer_D)
+        scaler.update()
+        loss_avg += loss.item()
+        loss_distil_avg += loss_distil.item()
+        loss_adv_avg += loss_adv.item()
+        loss_D_avg += loss_D.item()
+        iter_num += 1
+        # print("loss", loss.item())
+        # print("loss_distil", loss_distil.item())
+        # print("loss_adv", loss_adv.item())
+        # print("loss_D", loss_D.item())
+    scheduler.step()
+    loss_avg /= iter_num
+    loss_distil_avg /= iter_num
+    loss_adv_avg /= iter_num
+    loss_D_avg /= iter_num
+    log_data = "[%d/%d] Average loss: %f, distil loss: %f, adv loss: %f, D loss: %f, time cost: %.2fs, cur lr is %f." % (epoch_i, epoch, loss_avg, loss_distil_avg, loss_adv_avg, loss_D_avg, time() - start_time, scheduler.get_last_lr()[0])
+    if rank == 0:
+        print(log_data)
+        with open(log_path, "a") as log_file:
+            log_file.write(log_data + "\n")
+        if epoch_i % args.save_interval == 0:
+            torch.save(model.state_dict(), "./%s/net_params_%d.pkl" % (model_dir, epoch_i))

official/utils.py

@@ -0,0 +1,24 @@
+import torch
+from peft import LoraConfig
+
+def add_lora_to_unet(unet, rank=4):
+    l_target_modules_encoder, l_target_modules_decoder, l_modules_others = [], [], []
+    l_grep = ["to_k", "to_q", "to_v", "to_out.0", "conv", "conv1", "conv2", "conv_shortcut", "conv_out", "proj_out", "proj_in", "ff.net.2", "ff.net.0.proj"]
+    for n, p in unet.named_parameters():
+        check_flag = 0
+        if "bias" in n or "norm" in n:
+            continue
+        for pattern in l_grep:
+            if pattern in n and ("down_blocks" in n or "conv_in" in n):
+                l_target_modules_encoder.append(n.replace(".weight",""))
+                break
+            elif pattern in n and ("up_blocks" in n or "conv_out" in n):
+                l_target_modules_decoder.append(n.replace(".weight",""))
+                break
+            elif pattern in n:
+                l_modules_others.append(n.replace(".weight",""))
+                break
+    unet.add_adapter(LoraConfig(r=rank,init_lora_weights="gaussian",target_modules=l_target_modules_encoder), adapter_name="default_encoder")
+    unet.add_adapter(LoraConfig(r=rank,init_lora_weights="gaussian",target_modules=l_target_modules_decoder), adapter_name="default_decoder")
+    unet.add_adapter(LoraConfig(r=rank,init_lora_weights="gaussian",target_modules=l_modules_others), adapter_name="default_others")
+    return unet

requirements.txt

@@ -0,0 +1,17 @@
+# A100 训练环境依赖 (Python 3.10/3.11, CUDA 12.x)
+torch>=2.1
+torchvision>=0.16
+diffusers>=0.27
+accelerate>=0.28
+omegaconf>=2.3
+einops
+lpips
+pyiqa
+opencv-python
+scipy
+numpy
+pillow
+safetensors
+huggingface_hub
+pyyaml
+requests

scripts/archive_run.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+"""归档: 把 logs/<run> 与权重副本整理到 archive/<date>_<name>/, 生成 summary.txt
+用法: python scripts/archive_run.py --name s2_best --src_log logs/s2 --weights weight/s2/net_params_*.pkl
+"""
+import argparse, os, shutil, json, glob, datetime
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--name", required=True)
+    ap.add_argument("--src_log", default="logs")
+    ap.add_argument("--weights", nargs="*", default=[])
+    ap.add_argument("--eval_json", default="", help="eval_val.py 输出")
+    args = ap.parse_args()
+    dst = os.path.join("archive", f"{datetime.date.today().isoformat()}_{args.name}")
+    os.makedirs(dst, exist_ok=True)
+    if os.path.isdir(args.src_log):
+        shutil.copytree(args.src_log, os.path.join(dst, "log"), dirs_exist_ok=True)
+    wd = os.path.join(dst, "weights"); os.makedirs(wd, exist_ok=True)
+    for pat in args.weights:
+        for f in glob.glob(pat):
+            shutil.copy2(f, os.path.join(wd, os.path.basename(f)))
+    summary = {"name": args.name, "date": datetime.date.today().isoformat(), "weights": [os.path.basename(f) for pat in args.weights for f in glob.glob(pat)]}
+    if args.eval_json and os.path.exists(args.eval_json):
+        with open(args.eval_json, encoding="utf-8") as fh:
+            summary["eval_mean"] = json.load(fh).get("mean")
+    with open(os.path.join(dst, "summary.txt"), "w", encoding="utf-8") as fh:
+        fh.write(json.dumps(summary, ensure_ascii=False, indent=1))
+    print("archived to", dst)
+
+if __name__ == "__main__":
+    main()

scripts/download_weights.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Plan A: 服务器有外网时下载权重(在项目根执行)
+# 双方案: 无外网时在本地跑同脚本并把 weight/ 与 models/ 传上来(见 pack_upload)
+set -euo pipefail
+cd "$(dirname "$0")/.."
+export HF_ENDPOINT=https://hf-mirror.com
+mkdir -p weight/pretrained models/stable-diffusion-2-1-base weight/gdpo
+HF=https://huggingface.co/Guaishou74851/AdcSR/resolve/main
+for f in weight/net_params_200.pkl weight/pretrained/DAPE.pth weight/pretrained/halfDecoder.ckpt \
+         weight/pretrained/osediff.pkl weight/pretrained/ram_swin_large_14m.pth; do
+  echo "==> $f"; curl -fL --retry 5 --retry-delay 3 -o "$f" "$HF/$f"
+done
+MS=https://modelscope.cn/models/AI-ModelScope/stable-diffusion-2-1-base/resolve/master
+for f in model_index.json feature_extractor/preprocessor_config.json scheduler/scheduler_config.json \
+         tokenizer/merges.txt tokenizer/special_tokens_map.json tokenizer/tokenizer_config.json \
+         tokenizer/vocab.json text_encoder/config.json text_encoder/model.fp16.safetensors \
+         unet/config.json unet/diffusion_pytorch_model.fp16.safetensors \
+         vae/config.json vae/diffusion_pytorch_model.fp16.safetensors; do
+  echo "==> $f"; curl -fL --retry 5 --retry-delay 3 --create-dirs -o "models/stable-diffusion-2-1-base/$f" "$MS/$f"
+done
+# GDPO 教师 (S2 前下载; 失败可 --teacher osediff)
+echo "==> GDPO (Joypop/GDPO)"
+python - <<'EOF'
+import os
+try:
+    from huggingface_hub import snapshot_download
+    snapshot_download(repo_id="Joypop/GDPO", repo_type="model", local_dir="weight/gdpo", allow_patterns=["ckp/*", "ram/*"])
+    print("GDPO downloaded")
+except Exception as e:
+    print("GDPO download failed (可回退 osediff):", e)
+EOF
+du -sh weight models

scripts/env_check.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+"""环境自检: GPU/CUDA/pyiqa/dwt/官方import 冒烟(不加载大权重)"""
+import os, sys
+sys.path.insert(0, "official"); sys.path.insert(0, "src")
+import torch
+print("torch", torch.__version__, "cuda", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("gpu", torch.cuda.get_device_name(0), "mem_GB", round(torch.cuda.get_device_properties(0).total_memory/1e9,1))
+try:
+    import pyiqa; print("pyiqa ok")
+except Exception as e: print("pyiqa FAIL", e)
+try:
+    import pywt; print("pywt ok")
+except Exception as e: print("pywt missing(可选)", e)
+try:
+    import diffusers, transformers, peft, omegaconf
+    print("diffusers/transformers/peft ok")
+except Exception as e: print("libs FAIL", e)
+# 官方模块可 import
+try:
+    import model, dataset, utils, forward
+    print("official imports ok")
+except Exception as e:
+    print("official imports FAIL", e)
+print("env_check done")

scripts/probe_gdpo.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+"""GDPO 权重 probe: 打印键结构, 供 src/common.load_gdpo_teacher 排错。
+用法: python scripts/probe_gdpo.py weight/gdpo/ckp/diffusion_pytorch_model.safetensors
+"""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+from common import probe_gdpo
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("用法: probe_gdpo.py <gdpo权重路径>"); sys.exit(1)
+    probe_gdpo(sys.argv[1])

scripts/run_smoke.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+# S0 冒烟(每个训练配置 1 迭代)
+set -euo pipefail
+cd "$(dirname "$0")/.."
+source "$(conda info --base)/etc/profile.d/conda.sh"; conda activate AdcSR
+CUDA_VISIBLE_DEVICES=0 python src/train_lora.py --config configs/config_smoke.yml \
+  --manifest data/manifest_train.json --real_prob 0.0 --steps 2 --batch_size 1 --grad_accum 1 \
+  --init_net weight/net_params_200.pkl --out weight/smoke --log_dir logs/smoke_lora || true
+CUDA_VISIBLE_DEVICES=0 python src/train_stage2.py --config configs/config_smoke.yml \
+  --steps 2 --batch_size 1 --grad_accum 1 --teacher osediff --skip_ram \
+  --init_net weight/net_params_200.pkl --out weight/smoke --log_dir logs/smoke_s2 || true
+echo smoke done

scripts/run_stage1.sh

@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+# S1 LoRA 像素域适配 (A100 40G)
+set -euo pipefail
+cd "$(dirname "$0")/.."
+source "$(conda info --base)/etc/profile.d/conda.sh"; conda activate AdcSR
+CUDA_VISIBLE_DEVICES=0 python src/train_lora.py \
+  --config configs/config_s1_lora.yml \
+  --manifest data/manifest_train.json --real_prob 0.35 \
+  --model_id models/stable-diffusion-2-1-base \
+  --half_decoder weight/pretrained/halfDecoder.ckpt \
+  --init_net weight/net_params_200.pkl \
+  --out weight/s1 --log_dir logs/s1 \
+  --steps 20000 --batch_size 8 --grad_accum 2 --lr 5e-5 --lora_rank 64 \
+  --save_every 2000 --w_l1 1.0 --w_lpips 1.0 --w_dists 0.3 --w_wave 0.5 --w_color 0.2

scripts/run_stage2.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# S2 GDPO 教师蒸馏(先 probe; GDPO 不可用回退 osediff)
+set -euo pipefail
+cd "$(dirname "$0")/.."
+source "$(conda info --base)/etc/profile.d/conda.sh"; conda activate AdcSR
+TEACHER="${1:-gdpo}"
+if [ "$TEACHER" = "gdpo" ] && [ ! -d weight/gdpo ]; then
+  echo "未找到 weight/gdpo, 回退 osediff"; TEACHER=osediff
+fi
+CUDA_VISIBLE_DEVICES=0 python src/train_stage2.py \
+  --config configs/config_s2_distill.yml \
+  --model_id models/stable-diffusion-2-1-base \
+  --half_decoder weight/pretrained/halfDecoder.ckpt \
+  --teacher "$TEACHER" --gdpo_dir weight/gdpo \
+  --init_net weight/s1/net_params_BEST.pkl \
+  --out weight/s2 --log_dir logs/s2 \
+  --steps 30000 --batch_size 8 --grad_accum 2 --lr 5e-5 --lr_D 1e-6 \
+  --save_every 2000 --w_wave 0.5 --w_lpips 0.5 --w_dists 0.3 --w_color 0.2

scripts/run_stage3.sh

@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+# S3 域微调(S2 最优权重, 小 LR)
+set -euo pipefail
+cd "$(dirname "$0")/.."
+source "$(conda info --base)/etc/profile.d/conda.sh"; conda activate AdcSR
+CUDA_VISIBLE_DEVICES=0 python src/train_stage2.py \
+  --config configs/config_s3_ft.yml \
+  --model_id models/stable-diffusion-2-1-base \
+  --half_decoder weight/pretrained/halfDecoder.ckpt \
+  --teacher gdpo --gdpo_dir weight/gdpo \
+  --init_net weight/s2/net_params_BEST.pkl \
+  --out weight/s3 --log_dir logs/s3 \
+  --steps 20000 --batch_size 8 --grad_accum 2 --lr 2e-5 --lr_D 5e-7 \
+  --save_every 2000 --w_wave 0.5 --w_lpips 0.5 --w_dists 0.3 --w_color 0.2

scripts/setup_env.sh

@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# aspire2a 环境准备（A100 40G 单卡; python 3.10）
+set -euo pipefail
+cd "$(dirname "$0")/.."
+if ! command -v conda >/dev/null 2>&1; then echo "需要 conda"; exit 1; fi
+conda create -n AdcSR python=3.10 -y || true
+source "$(conda info --base)/etc/profile.d/conda.sh"
+conda activate AdcSR
+pip install --upgrade pip
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+pip install -r official/requirements.txt
+pip install diffusers transformers accelerate peft omegaconf tqdm pyyaml pyiqa lpips scikit-image imageio pillow numpy
+pip install huggingface_hub hf_transfer safetensors gdown
+pip install "dwt-pytorch" 2>/dev/null || pip install pywavelets || true
+python - <<'EOF'
+try:
+    import pywt; print("pywt ok")
+except Exception as e:
+    print("pywt missing:", e)
+EOF
+echo "setup done"

src/check_submission.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+"""提交包检查: 命名/数量/大小/结构/可加载/确定性。
+用法: python src/check_submission.py --zip xxx.zip [--names data/test_names.txt] [--test_runner]
+"""
+import argparse, json, os, sys, tempfile, zipfile
+import torch
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--zip", required=True)
+    ap.add_argument("--names", default="", help="每行一个期望文件名")
+    ap.add_argument("--test_runner", action="store_true", help="解压并实际加载 jit 跑 1 次")
+    args = ap.parse_args()
+    size = os.path.getsize(args.zip)
+    checks = {"exists": True, "size_gb": round(size / 1e9, 2), "le_10gb": size <= 10e9}
+    with zipfile.ZipFile(args.zip) as z:
+        names = z.namelist()
+        outs = [n for n in names if "/output_dir/" in n or n.startswith("output_dir/")]
+        jpgs = [n for n in outs if n.lower().endswith(".jpg")]
+        mods = [n for n in names if "/model_dir/" in n or n.startswith("model_dir/")]
+        checks["output_jpg_count"] = len(jpgs)
+        checks["has_model"] = any(n.endswith(".pt") for n in mods)
+        checks["has_runner"] = any(n.endswith("runner.py") for n in mods)
+        if args.names:
+            expect = [l.strip() for l in open(args.names, encoding="utf-8") if l.strip()]
+            got = {os.path.basename(n) for n in jpgs}
+            checks["missing"] = [e for e in expect if e not in got]
+            checks["extra"] = sorted(got - set(expect))[:10]
+        if args.test_runner:
+            tmp = tempfile.mkdtemp()
+            z.extractall(tmp)
+            pt = [os.path.join(tmp, n) for n in mods if n.endswith(".pt")][0]
+            m = torch.jit.load(pt, map_location="cpu")
+            m.eval()
+            x = torch.randn(1, 3, 512, 512).half() * 0.5
+            with torch.no_grad():
+                o1 = m(x); o2 = m(x)
+            checks["jit_load_ok"] = True
+            checks["out_shape"] = list(o1.shape)
+            checks["deterministic_maxdiff"] = float((o1 - o2).abs().max())
+    for k, v in checks.items():
+        print(f"{k}: {v}")
+    bad = [k for k, v in checks.items() if v is False] or \
+          (checks.get("missing") if checks.get("missing") else [])
+    print("PASS" if not bad else f"FAIL items: {bad}")
+
+if __name__ == "__main__":
+    main()

src/common.py

@@ -0,0 +1,213 @@
+#!/usr/bin/env python
+"""AdcSR 工程公共工具：路径、模型装配、手工 LoRA 注入、教师加载、GDPO probe。
+训练脚本统一从这里 import，禁止各自重复实现装配逻辑。
+"""
+import os, sys, copy, json, types
+from pathlib import Path
+
+REPO = Path(__file__).resolve().parents[1]
+OFFICIAL = REPO / "official"
+
+def ensure_official():
+    if str(OFFICIAL) not in sys.path:
+        sys.path.insert(0, str(OFFICIAL))
+
+ensure_official()
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# 模型装配（与 official/test.py 全链一致）
+# ---------------------------------------------------------------------------
+def load_diffusers_sd(model_id, dtype=torch.float32, device="cpu"):
+    from diffusers import StableDiffusionPipeline
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype).to(device)
+    return pipe.vae, pipe.unet, pipe.text_encoder, pipe.tokenizer
+
+def load_pruned_decoder(half_decoder_ckpt, device="cpu", dtype=torch.float32):
+    from diffusers.models.autoencoders.vae import Decoder
+    decoder = Decoder(in_channels=4, out_channels=3,
+                      up_block_types=["UpDecoderBlock2D"] * 4,
+                      block_out_channels=[64, 128, 256, 256], layers_per_block=2,
+                      norm_num_groups=32, act_fn="silu", norm_type="group",
+                      mid_block_add_attention=True).to(device=device, dtype=dtype)
+    ckpt = torch.load(half_decoder_ckpt, map_location="cpu", weights_only=False)
+    sd = {k.replace("decoder.", ""): v for k, v in ckpt["state_dict"].items() if k.startswith("decoder.")}
+    decoder.load_state_dict(sd, strict=True)
+    return decoder
+
+def build_net(unet, decoder):
+    from model import Net  # official
+    return Net(unet, copy.deepcopy(decoder))
+
+def assemble_full_student(unet, decoder, net_weights=None, device="cuda", dtype=torch.float32):
+    """???? 512 ???: Net(unet,decoder) + decoder ??(up_blocks...conv_out)?"""
+    net = build_net(unet, decoder)
+    if net_weights is not None:
+        sd = torch.load(net_weights, map_location="cpu", weights_only=False)
+        if any(k.startswith("module.") for k in sd):
+            sd = {k.replace("module.", "", 1): v for k, v in sd.items()}
+        net.load_state_dict(sd, strict=True)
+    net.to(device=device, dtype=dtype)
+    tail_mods = [decoder.up_blocks, decoder.conv_norm_out, decoder.conv_act, decoder.conv_out]
+    for m in tail_mods:
+        m.to(device=device, dtype=dtype)
+    full = nn.Sequential(net, *tail_mods)
+    return full
+
+
+def build_discriminator_unet(unet_copy, rank=4, dtype=torch.float32, device="cuda"):
+    """official 判别器: conv_in 4->256 + LoRA(unet)。"""
+    from utils import add_lora_to_unet
+    unet_D = copy.deepcopy(unet_copy).to(device=device, dtype=dtype)
+    cin = unet_D.conv_in
+    new_conv_in = nn.Conv2d(256, cin.out_channels, 3, padding=1).to(device=device, dtype=dtype)
+    new_conv_in.weight.data = cin.weight.data.repeat(1, 64, 1, 1) / 64
+    new_conv_in.bias.data = cin.bias.data
+    unet_D.conv_in = new_conv_in
+    unet_D = add_lora_to_unet(unet_D, rank=rank)
+    unet_D.set_adapters(["default_encoder", "default_decoder", "default_others"])
+    return unet_D
+
+# ---------------------------------------------------------------------------
+# 教师加载
+# ---------------------------------------------------------------------------
+def load_osediff_teacher(osediff_pkl, device="cuda", dtype=torch.float32):
+    ckpt = torch.load(osediff_pkl, map_location="cpu", weights_only=False)
+    return ckpt  # {"vae":..., "unet":...}
+
+def load_gdpo_teacher(gdpo_dir, device="cuda", dtype=torch.float32):
+    """GDPO ???????????????????????? probe_gdpo?
+    ?? diffusers UNet2DConditionModel?state dict ????? dict?"""
+    from diffusers import UNet2DConditionModel
+    if os.path.isdir(os.path.join(gdpo_dir, "unet")) and os.path.exists(
+            os.path.join(gdpo_dir, "unet", "diffusion_pytorch_model.safetensors")):
+        return UNet2DConditionModel.from_pretrained(os.path.join(gdpo_dir, "unet"),
+                                                    torch_dtype=dtype).to(device)
+    if os.path.isdir(gdpo_dir):
+        if os.path.exists(os.path.join(gdpo_dir, "diffusion_pytorch_model.safetensors")):
+            p = os.path.join(gdpo_dir, "diffusion_pytorch_model.safetensors")
+        else:
+            p = os.path.join(gdpo_dir, "ckp", "diffusion_pytorch_model.safetensors")
+        if os.path.exists(p):
+            try:
+                return UNet2DConditionModel.from_pretrained(os.path.dirname(p),
+                                                            torch_dtype=dtype).to(device)
+            except Exception:
+                return _load_raw(p)
+    if os.path.isfile(gdpo_dir):
+        return _load_raw(gdpo_dir)
+    raise RuntimeError("GDPO ??????????? python -m src.common --probe_gdpo <path> ?????"
+                       "??? --teacher osediff")
+
+def _load_raw(p):
+    if p.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        return load_file(p)
+    return torch.load(p, map_location="cpu", weights_only=False)
+
+def probe_gdpo(gdpo_path):
+    """打印权重键结构与前缀，帮助实现 GDPO->diffusers UNet 映射。"""
+    if gdpo_path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        sd = load_file(gdpo_path)
+    else:
+        sd = torch.load(gdpo_path, map_location="cpu", weights_only=False)
+        if isinstance(sd, dict) and "state_dict" in sd:
+            sd = sd["state_dict"]
+    keys = list(sd.keys())
+    print("num keys:", len(keys))
+    for k in keys[:40]:
+        print(k, tuple(sd[k].shape) if hasattr(sd[k], "shape") else type(sd[k]))
+    # 判断是否为完整 UNet（含 down_blocks）或 LoRA 或 Pipeline
+    has_unet = any("down_blocks" in k for k in keys)
+    has_lora = any("lora" in k.lower() for k in keys)
+    print("has_unet_blocks:", has_unet, "| has_lora:", has_lora)
+
+# ---------------------------------------------------------------------------
+# 手工 LoRA（对任意 Conv2d/Linear 注入，规避 peft 在剪枝/删模块后的解析问题）
+# ---------------------------------------------------------------------------
+class LoRAConv2d(nn.Module):
+    def __init__(self, conv: nn.Conv2d, r: int, alpha: float = 1.0):
+        super().__init__()
+        self.conv = conv
+        self.r = max(1, r)
+        self.alpha = alpha
+        cin, cout = conv.in_channels, conv.out_channels
+        self.lora_a = nn.Parameter(torch.zeros(cin, self.r))
+        self.lora_b = nn.Parameter(torch.zeros(self.r, cout))
+        nn.init.kaiming_uniform_(self.lora_a, a=5 ** 0.5)
+        nn.init.zeros_(self.lora_b)
+        self.requires_grad_(False)
+        self.lora_a.requires_grad_(True)
+        self.lora_b.requires_grad_(True)
+
+    def forward(self, x):
+        y = self.conv(x)
+        if self.training or True:
+            # 1x1 conv low-rank: 输入cin->r->cout, 保持空间尺寸
+            z = F.conv2d(x, self.lora_a.t().view(self.r, cin, 1, 1))
+            z = F.conv2d(z, self.lora_b.t().view(cout, self.r, 1, 1))
+            return y + self.alpha * z
+        return y
+
+class LoRALinear(nn.Module):
+    def __init__(self, lin: nn.Linear, r: int, alpha: float = 1.0):
+        super().__init__()
+        self.lin = lin
+        self.r = max(1, r)
+        self.alpha = alpha
+        cin, cout = lin.in_features, lin.out_features
+        self.lora_a = nn.Parameter(torch.zeros(cin, self.r))
+        self.lora_b = nn.Parameter(torch.zeros(self.r, cout))
+        nn.init.kaiming_uniform_(self.lora_a, a=5 ** 0.5)
+        nn.init.zeros_(self.lora_b)
+        self.requires_grad_(False)
+        self.lora_a.requires_grad_(True)
+        self.lora_b.requires_grad_(True)
+
+    def forward(self, x):
+        y = self.lin(x)
+        z = F.linear(x, self.lora_a.t())
+        z = F.linear(z, self.lora_b.t())
+        return y + self.alpha * z
+
+def _names(model):
+    for n, m in model.named_modules():
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            yield n, m
+
+def inject_lora(model, rank=64, alpha=1.0, skip_bias_norm=True, include=("conv", "to_q", "to_k", "to_v", "proj", "ff", "linear")):
+    """替换模型内所有 Conv2d/Linear 为 LoRA 包装（原始权重冻结，仅训练 lora_a/b）。
+    include: 子串过滤，None=全部。"""
+    for n, m in list(_names(model)):
+        if include is not None and not any(s in n for s in include):
+            continue
+        parent, attr = _find_parent(model, n)
+        if isinstance(m, nn.Conv2d):
+            setattr(parent, attr, LoRAConv2d(m, rank, alpha))
+        elif isinstance(m, nn.Linear):
+            setattr(parent, attr, LoRALinear(m, rank, alpha))
+    return model
+
+def _find_parent(model, name):
+    parts = name.split(".")
+    node = model
+    for p in parts[:-1]:
+        node = getattr(node, p)
+    return node, parts[-1]
+
+def lora_params(model):
+    for p in model.parameters():
+        if p.requires_grad:
+            yield p
+
+def count_params(model, only_trainable=False):
+    if only_trainable:
+        return sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return sum(p.numel() for p in model.parameters())
+
+if __name__ == "__main__":
+    print("common module OK; official dir:", OFFICIAL)

src/eval_val.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+"""Val 评测 + 代理分: 在隔离 val(真实对/官方对) 上计算 8 指标 + proxy + 保真护栏。
+用法:
+  python src/eval_val.py --weights weight/s2/net_params_X.pkl --pairs_json data/manifest_val.json
+  python src/eval_val.py --weights ... --official_lq <val lq dir> --official_gt <val gt dir>
+"""
+import argparse, json, os, sys
+from pathlib import Path
+REPO = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO)); sys.path.insert(0, str(REPO / "src")); sys.path.insert(0, str(REPO / "official"))
+import numpy as np, torch
+from PIL import Image
+from inference_4k import build_model, infer_4k_single
+
+def load_metrics(device):
+    import pyiqa
+    return {k: pyiqa.create_metric(k, device=device) for k in
+            ["psnr", "ssim", "lpips", "dists", "niqe", "maniqa", "musiq", "clipiqa"]}
+
+def resize_pair(sr, hr, max_side=1024):
+    sr = sr.convert("RGB"); hr = hr.convert("RGB")
+    for im in (sr, hr):
+        w, h = im.size
+        if max(w, h) > max_side:
+            s = max_side / max(w, h)
+            im.thumbnail((int(w * s), int(h * s)), Image.LANCZOS)
+    return sr, hr
+
+def metric_dict(metrics, sr, hr, device):
+    import torchvision.transforms as T
+    def t(im):
+        return T.ToTensor()(im).to(device).unsqueeze(0)
+    a, b = t(sr), t(hr)
+    out = {"psnr": metrics["psnr"](a, b).item(), "ssim": metrics["ssim"](a, b).item(),
+           "lpips": metrics["lpips"](a, b).item(), "dists": metrics["dists"](a, b).item(),
+           "niqe": metrics["niqe"](a).item(), "maniqa": metrics["maniqa"](a).item(),
+           "musiq": metrics["musiq"](a).item(), "clipiqa": metrics["clipiqa"](a).item()}
+    return out
+
+def proxy(m):
+    return (0.25 * m["clipiqa"] + 0.2 * (m["musiq"] / 100.0) + 0.15 * m["maniqa"]
+            + 0.15 * (1 - m["niqe"] / 10.0) + 0.25 * (1 - m["lpips"]))
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--weights", required=True)
+    ap.add_argument("--half_decoder", default="weight/pretrained/halfDecoder.ckpt")
+    ap.add_argument("--model_id", default="models/stable-diffusion-2-1-base")
+    ap.add_argument("--pairs_json", default="", help='[{lr,hr}] 或 manifest 结构')
+    ap.add_argument("--official_lq", default="", help="同分辨率官方对(整图滑窗)")
+    ap.add_argument("--official_gt", default="")
+    ap.add_argument("--out", default="logs/eval_result.json")
+    ap.add_argument("--max_side", type=int, default=1024)
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    net, tail = build_model(args.weights, args.half_decoder, args.model_id, device, bf16=True)
+    metrics = load_metrics(device)
+    rows = []
+    if args.pairs_json:
+        with open(args.pairs_json, encoding="utf-8") as fh:
+            data = json.load(fh)
+        pairs = data.get("real_pairs", data if isinstance(data, list) else [])
+        for i, p in enumerate(pairs):
+            with torch.no_grad():
+                lr = Image.open(p["lr"]).convert("RGB")
+                hr = Image.open(p["hr"]).convert("RGB")
+                lr128 = lr.resize((max(1, lr.width // 4), max(1, lr.height // 4)), Image.LANCZOS)
+                t = torch.from_numpy(np.asarray(lr128, dtype=np.float32).transpose(2, 0, 1) / 255.0 * 2 - 1)[None].to(device)
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    z = net(t); sr_arr = tail(z)
+                sr = Image.fromarray(((sr_arr[0].float().cpu().numpy().transpose(1, 2, 0) + 1) / 2 * 255).clip(0, 255).astype(np.uint8))
+            sr, hr = resize_pair(sr, hr, args.max_side)
+            m = metric_dict(metrics, sr, hr, device)
+            m["name"] = os.path.basename(p["lr"]); m["proxy"] = proxy(m)
+            rows.append(m)
+            print(f"  [{i+1}] {m['name']} proxy {m['proxy']:.4f}", flush=True)
+    if args.official_lq and args.official_gt:
+        lq_files = sorted(os.listdir(args.official_lq))
+        for f in lq_files:
+            lq_p = os.path.join(args.official_lq, f)
+            gt_p = os.path.join(args.official_gt, f.replace("_lq", "_gt").replace("lq.jpg", "gt.png"))
+            if not os.path.exists(gt_p):
+                gt_p = os.path.join(args.official_gt, f.replace("_lq.jpg", "_gt.jpg"))
+            if not os.path.exists(gt_p):
+                continue
+            sr = infer_4k_single(lq_p, net, tail, device)
+            gt = Image.open(gt_p).convert("RGB")
+            sr, gt = resize_pair(sr, gt, args.max_side)
+            m = metric_dict(metrics, sr, gt, device)
+            m["name"] = f; m["proxy"] = proxy(m)
+            rows.append(m)
+            print(f"  [official] {f} proxy {m['proxy']:.4f}", flush=True)
+    if not rows:
+        print("无评测样本"); return
+    keys = ["psnr", "ssim", "lpips", "dists", "niqe", "maniqa", "musiq", "clipiqa", "proxy"]
+    agg = {k: float(np.mean([r[k] for r in rows])) for k in keys}
+    result = {"per_image": rows, "mean": agg}
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    with open(args.out, "w", encoding="utf-8") as fh:
+        json.dump(result, fh, ensure_ascii=False, indent=1)
+    print(json.dumps(agg, indent=1))
+
+if __name__ == "__main__":
+    main()

src/export_jit.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+"""导出 torch.jit(fp16) 512x512 模型: 输入 LR[1,3,512,512][-1,1] -> 输出同尺寸。
+内部: bicubic 512->128 -> 官方 4x 学生全链 -> 512。 (AdaIN 后处理不进测速模型)
+用法: python src/export_jit.py --net weight/s2/net_params_X.pkl --out model_dir
+"""
+import argparse, os, sys
+from pathlib import Path
+REPO = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO)); sys.path.insert(0, str(REPO / "src")); sys.path.insert(0, str(REPO / "official"))
+import torch, torch.nn as nn, torch.nn.functional as F
+from common import load_diffusers_sd, load_pruned_decoder, build_net
+
+class SR512(nn.Module):
+    def __init__(self, net, tail):
+        super().__init__()
+        self.net = net
+        self.tail = tail
+
+    def forward(self, x512):
+        x128 = F.interpolate(x512, size=(128, 128), mode="bicubic", align_corners=False)
+        z = self.net(x128)
+        return self.tail(z)
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--net", required=True)
+    ap.add_argument("--out", default="model_dir")
+    ap.add_argument("--half_decoder", default="weight/pretrained/halfDecoder.ckpt")
+    ap.add_argument("--model_id", default="models/stable-diffusion-2-1-base")
+    ap.add_argument("--name", default="your_model.pt")
+    args = ap.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    vae, unet, _, _ = load_diffusers_sd(args.model_id, dtype=torch.float32, device="cpu")
+    del vae
+    decoder = load_pruned_decoder(args.half_decoder, device="cpu", dtype=torch.float32)
+    net = build_net(unet, decoder)
+    sd = torch.load(args.net, map_location="cpu", weights_only=False)
+    if any(k.startswith("module.") for k in sd):
+        sd = {k.replace("module.", "", 1): v for k, v in sd.items()}
+    net.load_state_dict(sd, strict=True)
+    net.eval()
+    tail = nn.Sequential(*decoder.up_blocks, decoder.conv_norm_out,
+                         decoder.conv_act, decoder.conv_out).eval()
+    model = SR512(net, tail).to(device).half().eval()
+    # trace on fixed 512 fp16
+    dummy = torch.randn(1, 3, 512, 512, device=device).half() * 0.5
+    with torch.no_grad():
+        traced = torch.jit.trace(model, dummy, check_trace=False)
+    traced = torch.jit.freeze(traced)
+    out_path = os.path.join(args.out, args.name)
+    traced.save(out_path)
+    # 自检: 两次前向一致性 + 形状
+    with torch.no_grad():
+        o1 = traced(dummy); o2 = traced(dummy)
+    assert o1.shape == dummy.shape, o1.shape
+    err = (o1 - o2).abs().max().item()
+    print("saved", out_path, "| deterministic max-diff:", err)
+    print("output range sample:", float(o1.min()), float(o1.max()))
+
+if __name__ == "__main__":
+    main()

src/inference_4k.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+"""4K 整图增强: 512 窗 + overlap 128 线性融合 + AdaIN 对齐 LR(禁止 zero-padding, 边缘 reflect)。
+x1 语义: 对每个 512 窗先缩到 128, 走官方 4x 学生模型, 回 512。
+用法: python src/inference_4k.py --lr_dir <dir> --net weight/s2/net_params_X.pkl --out output_dir
+"""
+import argparse, copy, os, sys
+from pathlib import Path
+REPO = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO)); sys.path.insert(0, str(REPO / "src")); sys.path.insert(0, str(REPO / "official"))
+import torch, torch.nn.functional as F
+import numpy as np
+from PIL import Image
+from common import load_diffusers_sd, load_pruned_decoder, build_net
+
+def build_model(net_pkl, half_decoder, model_id, device, bf16=True):
+    dtype = torch.bfloat16 if (bf16 and device == "cuda") else torch.float32
+    vae, unet, _, _ = load_diffusers_sd(model_id, dtype=torch.float32, device="cpu")
+    del vae
+    decoder = load_pruned_decoder(half_decoder, device="cpu", dtype=torch.float32)
+    net = build_net(unet, decoder)
+    net.to(device=device, dtype=dtype)
+    sd = torch.load(net_pkl, map_location="cpu", weights_only=False)
+    if any(k.startswith("module.") for k in sd):
+        sd = {k.replace("module.", "", 1): v for k, v in sd.items()}
+    net.load_state_dict(sd, strict=True)
+    net.eval()
+    tail = torch.nn.Sequential(*decoder.up_blocks, decoder.conv_norm_out,
+                               decoder.conv_act, decoder.conv_out).to(device=device, dtype=dtype).eval()
+    return net, tail
+
+def infer_tile(net, tail, lr128, lr_stats):
+    """lr128: [1,3,128,128] [-1,1]; 返回 [1,3,512,512] [-1,1] 并 AdaIN 对齐 lr_stats"""
+    with torch.no_grad():
+        z = net(lr128)
+        out = tail(z)
+    out = out.float()
+    mu, std = out.mean(dim=(2, 3), keepdim=True), out.std(dim=(2, 3), keepdim=True)
+    lmu, lstd = lr_stats
+    out = (out - mu) / (std + 1e-6) * lstd + lmu
+    return out.clamp(-1, 1)
+
+def infer_4k_single(lr_path, net, tail, device, tile=512, overlap=128, bf16=True):
+    im = Image.open(lr_path).convert("RGB")
+    w, h = im.size
+    a = np.asarray(im, dtype=np.float32) / 255.0  # [H,W,3]
+    # reflect pad so tile 整除
+    stride = tile - overlap
+    pad_r = (stride - w % stride) % stride + (tile - stride) if (w % stride) else max(0, tile - stride)
+    pad_b = (stride - h % stride) % stride + (tile - stride) if (h % stride) else max(0, tile - stride)
+    a = np.pad(a, ((0, pad_b), (0, pad_r), (0, 0)), mode="reflect")
+    H, W = a.shape[:2]
+    acc = np.zeros((H, W, 3), dtype=np.float64)
+    wsum = np.zeros((H, W, 1), dtype=np.float64)
+    rows = list(range(0, H - tile + 1, stride)) or [0]
+    cols = list(range(0, W - tile + 1, stride)) or [0]
+    if rows[-1] + tile < H: rows.append(H - tile)
+    if cols[-1] + tile < W: cols.append(W - tile)
+    # 线性窗权重
+    ramp = np.minimum(np.arange(tile), np.arange(tile)[::-1]) / (tile // 2)
+    w2d = (ramp[None, :] * ramp[:, None])[..., None].astype(np.float64)
+    for y in rows:
+        for x in cols:
+            crop = a[y:y + tile, x:x + tile]
+            lr128 = np.asarray(Image.fromarray((crop * 255).astype(np.uint8)).resize((128, 128), Image.BICUBIC),
+                               dtype=np.float32) / 255.0
+            lr_t = torch.from_numpy(lr128.transpose(2, 0, 1))[None].to(device) * 2 - 1
+            stats = (torch.tensor(crop.mean(axis=(0, 1))[None, :, None, None], device=device) * 2 - 1,
+                     torch.tensor(crop.std(axis=(0, 1))[None, :, None, None], device=device))
+            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=bf16):
+                out = infer_tile(net, tail, lr_t, stats)
+            o = out[0].float().cpu().numpy().transpose(1, 2, 0)  # [512,512,3] [-1,1]
+            o = (o + 1) / 2.0
+            acc[y:y + tile, x:x + tile] += o.astype(np.float64) * w2d
+            wsum[y:y + tile, x:x + tile] += w2d
+    res = acc / (wsum + 1e-9)
+    res = res[:h, :w]
+    return Image.fromarray((np.clip(res, 0, 1) * 255).astype(np.uint8))
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--lr_dir", required=True)
+    ap.add_argument("--net", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--half_decoder", default="weight/pretrained/halfDecoder.ckpt")
+    ap.add_argument("--model_id", default="models/stable-diffusion-2-1-base")
+    ap.add_argument("--tile", type=int, default=512)
+    ap.add_argument("--overlap", type=int, default=128)
+    ap.add_argument("--no_bf16", action="store_true")
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    os.makedirs(args.out, exist_ok=True)
+    net, tail = build_model(args.net, args.half_decoder, args.model_id, device, bf16=not args.no_bf16)
+    files = sorted([f for f in os.listdir(args.lr_dir) if f.lower().endswith((".jpg", ".jpeg", ".png"))])
+    import time
+    t0 = time.time()
+    for i, f in enumerate(files):
+        out_im = infer_4k_single(os.path.join(args.lr_dir, f), net, tail, device,
+                                 tile=args.tile, overlap=args.overlap, bf16=not args.no_bf16)
+        out_im.save(os.path.join(args.out, f), quality=95)
+        if (i + 1) % 10 == 0:
+            print(f"  {i+1}/{len(files)} elapsed {time.time()-t0:.1f}s", flush=True)
+    print(f"done {len(files)} imgs in {time.time()-t0:.1f}s")
+
+if __name__ == "__main__":
+    main()

src/latency_test.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+"""512x512 fp16 时延测速: 加载 torch.jit 模型, warmup + N 次前向。
+用法: python src/latency_test.py --model model_dir/your_model.pt [--osediff_latency 0.168]
+"""
+import argparse, time, statistics
+import torch
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--n", type=int, default=100)
+    ap.add_argument("--warmup", type=int, default=10)
+    ap.add_argument("--osediff_latency", type=float, default=0.0, help="秒; 提供则打印加速比")
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    m = torch.jit.load(args.model, map_location=device)
+    m.eval()
+    x = torch.randn(1, 3, 512, 512, device=device).half()
+    with torch.no_grad():
+        for _ in range(args.warmup):
+            m(x)
+        torch.cuda.synchronize() if device == "cuda" else None
+        times = []
+        for _ in range(args.n):
+            if device == "cuda":
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            with torch.no_grad():
+                m(x)
+            if device == "cuda":
+                torch.cuda.synchronize()
+            times.append(time.perf_counter() - t0)
+    mean = statistics.mean(times)
+    med = statistics.median(times)
+    print(f"mean {mean*1000:.3f} ms | median {med*1000:.3f} ms | n={args.n}")
+    if args.osediff_latency > 0:
+        print(f"speedup vs OSEDiff({args.osediff_latency*1000:.1f}ms): {args.osediff_latency/mean:.2f}x")
+
+if __name__ == "__main__":
+    main()