Spaces:
Sleeping
Sleeping
| import argparse | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import tensorflow as tf | |
| from PIL import Image | |
| sys.path.insert(0, str(Path(__file__).parent)) | |
| from model import PixelShuffle | |
| def upscale( | |
| input_path, | |
| checkpoint_path=None, | |
| output_path=None, | |
| comparison_path=None, | |
| scale=2, | |
| arch='espcn', | |
| checkpoint_dir='outputs/checkpoints', | |
| ): | |
| """ | |
| Upscale a single image using a trained ESPCN/SRCNN model. | |
| Parameters | |
| ---------- | |
| input_path : path to the low-resolution input image | |
| checkpoint_path : path to the .keras model file | |
| if None, auto-loads from checkpoint_dir using arch + scale | |
| output_path : where to save the SR image | |
| comparison_path : where to save a side-by-side comparison image | |
| scale : upscale factor (2, 3, 4, or 8) | |
| arch : model architecture — 'espcn' or 'srcnn' | |
| checkpoint_dir : folder containing trained checkpoints | |
| """ | |
| input_path = Path(input_path) | |
| if checkpoint_path is None: | |
| checkpoint_path = Path(checkpoint_dir) / f"{arch}_x{scale}_best.keras" | |
| else: | |
| checkpoint_path = Path(checkpoint_path) | |
| if not checkpoint_path.exists(): | |
| raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") | |
| model = tf.keras.models.load_model( | |
| checkpoint_path, | |
| custom_objects={'PixelShuffle': PixelShuffle}, | |
| compile=False, | |
| ) | |
| print(f"Loaded : {checkpoint_path} ({arch.upper()}, x{scale})") | |
| lr_img = Image.open(input_path).convert('RGB') | |
| orig_w, orig_h = lr_img.size | |
| print(f"Input : {input_path} ({orig_w}×{orig_h})") | |
| lr_arr = np.array(lr_img, dtype=np.float32) / 255.0 | |
| lr_tensor = lr_arr[np.newaxis, ...] # (1, H, W, 3) | |
| if arch == 'srcnn': | |
| lr_tensor = tf.image.resize( | |
| lr_tensor, | |
| [orig_h * scale, orig_w * scale], | |
| method='bicubic' | |
| ).numpy() | |
| lr_tensor = np.clip(lr_tensor, 0.0, 1.0) | |
| sr_arr = model.predict(lr_tensor, verbose=0)[0] | |
| sr_arr = np.clip(sr_arr, 0.0, 1.0) | |
| sr_img = Image.fromarray((sr_arr * 255).astype(np.uint8)) | |
| sr_w, sr_h = sr_img.size | |
| input_stem = input_path.stem | |
| input_ext = input_path.suffix if input_path.suffix else '.png' | |
| if output_path is None: | |
| output_path = input_path.parent / f"{input_stem}_{arch}_x{scale}_sr{input_ext}" | |
| else: | |
| output_path = Path(output_path) | |
| sr_img.save(output_path) | |
| print(f"SR image: {output_path} ({sr_w}×{sr_h})") | |
| # Bicubic baseline for comparison | |
| bicubic_img = lr_img.resize((orig_w * scale, orig_h * scale), Image.BICUBIC) | |
| # Resize LR preview so all three panels match in displayed size | |
| lr_preview = lr_img.resize((orig_w * scale, orig_h * scale), Image.NEAREST) | |
| padding = 12 | |
| label_space = 28 | |
| panel_w = lr_preview.width | |
| panel_h = lr_preview.height | |
| comp_w = panel_w * 3 + padding * 4 | |
| comp_h = panel_h + padding * 2 + label_space | |
| comparison = Image.new('RGB', (comp_w, comp_h), (30, 30, 30)) | |
| x = padding | |
| y = padding | |
| comparison.paste(lr_preview, (x, y)) | |
| x += panel_w + padding | |
| comparison.paste(bicubic_img, (x, y)) | |
| x += panel_w + padding | |
| comparison.paste(sr_img, (x, y)) | |
| try: | |
| from PIL import ImageDraw | |
| draw = ImageDraw.Draw(comparison) | |
| labels = [ | |
| (padding, panel_h + padding + 6, 'LR Input'), | |
| (panel_w + padding * 2, panel_h + padding + 6, 'Bicubic'), | |
| (panel_w * 2 + padding * 3, panel_h + padding + 6, f'{arch.upper()} x{scale}'), | |
| ] | |
| for lx, ly, text in labels: | |
| draw.text((lx, ly), text, fill=(220, 220, 220)) | |
| except Exception as e: | |
| print(f"Warning: could not draw labels on comparison image: {e}") | |
| if comparison_path is None: | |
| comparison_path = input_path.parent / f"{input_stem}_{arch}_x{scale}_compare{input_ext}" | |
| else: | |
| comparison_path = Path(comparison_path) | |
| comparison.save(comparison_path) | |
| print(f"Compare : {comparison_path} (LR | Bicubic | {arch.upper()} x{scale})") | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser( | |
| description='Upscale a single image with a trained ESPCN/SRCNN model' | |
| ) | |
| parser.add_argument( | |
| '--input', | |
| type=str, | |
| required=True, | |
| help='Path to the low-resolution input image' | |
| ) | |
| parser.add_argument( | |
| '--checkpoint', | |
| type=str, | |
| default=None, | |
| help='Optional path to .keras model file; if omitted, uses outputs/checkpoints/{arch}_x{scale}_best.keras' | |
| ) | |
| parser.add_argument( | |
| '--checkpoint_dir', | |
| type=str, | |
| default='outputs/checkpoints', | |
| help='Directory containing trained checkpoints' | |
| ) | |
| parser.add_argument( | |
| '--output', | |
| type=str, | |
| default=None, | |
| help='Output path for the SR image' | |
| ) | |
| parser.add_argument( | |
| '--comparison', | |
| type=str, | |
| default=None, | |
| help='Output path for the side-by-side comparison image' | |
| ) | |
| parser.add_argument( | |
| '--scale', | |
| type=int, | |
| default=2, | |
| choices=[2, 3, 4, 8], | |
| help='Upscale factor' | |
| ) | |
| parser.add_argument( | |
| '--arch', | |
| type=str, | |
| default='espcn', | |
| choices=['espcn', 'srcnn'], | |
| help='Model architecture' | |
| ) | |
| args = parser.parse_args() | |
| upscale( | |
| input_path=args.input, | |
| checkpoint_path=args.checkpoint, | |
| output_path=args.output, | |
| comparison_path=args.comparison, | |
| scale=args.scale, | |
| arch=args.arch, | |
| checkpoint_dir=args.checkpoint_dir, | |
| ) | |