Spaces:
Running
Running
File size: 7,573 Bytes
c4a8c0e | 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 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 | 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)}")
|