|
|
| """Upscale di un'immagine con lo Student 9x distillato (ONNX) alla scala richiesta.
|
|
|
| Lo Student esportato dal notebook Colab e` nativamente 9x (staged upsample 3x3).
|
| Per una scala diversa il modello non viene riaddestrato: si lavora sull'output 9x
|
|
|
| * scala < 9 -> inferenza a 9x, poi downsampling d'area sulla scala richiesta;
|
| * scala > 9 -> inferenza a 9x, poi upsampling bicubico.
|
|
|
| L'inferenza e` a tile in spazio LR (tile e pad si intendono in pixel LR), come nella
|
| cella "Usare il modello esportato" del notebook. L'eventuale canale alfa dei PNG non
|
| attraversa la rete, che e` a 3 canali: viene ricampionato a parte e riaccodato.
|
|
|
| Esempio:
|
| python scripts/upscale_onnx.py a.png --scale 6
|
| """
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import time
|
| from pathlib import Path
|
|
|
| import cv2
|
| import numpy as np
|
| import onnxruntime as ort
|
|
|
|
|
| def build_session(onnx_path: Path, providers_mode: str) -> ort.InferenceSession:
|
| """Session ONNX con arena disattivata: senza, la RAM cresce a ogni tile.
|
|
|
| Con l'inferenza a tile le shape d'input cambiano da tile a tile; il pattern
|
| allocator di ONNX Runtime trattiene i buffer del tile piu` grande e la RSS sale
|
| run dopo run fino a esaurire la memoria (processo ucciso senza traceback).
|
| """
|
| options = ort.SessionOptions()
|
| options.enable_cpu_mem_arena = False
|
| options.enable_mem_pattern = False
|
| options.log_severity_level = 3
|
| providers = (["CUDAExecutionProvider", "CPUExecutionProvider"]
|
| if providers_mode == "auto" else ["CPUExecutionProvider"])
|
| return ort.InferenceSession(str(onnx_path), sess_options=options, providers=providers)
|
|
|
|
|
| def load_image(path: Path) -> tuple[np.ndarray, np.ndarray | None]:
|
| """Carica l'immagine in BGR separando l'alpha, se presente."""
|
| img = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
|
| if img is None:
|
| raise SystemExit(f"Impossibile leggere {path}")
|
| if img.dtype != np.uint8:
|
| img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
|
| if img.ndim == 2:
|
| img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
| alpha = None
|
| if img.shape[2] == 4:
|
| img, alpha = img[:, :, :3].copy(), img[:, :, 3].copy()
|
| return img, alpha
|
|
|
|
|
| def native_scale(session: ort.InferenceSession) -> int:
|
| """Deduce la scala nativa del modello da una probe 8x9 (evita di fissarla a mano)."""
|
| name = session.get_inputs()[0].name
|
| out = session.run(None, {name: np.zeros((1, 3, 8, 9), dtype=np.float32)})[0]
|
| ratio = max(out.shape[-2] / 8.0, out.shape[-1] / 9.0)
|
| if abs(ratio - round(ratio)) > 1e-6 or ratio < 1:
|
| raise SystemExit(f"Scala nativa non intera: input 8x9 -> output {out.shape}")
|
| return int(round(ratio))
|
|
|
|
|
| def upscale_tiled(session, bgr: np.ndarray, scale: float, native: int,
|
| tile: int, pad: int) -> np.ndarray:
|
| """Inferenza a tile in spazio LR; ogni tile viene ricampionato native -> scale."""
|
| input_name = session.get_inputs()[0].name
|
| height, width = bgr.shape[:2]
|
| out_h, out_w = int(round(height * scale)), int(round(width * scale))
|
| canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8)
|
| interp = cv2.INTER_AREA if native > scale else cv2.INTER_CUBIC
|
| tiles = 0
|
| t0 = time.time()
|
| for top in range(0, height, tile):
|
| for left in range(0, width, tile):
|
| bottom, right = min(top + tile, height), min(left + tile, width)
|
|
|
|
|
| y0, y1 = max(top - pad, 0), min(bottom + pad, height)
|
| x0, x1 = max(left - pad, 0), min(right + pad, width)
|
| patch = bgr[y0:y1, x0:x1]
|
| tensor = (cv2.cvtColor(patch, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0)
|
| tensor = np.ascontiguousarray(tensor.transpose(2, 0, 1)[None, :, :, :])
|
| sr = np.clip(session.run(None, {input_name: tensor})[0][0], 0.0, 1.0)
|
| sr = cv2.cvtColor((sr.transpose(1, 2, 0) * 255.0).round().astype(np.uint8),
|
| cv2.COLOR_RGB2BGR)
|
|
|
|
|
| full_h, full_w = int(round((y1 - y0) * scale)), int(round((x1 - x0) * scale))
|
| sr = cv2.resize(sr, (full_w, full_h), interpolation=interp)
|
| th, tw = int(round((bottom - top) * scale)), int(round((right - left) * scale))
|
| oy = min(int(round((top - y0) * scale)), max(0, full_h - th))
|
| ox = min(int(round((left - x0) * scale)), max(0, full_w - tw))
|
| th = min(th, full_h - oy)
|
| tw = min(tw, full_w - ox)
|
| oy_px, ox_px = int(round(top * scale)), int(round(left * scale))
|
| th = min(th, out_h - oy_px)
|
| tw = min(tw, out_w - ox_px)
|
| canvas[oy_px:oy_px + th, ox_px:ox_px + tw] = sr[oy:oy + th, ox:ox + tw]
|
| tiles += 1
|
| print(f" tile {tiles}: LR {y1 - y0}x{x1 - x0} -> {th}x{tw} "
|
| f"[{time.time() - t0:.1f}s]", flush=True)
|
| return canvas
|
|
|
|
|
| def psnr(a: np.ndarray, b: np.ndarray) -> float:
|
| mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
| return 99.0 if mse == 0 else float(10.0 * np.log10(255.0 ** 2 / mse))
|
|
|
|
|
| def main() -> None:
|
| ap = argparse.ArgumentParser(description=__doc__,
|
| formatter_class=argparse.RawDescriptionHelpFormatter)
|
| ap.add_argument("input", help="immagine da ingrandire")
|
| ap.add_argument("--onnx", default="student_sr9x.onnx", help="ONNX dello Student distillato")
|
| ap.add_argument("--scale", type=float, default=6.0, help="scala richiesta (il modello e` 9x)")
|
| ap.add_argument("--output", default=None, help="percorso output (default output_<nome>_<scala>x.png)")
|
| ap.add_argument("--tile", type=int, default=64, help="lato tile in pixel LR (0 = immagine intera)")
|
| ap.add_argument("--pad", type=int, default=6, help="overlap di padding in pixel LR")
|
| ap.add_argument("--alpha-interp", choices=["nearest", "linear"], default="nearest",
|
| help="ricampionamento dell'alpha (nearest evita aloni sui bordi)")
|
| ap.add_argument("--no-baseline", action="store_true", help="non generare il confronto bicubico")
|
| ap.add_argument("--providers", choices=["auto", "cpu"], default="auto")
|
| args = ap.parse_args()
|
|
|
| src = Path(args.input)
|
| if not src.is_file():
|
| raise SystemExit(f"File non trovato: {src}")
|
| onnx_path = Path(args.onnx)
|
| if not onnx_path.is_file():
|
| raise SystemExit(f"ONNX non trovato: {onnx_path}")
|
| if args.scale <= 0:
|
| raise SystemExit("--scale deve essere > 0")
|
|
|
| session = build_session(onnx_path, args.providers)
|
| native = native_scale(session)
|
| bgr, alpha = load_image(src)
|
| height, width = bgr.shape[:2]
|
| out_h, out_w = int(round(height * args.scale)), int(round(width * args.scale))
|
| tile = height if args.tile <= 0 else min(int(args.tile), height)
|
| print(f"[modello] {onnx_path} | scala nativa {native}x | tile {tile}+pad {args.pad}")
|
| print(f"[input] {src} {width}x{height}"
|
| + (" (con alpha)" if alpha is not None else ""))
|
| print(f"[output] {out_w}x{out_h} ({args.scale}x)")
|
|
|
| result = upscale_tiled(session, bgr, args.scale, native, tile, int(args.pad))
|
|
|
| out_path = Path(args.output) if args.output else src.with_name(
|
| f"output_{src.stem}_{args.scale:g}x.png")
|
| if alpha is not None:
|
| interp = cv2.INTER_NEAREST if args.alpha_interp == "nearest" else cv2.INTER_LINEAR
|
| alpha_out = cv2.resize(alpha, (out_w, out_h), interpolation=interp)
|
| result = np.dstack([result, alpha_out])
|
| if not cv2.imwrite(str(out_path), result):
|
| raise SystemExit(f"Scrittura fallita: {out_path}")
|
| print(f"[salvato] {out_path}")
|
|
|
| if not args.no_baseline:
|
| interp = cv2.INTER_AREA if native > args.scale else cv2.INTER_CUBIC
|
| bicubic = cv2.resize(bgr, (out_w, out_h), interpolation=interp)
|
| bic_path = out_path.with_name(f"{out_path.stem}_bicubic{out_path.suffix}")
|
| if alpha is not None:
|
| bicubic = np.dstack([bicubic, cv2.resize(alpha, (out_w, out_h),
|
| interpolation=cv2.INTER_NEAREST)])
|
| cv2.imwrite(str(bic_path), bicubic)
|
| print(f"[baseline] {bic_path} | PSNR modello vs bicubica "
|
| f"{psnr(result if result.shape[2] == 3 else result[:, :, :3], bicubic[:, :, :3]):.2f} dB")
|
|
|
|
|
| if __name__ == "__main__":
|
| main() |
