File size: 3,933 Bytes
0193004 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """Create a deterministic, held-out synthetic restoration benchmark.
Only real-portrait folders are selected. The generated manifest and holdout
list are ignored by Git; pass the holdout list to training to avoid leakage.
"""
from __future__ import annotations
import argparse
import csv
import hashlib
import shutil
from pathlib import Path
import cv2
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SOURCE = ROOT / "models" / "CodeFormer" / "datasets" / "ffhq" / "ffhq_512"
REAL_FOLDERS = {"faces", "pinterest"}
EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
def candidates(source: Path) -> list[Path]:
files = []
for folder in REAL_FOLDERS:
directory = source / folder
if directory.is_dir():
files.extend(p for p in directory.rglob("*") if p.suffix.lower() in EXTENSIONS)
return sorted(files)
def degrade(image: np.ndarray, seed: int) -> np.ndarray:
rng = np.random.default_rng(seed)
blur_size = int(rng.choice([3, 5, 7]))
output = cv2.GaussianBlur(image, (blur_size, blur_size), float(rng.uniform(0.4, 2.5)))
h, w = output.shape[:2]
scale = float(rng.uniform(1.5, 6.0))
small = cv2.resize(output, (max(32, int(w / scale)), max(32, int(h / scale))), interpolation=cv2.INTER_AREA)
output = cv2.resize(small, (w, h), interpolation=cv2.INTER_CUBIC)
output = np.clip(output.astype(np.float32) + rng.normal(0, rng.uniform(2, 25), output.shape), 0, 255).astype(np.uint8)
quality = int(rng.integers(10, 71))
ok, encoded = cv2.imencode(".jpg", output, [cv2.IMWRITE_JPEG_QUALITY, quality])
if not ok:
raise RuntimeError("JPEG degradation failed")
return cv2.imdecode(encoded, cv2.IMREAD_COLOR)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
parser.add_argument("--count", type=int, default=500)
parser.add_argument("--seed", type=int, default=20260723)
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
source = args.source.resolve()
ranked = sorted(candidates(source), key=lambda p: hashlib.sha256(f"{args.seed}:{p.relative_to(source)}".encode()).hexdigest())
selected = ranked[:args.count]
if len(selected) < args.count:
raise SystemExit(f"Need {args.count} real portraits but found {len(selected)} in {source}")
print(f"Selected {len(selected)} held-out real portraits from {len(ranked)} candidates.")
if args.dry_run:
return
benchmark = ROOT / "benchmarks"
inputs, references = benchmark / "inputs", benchmark / "references"
inputs.mkdir(parents=True, exist_ok=True); references.mkdir(parents=True, exist_ok=True)
manifest = benchmark / "manifest.csv"
holdout = benchmark / "holdout_paths.txt"
with manifest.open("w", newline="", encoding="utf-8") as csv_file, holdout.open("w", encoding="utf-8") as holdout_file:
writer = csv.DictWriter(csv_file, fieldnames=["id", "input_path", "reference_path", "category", "notes"]); writer.writeheader()
for index, path in enumerate(selected, 1):
image = cv2.imread(str(path), cv2.IMREAD_COLOR)
if image is None: raise RuntimeError(f"Unreadable image: {path}")
sample_id = f"portrait_{index:04d}"
cv2.imwrite(str(inputs / f"{sample_id}.jpg"), degrade(image, args.seed + index), [cv2.IMWRITE_JPEG_QUALITY, 95])
shutil.copy2(path, references / f"{sample_id}{path.suffix.lower()}")
writer.writerow({"id": sample_id, "input_path": f"inputs/{sample_id}.jpg", "reference_path": f"references/{sample_id}{path.suffix.lower()}", "category": "synthetic_real_portrait", "notes": "deterministic holdout"})
holdout_file.write(path.relative_to(source).as_posix() + "\n")
print(f"Created {manifest} and {holdout}. Do not train on the listed paths.")
if __name__ == "__main__": main()
|