File size: 6,328 Bytes
b4ed648 | 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 | """Minimal inference example for ESRGAN 4x Super-Resolution using ExecuTorch.
Loads a quantized .pte model and runs super-resolution inference on a single image.
The model upscales the input image by 4x. For images larger than 128x128, the
input is automatically split into overlapping tiles, each tile is super-resolved,
and the results are blended and stitched into the final output.
"""
import json
from pathlib import Path
import numpy as np
import torch
from executorch.runtime import Runtime
from PIL import Image
from torchvision import transforms
# ── Configuration ──────────────────────────────────────────────────────────────
MODEL_PATH = "esrgan-x4-int8-executorch.pte"
IMAGE_PATH = "sample_input.jpg"
TILE_SIZE = 128 # model was exported with 128x128 input tiles
SCALE = 4 # 4x upscaling factor
TILE_OVERLAP = 8 # pixel overlap between tiles for seamless blending
# ── Model Loading ──────────────────────────────────────────────────────────────
def load_model(pte_path: str):
"""Load ExecuTorch .pte model and return the forward method."""
runtime = Runtime.get()
program = runtime.load_program(pte_path)
return program.load_method("forward")
# ── Preprocessing ──────────────────────────────────────────────────────────────
def preprocess(image_path: str) -> tuple[torch.Tensor, tuple[int, int]]:
"""Load and preprocess image for model input.
The model expects [0, 1] float32 RGB tensors — no normalization is applied.
Returns the input tensor and the original image size for reference.
"""
image = Image.open(image_path).convert("RGB")
original_size = (image.width, image.height)
to_tensor = transforms.ToTensor() # converts [0-255] uint8 -> [0, 1] float32
input_tensor = to_tensor(image).unsqueeze(0) # [1, 3, H, W]
return input_tensor, original_size
# ── Tiled Inference ────────────────────────────────────────────────────────────
def run_tiled_inference(method, input_tensor: torch.Tensor) -> torch.Tensor:
"""Run super-resolution inference with overlapping tile stitching.
Splits the input into overlapping 128x128 tiles, runs each through the
model, and blends overlapping regions using pixel-level averaging.
"""
_, _, h, w = input_tensor.shape
out_h, out_w = h * SCALE, w * SCALE
output = torch.zeros(1, 3, out_h, out_w)
weights = torch.zeros(1, 1, out_h, out_w)
stride = TILE_SIZE - TILE_OVERLAP
y_positions = list(range(0, max(1, h - TILE_SIZE + 1), stride))
if not y_positions or y_positions[-1] + TILE_SIZE < h:
y_positions.append(max(0, h - TILE_SIZE))
x_positions = list(range(0, max(1, w - TILE_SIZE + 1), stride))
if not x_positions or x_positions[-1] + TILE_SIZE < w:
x_positions.append(max(0, w - TILE_SIZE))
for y in y_positions:
for x in x_positions:
tile = input_tensor[:, :, y:y + TILE_SIZE, x:x + TILE_SIZE].contiguous()
sr_tile = method.execute([tile])[0]
oy, ox = y * SCALE, x * SCALE
oh, ow = TILE_SIZE * SCALE, TILE_SIZE * SCALE
output[:, :, oy:oy + oh, ox:ox + ow] += sr_tile
weights[:, :, oy:oy + oh, ox:ox + ow] += 1.0
return output / weights.clamp(min=1.0)
# ── Postprocessing ─────────────────────────────────────────────────────────────
def postprocess(raw_output: torch.Tensor) -> np.ndarray:
"""Clamp output to [0, 1] and convert to uint8 numpy array (H, W, 3)."""
sr_tensor = raw_output.clamp(0.0, 1.0).squeeze(0) # [3, H*4, W*4]
sr_array = (sr_tensor.permute(1, 2, 0).numpy() * 255.0).round().astype(np.uint8)
return sr_array
# ── Save Results ──────────────────────────────────────────────────────────────
def save_results(sr_array: np.ndarray, original_size: tuple[int, int]) -> None:
"""Save the super-resolved image and a JSON summary to the script directory."""
script_dir = Path(__file__).parent
# Save super-resolved image
output_image = Image.fromarray(sr_array)
output_path = script_dir / "sample_output.png"
output_image.save(output_path)
print(f"Super-resolved image saved to: {output_path}")
# Save JSON summary
summary = {
"input_size": {"width": original_size[0], "height": original_size[1]},
"output_size": {"width": sr_array.shape[1], "height": sr_array.shape[0]},
"scale_factor": SCALE,
}
json_path = script_dir / "super_resolution.json"
with open(json_path, "w") as f:
json.dump(summary, f, indent=2)
print(f"Summary saved to: {json_path}")
# ── Main ───────────────────────────────────────────────────────────────────────
def main():
script_dir = Path(__file__).parent
model_path = script_dir / MODEL_PATH
image_path = script_dir / IMAGE_PATH
print(f"Loading model from: {model_path}")
method = load_model(str(model_path))
print(f"Preprocessing image: {image_path}")
input_tensor, original_size = preprocess(str(image_path))
_, _, h, w = input_tensor.shape
print(f" Input size: {w}x{h} -> output will be {w * SCALE}x{h * SCALE}")
print("Running tiled super-resolution inference...")
raw_output = run_tiled_inference(method, input_tensor)
sr_array = postprocess(raw_output)
print(f" Output shape: {sr_array.shape[1]}x{sr_array.shape[0]} (WxH)")
save_results(sr_array, original_size)
if __name__ == "__main__":
main()
|