Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import argparse | |
| 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 build_model, PixelShuffle | |
| class SRDataset: | |
| """ | |
| Super-resolution dataset. | |
| Supports two data sources (in priority order): | |
| 1. Pre-extracted NumPy patch arrays | |
| 2. Image directories | |
| """ | |
| def __init__( | |
| self, | |
| hr_patches_path=None, | |
| lr_patches_path=None, | |
| hr_dir=None, | |
| lr_dir=None, | |
| scale_factor=2, | |
| architecture='espcn', | |
| ): | |
| self.architecture = architecture | |
| self.scale_factor = scale_factor | |
| if hr_patches_path and lr_patches_path: | |
| self.hr_data = np.load(hr_patches_path).astype(np.float32) | |
| self.lr_data = np.load(lr_patches_path).astype(np.float32) | |
| elif hr_dir and lr_dir: | |
| self.hr_data, self.lr_data = self._load_from_dirs(hr_dir, lr_dir) | |
| else: | |
| raise ValueError("Provide either (hr_patches_path, lr_patches_path) or (hr_dir, lr_dir).") | |
| def _load_from_dirs(self, hr_dir, lr_dir): | |
| exts = {'.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff'} | |
| hr_files = sorted([f for f in Path(hr_dir).iterdir() if f.suffix.lower() in exts]) | |
| lr_files = sorted([f for f in Path(lr_dir).iterdir() if f.suffix.lower() in exts]) | |
| assert len(hr_files) == len(lr_files), ( | |
| f"HR and LR directories must have the same number of images " | |
| f"({len(hr_files)} vs {len(lr_files)})" | |
| ) | |
| hr_data, lr_data = [], [] | |
| for hf, lf in zip(hr_files, lr_files): | |
| hr_data.append(np.array(Image.open(hf).convert('RGB'), dtype=np.float32) / 255.0) | |
| lr_data.append(np.array(Image.open(lf).convert('RGB'), dtype=np.float32) / 255.0) | |
| return np.array(hr_data), np.array(lr_data) | |
| def get_tf_datasets(self, batch_size=16, val_split=0.1): | |
| """ | |
| Split into train/val and return tf.data.Dataset objects. | |
| For SRCNN, the LR images are bicubic-upsampled to match HR size. | |
| """ | |
| n = len(self.hr_data) | |
| val_size = max(1, int(n * val_split)) | |
| train_size = n - val_size | |
| lr_input = self.lr_data | |
| if self.architecture == 'srcnn': | |
| hr_h, hr_w = self.hr_data.shape[1], self.hr_data.shape[2] | |
| lr_input = tf.image.resize( | |
| self.lr_data, [hr_h, hr_w], method='bicubic' | |
| ).numpy().astype(np.float32) | |
| lr_input = np.clip(lr_input, 0.0, 1.0) | |
| print(f"Samples — train: {train_size} val: {val_size}") | |
| train_ds = ( | |
| tf.data.Dataset.from_tensor_slices((lr_input[:train_size], self.hr_data[:train_size])) | |
| .shuffle(buffer_size=1000, seed=42) | |
| .batch(batch_size) | |
| .prefetch(tf.data.AUTOTUNE) | |
| ) | |
| val_ds = ( | |
| tf.data.Dataset.from_tensor_slices((lr_input[train_size:], self.hr_data[train_size:])) | |
| .batch(batch_size) | |
| .prefetch(tf.data.AUTOTUNE) | |
| ) | |
| return train_ds, val_ds | |
| def train(args): | |
| use_patches = args.use_patches | |
| if use_patches: | |
| patch_hr = Path(args.hr_patches_path) | |
| patch_lr = Path(args.lr_patches_path) | |
| if not patch_hr.exists() or not patch_lr.exists(): | |
| raise FileNotFoundError( | |
| f"Patch files not found:\nHR: {patch_hr}\nLR: {patch_lr}" | |
| ) | |
| print("Loading pre-extracted patch arrays...") | |
| dataset = SRDataset( | |
| hr_patches_path=str(patch_hr), | |
| lr_patches_path=str(patch_lr), | |
| scale_factor=args.scale, | |
| architecture=args.arch, | |
| ) | |
| else: | |
| print("Loading images from directories...") | |
| print(f"HR dir: {args.hr_dir}") | |
| print(f"LR dir: {args.lr_dir}") | |
| dataset = SRDataset( | |
| hr_dir=args.hr_dir, | |
| lr_dir=args.lr_dir, | |
| scale_factor=args.scale, | |
| architecture=args.arch, | |
| ) | |
| train_ds, val_ds = dataset.get_tf_datasets(batch_size=args.batch_size) | |
| resume_path = args.resume | |
| if resume_path and Path(resume_path).exists(): | |
| print(f"Resuming from checkpoint: {resume_path}") | |
| model = tf.keras.models.load_model( | |
| resume_path, | |
| custom_objects={'PixelShuffle': PixelShuffle}, | |
| compile=False, | |
| ) | |
| else: | |
| model = build_model(args.arch, scale_factor=args.scale) | |
| model.summary() | |
| model.compile( | |
| optimizer=tf.keras.optimizers.Adam(learning_rate=args.lr), | |
| loss=tf.keras.losses.MeanAbsoluteError(), | |
| metrics=[tf.keras.metrics.MeanSquaredError(name='mse')], | |
| ) | |
| os.makedirs(args.checkpoint_dir, exist_ok=True) | |
| best_ckpt = os.path.join( | |
| args.checkpoint_dir, f'{args.arch}_x{args.scale}_best.keras' | |
| ) | |
| final_ckpt = os.path.join( | |
| args.checkpoint_dir, f'{args.arch}_x{args.scale}_final.keras' | |
| ) | |
| callbacks = [ | |
| tf.keras.callbacks.ModelCheckpoint( | |
| best_ckpt, | |
| monitor='val_loss', | |
| save_best_only=True, | |
| verbose=1, | |
| ), | |
| tf.keras.callbacks.ReduceLROnPlateau( | |
| monitor='val_loss', | |
| factor=0.5, | |
| patience=5, | |
| verbose=1, | |
| ), | |
| ] | |
| print(f"\nTraining {args.arch.upper()} x{args.scale} for up to {args.epochs} epochs\n{'-' * 60}") | |
| history = model.fit( | |
| train_ds, | |
| epochs=args.epochs, | |
| validation_data=val_ds, | |
| callbacks=callbacks, | |
| verbose=1, | |
| ) | |
| model.save(final_ckpt) | |
| print(f"\nDone. Best model saved: {best_ckpt}") | |
| print(f"Done. Final model saved: {final_ckpt}") | |
| return history | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser(description='Train super-resolution ANN') | |
| parser.add_argument('--arch', type=str, default='espcn', choices=['srcnn', 'espcn']) | |
| parser.add_argument('--scale', type=int, default=2, choices=[2, 3, 4, 8]) | |
| # Directory-based loading | |
| 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) | |
| parser.add_argument('--lr_dir', type=str, default=None) | |
| parser.add_argument('--use_valid_as_main', action='store_true') | |
| # Optional patch-based loading | |
| parser.add_argument('--use_patches', action='store_true') | |
| parser.add_argument('--hr_patches_path', type=str, default=None) | |
| parser.add_argument('--lr_patches_path', type=str, default=None) | |
| # Training params | |
| parser.add_argument('--epochs', type=int, default=50) | |
| parser.add_argument('--batch_size', type=int, default=16) | |
| parser.add_argument('--lr', type=float, default=1e-4) | |
| parser.add_argument('--checkpoint_dir', type=str, default='outputs/checkpoints') | |
| parser.add_argument('--resume', type=str, default=None, | |
| help='Path to a .keras checkpoint to resume training from') | |
| args = parser.parse_args() | |
| # Auto-build patch paths if patch mode is enabled | |
| if args.use_patches: | |
| if args.hr_patches_path is None: | |
| args.hr_patches_path = f'Data/hr_patches_x{args.scale}.npy' | |
| if args.lr_patches_path is None: | |
| args.lr_patches_path = f'Data/lr_patches_x{args.scale}.npy' | |
| else: | |
| # Auto-build image directory paths | |
| 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}' | |
| train(args) | |