PSRGAN v1.0

README.md

@@ -109,7 +109,7 @@ The model expects Sentinel-2 L2A RGB (B04, B03, B02) after a per-image 1 to 99%
 | `psrgan_x4_swag.safetensors` | SWAG mean and variance (32 snapshots) |
 | `config.json`, `psrgan.py` | model hyperparameters and code |
 | `eval/metrics.json` | every number on this page |
-| `notebooks/` | dataset creation and training notebooks |
+| `degradation.py` | randomized degradation that made the LR images of the dataset |
 
 # Evaluation
 

degradation.py

@@ -0,0 +1,90 @@
+"""Randomized degradation that made the LR images of the PSRGAN Sentinel-2 dataset.
+
+Each HR image gets one random blur (Gaussian, horizontal or vertical motion, defocus)
+and is then downsampled 4x with a random interpolation (bicubic, bilinear, area).
+
+    python degradation.py hr_dir lr_dir --seed 0
+
+Writes one LR PNG per HR image and degradations.csv with the drawn parameters.
+"""
+import argparse
+import csv
+import random
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+INTERP_NAMES = {cv2.INTER_CUBIC: "bicubic", cv2.INTER_LINEAR: "bilinear", cv2.INTER_AREA: "area"}
+
+
+def degrade(hr, scale=4, rng=random):
+    """hr: (H, W, 3) float32 in [0, 1]. Returns the LR image and the drawn parameters."""
+    h, w = hr.shape[:2]
+    filter_type = rng.choice(["gaussian", "motion", "defocus"])
+    sigma = 0.0
+
+    if filter_type == "gaussian":
+        # Isotropic Gaussian blur
+        k = rng.choice([3, 5, 7])
+        sigma = rng.uniform(0.5, 1.5)
+        blurred = cv2.GaussianBlur(hr, (k, k), sigmaX=sigma)
+
+    elif filter_type == "motion":
+        # Line kernel, horizontal or vertical
+        k = rng.choice([5, 7, 9])
+        kernel = np.zeros((k, k), np.float32)
+        if rng.random() < 0.5:
+            kernel[k // 2, :] = 1.0
+            filter_type = "motion_h"
+        else:
+            kernel[:, k // 2] = 1.0
+            filter_type = "motion_v"
+        blurred = cv2.filter2D(hr, -1, kernel / k)
+
+    else:
+        # Defocus as a disk kernel
+        k = rng.choice([5, 7, 9])
+        c = k // 2
+        yy, xx = np.mgrid[:k, :k]
+        kernel = ((yy - c) ** 2 + (xx - c) ** 2 <= c ** 2).astype(np.float32)
+        blurred = cv2.filter2D(hr, -1, kernel / kernel.sum())
+
+    interpolation = rng.choice([cv2.INTER_CUBIC, cv2.INTER_LINEAR, cv2.INTER_AREA])
+    lr = cv2.resize(blurred, (w // scale, h // scale), interpolation=interpolation)
+    return lr, {"blur": filter_type, "blur_kernel": k, "blur_sigma": sigma,
+                "interpolation": INTERP_NAMES[interpolation]}
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    p.add_argument("hr_dir")
+    p.add_argument("lr_dir")
+    p.add_argument("--scale", type=int, default=4)
+    p.add_argument("--seed", type=int, default=None)
+    args = p.parse_args()
+
+    rng = random.Random(args.seed)
+    out = Path(args.lr_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    rows = []
+    for path in sorted(Path(args.hr_dir).iterdir()):
+        img = cv2.imread(str(path), cv2.IMREAD_COLOR)
+        if img is None:
+            continue
+        # BGR stays BGR, the filters act per channel
+        hr = img.astype(np.float32) / 255
+        lr, params = degrade(hr, args.scale, rng)
+        lr = np.round(np.clip(lr, 0, 1) * 255).astype(np.uint8)
+        cv2.imwrite(str(out / f"{path.stem}.png"), lr)
+        rows.append({"file": path.name, **params})
+
+    with open(out / "degradations.csv", "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=["file", "blur", "blur_kernel", "blur_sigma", "interpolation"])
+        w.writeheader()
+        w.writerows(rows)
+    print(f"{len(rows)} LR images in {out}")
+
+
+if __name__ == "__main__":
+    main()