File size: 7,474 Bytes
45d3d29 | 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | #!/usr/bin/env python3
"""Run the released RAESR or RAESR-D 512px inference protocol on an image folder.
Both models use the same input protocol and the paper-reported mean-only color
matching applied to the float decoder output before PNG quantization.
"""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from raesr.artifacts import (
DELIVERY_ROOT as PACKAGE_ROOT,
DINO_S_WEIGHT_NAME,
DINO_WEIGHT_NAME,
build_paper_model,
build_raesr_d_model,
configure_dinov3,
file_record,
runtime_environment,
runtime_record,
verify_decoder,
verify_frontend,
verify_runtime,
write_json,
)
MEAN_ONLY_PROTOCOL = "paper_mean_only_color_match_512px_bicubic_input_resize_v1"
COLOR_MATCH_FORMULA = "clamp(sr - mean(sr,H,W) + mean(lr,H,W), 0, 1)"
IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}
MODEL_CHOICES = ("raesr", "raesr-d")
def match_color_mean(sr: torch.Tensor, lr: torch.Tensor) -> torch.Tensor:
"""Shift each SR channel to the LR mean without changing its contrast."""
sr = sr.float()
lr = lr.float()
mean_sr = sr.mean(dim=(-2, -1), keepdim=True)
mean_lr = lr.mean(dim=(-2, -1), keepdim=True)
return (sr - mean_sr + mean_lr).clamp(0.0, 1.0)
def collect_images(directory: Path) -> list[Path]:
if not directory.is_dir():
raise FileNotFoundError(f"input directory does not exist: {directory}")
images = sorted(path for path in directory.iterdir() if path.suffix.lower() in IMAGE_SUFFIXES)
if not images:
raise ValueError(f"no supported images in {directory}")
stems: set[str] = set()
for path in images:
if path.stem in stems:
raise ValueError(f"ambiguous duplicate input stem: {path.stem}")
stems.add(path.stem)
with Image.open(path) as image:
image.convert("RGB").load()
return images
def is_valid_output(path: Path) -> bool:
try:
with Image.open(path) as image:
image.verify()
with Image.open(path) as image:
return image.format == "PNG" and image.size == (512, 512) and image.mode in {"RGB", "RGBA"}
except (OSError, ValueError):
return False
def load_input(path: Path, device: str) -> torch.Tensor:
array = np.asarray(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0
tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).contiguous().to(device)
if tensor.shape[-2:] != (512, 512):
tensor = F.interpolate(tensor, size=(512, 512), mode="bicubic", align_corners=False).clamp(0, 1)
return tensor
def save_png(tensor: torch.Tensor, destination: Path) -> None:
array = tensor[0].detach().float().cpu().numpy().transpose(1, 2, 0)
image = Image.fromarray((array * 255.0).round().astype(np.uint8))
temporary = destination.with_name(f".{destination.name}.tmp.{os.getpid()}")
try:
image.save(temporary, format="PNG")
temporary.replace(destination)
finally:
if temporary.exists():
temporary.unlink()
def build_model(args: argparse.Namespace):
if args.model == "raesr-d":
return build_raesr_d_model(
args.frontend, args.decoder, args.dinov3_root, args.dinov3_weights, args.device
)
return build_paper_model(args.decoder, args.dinov3_root, args.dinov3_weights, args.device)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--input-dir", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model", choices=MODEL_CHOICES, default="raesr")
parser.add_argument("--decoder", type=Path, default=PACKAGE_ROOT / "artifacts" / "ema_decoder.pt")
parser.add_argument("--frontend", type=Path, default=PACKAGE_ROOT / "weights" / "raesr_d_frontend.pt",
help="RAESR-D frontend weights (used with --model raesr-d)")
parser.add_argument("--dinov3-root", type=Path, required=True)
parser.add_argument("--dinov3-weights", type=Path, required=True)
parser.add_argument("--device", default="cuda")
parser.add_argument("--num-shards", type=int, default=1)
parser.add_argument("--shard-index", type=int, default=0)
parser.add_argument("--overwrite", action="store_true")
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards:
parser.error("--shard-index must be in [0, --num-shards)")
images = collect_images(args.input_dir.expanduser().resolve())
args.output_dir.mkdir(parents=True, exist_ok=True)
selected = images[args.shard_index::args.num_shards]
pending = [
path for path in selected
if args.overwrite or not is_valid_output(args.output_dir / f"{path.stem}.png")
]
dinov3_weight_name = DINO_S_WEIGHT_NAME if args.model == "raesr-d" else DINO_WEIGHT_NAME
base_manifest = {
"model": args.model,
"protocol": MEAN_ONLY_PROTOCOL,
"color_match": "mean-only",
"color_match_formula": COLOR_MATCH_FORMULA,
"input_dir": str(args.input_dir.expanduser().resolve()),
"output_dir": str(args.output_dir.resolve()),
"shard": {"index": args.shard_index, "count": args.num_shards},
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
"input_count": len(images),
"selected_count": len(selected),
"pending_count": len(pending),
"inputs": [file_record(path) for path in selected],
"decoder": verify_decoder(args.decoder),
"dinov3": configure_dinov3(args.dinov3_root, args.dinov3_weights, dinov3_weight_name),
"runtime": runtime_record(),
"runtime_verification": verify_runtime(),
"environment": runtime_environment(),
}
if args.model == "raesr-d":
base_manifest["frontend"] = verify_frontend(args.frontend)
manifest_name = f"run_manifest.shard_{args.shard_index:02d}.json"
if args.dry_run:
base_manifest["status"] = "DRY_RUN"
write_json(args.output_dir / manifest_name, base_manifest)
print(f"validated {len(selected)} inputs; manifest -> {args.output_dir / manifest_name}")
return
model = build_model(args)
for index, path in enumerate(pending, start=1):
lr = load_input(path, args.device)
with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
output = model(lr).clamp(0.0, 1.0)
output = match_color_mean(output, lr)
save_png(output, args.output_dir / f"{path.stem}.png")
if index % 25 == 0 or index == len(pending):
print(f"{index}/{len(pending)}")
outputs = [args.output_dir / f"{path.stem}.png" for path in selected]
if not all(is_valid_output(path) for path in outputs):
raise RuntimeError("one or more expected outputs are missing or invalid after inference")
base_manifest["status"] = "COMPLETE"
base_manifest["outputs"] = [file_record(path) for path in outputs]
write_json(args.output_dir / manifest_name, base_manifest)
print(f"wrote {len(pending)} images; manifest -> {args.output_dir / manifest_name}")
if __name__ == "__main__":
main()
|