| """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": |
| |
| 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": |
| |
| 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: |
| |
| 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 |
| |
| 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() |
|
|
