Spaces:
Sleeping
Sleeping
| import os | |
| import argparse | |
| from pathlib import Path | |
| import numpy as np | |
| from PIL import Image | |
| def _find_lr_file(hr_path: Path, lr_dir: Path, scale: int) -> Path: | |
| """ | |
| Locate the LR counterpart of an HR file. | |
| DIV2K names LR files as <stem>x<scale>.png (e.g. 0801x2.png). | |
| Falls back to the plain stem name if the suffixed version is absent. | |
| """ | |
| suffixed = lr_dir / f"{hr_path.stem}x{scale}{hr_path.suffix}" | |
| if suffixed.exists(): | |
| return suffixed | |
| for ext in [".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"]: | |
| alt = lr_dir / f"{hr_path.stem}x{scale}{ext}" | |
| if alt.exists(): | |
| return alt | |
| plain = lr_dir / hr_path.name | |
| if plain.exists(): | |
| return plain | |
| return None | |
| def create_lr_images(hr_dir: str, lr_dir: str, scale_factor: int = 2) -> None: | |
| """Downsample HR images to create LR counterparts via bicubic interpolation.""" | |
| hr_path = Path(hr_dir) | |
| lr_path = Path(lr_dir) | |
| lr_path.mkdir(parents=True, exist_ok=True) | |
| exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff'} | |
| hr_images = sorted([f for f in hr_path.iterdir() if f.suffix.lower() in exts]) | |
| if not hr_images: | |
| print(f"No images found in {hr_dir}") | |
| return | |
| print(f"Creating LR images (x{scale_factor} downscale) from {len(hr_images)} HR images...") | |
| for img_path in hr_images: | |
| hr_img = Image.open(img_path).convert('RGB') | |
| w, h = hr_img.size | |
| lr_img = hr_img.resize((w // scale_factor, h // scale_factor), Image.BICUBIC) | |
| lr_img.save(lr_path / img_path.name) | |
| print(f"Saved {len(hr_images)} LR images to {lr_dir}") | |
| def extract_patches( | |
| hr_dir: str, | |
| lr_dir: str = None, | |
| patch_size: int = 64, | |
| scale_factor: int = 2, | |
| stride: int = 64, | |
| max_patches: int = None, | |
| num_images: int = None, | |
| ) -> tuple: | |
| """ | |
| Extract aligned HR / LR patch pairs. | |
| If lr_dir is None, LR patches are created on-the-fly by downsampling | |
| each HR patch. | |
| patch_size is the HR patch size. | |
| LR patch size is patch_size // scale_factor. | |
| """ | |
| hr_path = Path(hr_dir) | |
| lr_path = Path(lr_dir) if lr_dir else None | |
| exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff'} | |
| hr_images = sorted([f for f in hr_path.iterdir() if f.suffix.lower() in exts]) | |
| if num_images: | |
| hr_images = hr_images[:num_images] | |
| lr_patch_size = patch_size // scale_factor | |
| hr_patches, lr_patches = [], [] | |
| for img_path in hr_images: | |
| lr_file = _find_lr_file(img_path, lr_path, scale_factor) if lr_path else None | |
| hr_img = np.array(Image.open(img_path).convert('RGB'), dtype=np.float32) / 255.0 | |
| if lr_file: | |
| lr_img = np.array(Image.open(lr_file).convert('RGB'), dtype=np.float32) / 255.0 | |
| else: | |
| pil_hr = Image.fromarray((hr_img * 255).astype(np.uint8)) | |
| w, h = pil_hr.size | |
| lr_img = np.array( | |
| pil_hr.resize((w // scale_factor, h // scale_factor), Image.BICUBIC), | |
| dtype=np.float32, | |
| ) / 255.0 | |
| h_hr, w_hr = hr_img.shape[:2] | |
| for y in range(0, h_hr - patch_size + 1, stride): | |
| for x in range(0, w_hr - patch_size + 1, stride): | |
| hr_patches.append(hr_img[y:y + patch_size, x:x + patch_size]) | |
| ly, lx = y // scale_factor, x // scale_factor | |
| lr_patches.append(lr_img[ly:ly + lr_patch_size, lx:lx + lr_patch_size]) | |
| if max_patches and len(hr_patches) >= max_patches: | |
| break | |
| if max_patches and len(hr_patches) >= max_patches: | |
| break | |
| if max_patches and len(hr_patches) >= max_patches: | |
| break | |
| hr_arr = np.array(hr_patches, dtype=np.float32) | |
| lr_arr = np.array(lr_patches, dtype=np.float32) | |
| print(f"Extracted {len(hr_arr)} patch pairs | HR: {hr_arr.shape} | LR: {lr_arr.shape}") | |
| return hr_arr, lr_arr | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser( | |
| description='Prepare dynamic LR-HR pairs and extract SR training patches' | |
| ) | |
| parser.add_argument('--scale', type=int, default=2, choices=[2, 3, 4, 8]) | |
| parser.add_argument('--train_hr_dir', type=str, default='Data/DIV2K_train_HR') | |
| parser.add_argument('--valid_hr_dir', type=str, default='Data/DIV2K_valid_HR') | |
| parser.add_argument('--train_lr_root', type=str, default='Data/DIV2K_train_LR_bicubic') | |
| parser.add_argument('--valid_lr_root', type=str, default='Data/DIV2K_valid_LR_bicubic') | |
| parser.add_argument('--hr_dir', type=str, default=None, | |
| help='Optional override for HR folder') | |
| parser.add_argument('--lr_dir', type=str, default=None, | |
| help='Optional override for LR folder') | |
| parser.add_argument('--use_valid_as_main', action='store_true', | |
| help='Use validation directories instead of training directories') | |
| parser.add_argument('--out_dir', type=str, default='Data', | |
| help='Directory to save patch arrays') | |
| parser.add_argument('--hr_out_name', type=str, default=None, | |
| help='Optional output filename for HR patches') | |
| parser.add_argument('--lr_out_name', type=str, default=None, | |
| help='Optional output filename for LR patches') | |
| parser.add_argument('--patch_size', type=int, default=64, | |
| help='HR patch side length in pixels') | |
| parser.add_argument('--stride', type=int, default=64, | |
| help='Patch extraction stride in HR pixels') | |
| parser.add_argument('--max_patches', type=int, default=None, | |
| help='Cap on total patches') | |
| parser.add_argument('--num_images', type=int, default=80, | |
| help='Number of HR images to use') | |
| parser.add_argument('--create_lr', action='store_true', | |
| help='Create bicubic LR images if needed') | |
| args = parser.parse_args() | |
| if args.hr_dir is None or args.lr_dir is None: | |
| if args.use_valid_as_main: | |
| if args.hr_dir is None: | |
| args.hr_dir = args.valid_hr_dir | |
| if args.lr_dir is None: | |
| args.lr_dir = f'{args.valid_lr_root}/X{args.scale}' | |
| else: | |
| if args.hr_dir is None: | |
| args.hr_dir = args.train_hr_dir | |
| if args.lr_dir is None: | |
| args.lr_dir = f'{args.train_lr_root}/X{args.scale}' | |
| if args.hr_out_name is None: | |
| args.hr_out_name = f'hr_patches_x{args.scale}.npy' | |
| if args.lr_out_name is None: | |
| args.lr_out_name = f'lr_patches_x{args.scale}.npy' | |
| print(f"Scale: {args.scale}") | |
| print(f"HR dir: {args.hr_dir}") | |
| print(f"LR dir: {args.lr_dir}") | |
| print(f"HR output: {os.path.join(args.out_dir, args.hr_out_name)}") | |
| print(f"LR output: {os.path.join(args.out_dir, args.lr_out_name)}") | |
| if args.create_lr: | |
| create_lr_images(args.hr_dir, args.lr_dir, scale_factor=args.scale) | |
| hr_arr, lr_arr = extract_patches( | |
| hr_dir=args.hr_dir, | |
| lr_dir=args.lr_dir, | |
| patch_size=args.patch_size, | |
| scale_factor=args.scale, | |
| stride=args.stride, | |
| max_patches=args.max_patches, | |
| num_images=args.num_images, | |
| ) | |
| os.makedirs(args.out_dir, exist_ok=True) | |
| np.save(os.path.join(args.out_dir, args.hr_out_name), hr_arr) | |
| np.save(os.path.join(args.out_dir, args.lr_out_name), lr_arr) | |
| print(f"Saved {os.path.join(args.out_dir, args.hr_out_name)}") | |
| print(f"Saved {os.path.join(args.out_dir, args.lr_out_name)}") | |