Реализация AI-увеличения разрешения изображений (Super Resolution / Upscale)

AI Super-Resolution — апскейл изображений

Бикубическая интерполяция даёт 4x апскейл с размытием. AI super-resolution восстанавливает детали: текстуры кожи, текст на вывесках, структуру ткани. Разница видна при сравнении PSNR: бикубик — 28–30 dB, Real-ESRGAN — 32–36 dB на фотографиях.

Real-ESRGAN — практический стандарт

import torch
import numpy as np
from PIL import Image
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

def upscale_image(
    image_path: str,
    scale: int = 4,
    model_name: str = 'RealESRGAN_x4plus',  # или 'RealESRGAN_x4plus_anime_6B'
    tile_size: int = 512,    # для больших изображений — обработка тайлами
    half_precision: bool = True
) -> np.ndarray:
    """
    tile_size=512 при VRAM 6GB, tile_size=0 (whole image) при VRAM 24GB.
    half=True — FP16, экономит ~50% VRAM.
    """
    model = RRDBNet(
        num_in_ch=3, num_out_ch=3,
        num_feat=64, num_block=23, num_grow_ch=32,
        scale=scale
    )
    upsampler = RealESRGANer(
        scale=scale,
        model_path=f'weights/{model_name}.pth',
        model=model,
        tile=tile_size,
        tile_pad=10,      # перекрытие тайлов для сглаживания швов
        pre_pad=0,
        half=half_precision,
        device='cuda'
    )

    img = np.array(Image.open(image_path).convert('RGB'))
    output, _ = upsampler.enhance(img, outscale=scale)
    return output

GFPGAN для восстановления лиц

Real-ESRGAN на портретах иногда создаёт артефакты на лице. GFPGAN добавляет face restoration поверх SR:

from gfpgan import GFPGANer

def restore_face_photo(
    degraded_image: np.ndarray,
    upscale: int = 2,
    arch: str = 'clean',         # 'clean' | 'RestoreFormer'
    channel_multiplier: int = 2,
    weight: float = 0.5          # 0=чистый GFPGAN, 1=без face enhancement
) -> np.ndarray:
    """
    weight=0.5 — компромисс между восстановлением и сохранением
    индивидуальных черт. При weight=0 лица «глянцевые».
    """
    restorer = GFPGANer(
        model_path='weights/GFPGANv1.4.pth',
        upscale=upscale,
        arch=arch,
        channel_multiplier=channel_multiplier,
        bg_upsampler=None   # можно передать RealESRGANer для фона
    )

    _, _, restored = restorer.enhance(
        degraded_image,
        has_aligned=False,
        only_center_face=False,
        paste_back=True,
        weight=weight
    )
    return restored

Метрики и сравнение моделей

PSNR — не единственный критерий: человеческое восприятие коррелирует с LPIPS (perceptual loss). Real-ESRGAN при PSNR ниже SwinIR часто выглядит лучше субъективно из-за более высокочастотных деталей.

Батчевая обработка больших объёмов

from pathlib import Path
import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

class ImageDataset(Dataset):
    def __init__(self, image_paths: list[str], size: int = 256):
        self.paths = image_paths
        self.transform = transforms.Compose([
            transforms.Resize((size, size)),
            transforms.ToTensor()
        ])

    def __len__(self): return len(self.paths)

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert('RGB')
        return self.transform(img), self.paths[idx]

def batch_upscale_pipeline(
    input_dir: str,
    output_dir: str,
    batch_size: int = 4,   # при VRAM 12GB и tile_size=0
    scale: int = 4
):
    paths = list(Path(input_dir).glob('*.{jpg,jpeg,png}'))
    Path(output_dir).mkdir(exist_ok=True)

    # Для батч-инференса используем прямой forward
    # (RealESRGANer не поддерживает батчи, нужен прямой вызов модели)
    model = RRDBNet(
        num_in_ch=3, num_out_ch=3,
        num_feat=64, num_block=23, num_grow_ch=32, scale=scale
    )
    model.load_state_dict(
        torch.load(f'weights/RealESRGAN_x4plus.pth')['params_ema']
    )
    model.eval().cuda().half()

    for path in paths:
        with torch.no_grad(), torch.cuda.amp.autocast():
            img_t = transforms.ToTensor()(
                Image.open(path).convert('RGB')
            ).unsqueeze(0).half().cuda()
            out = model(img_t).squeeze(0).float().cpu()
            out_img = transforms.ToPILImage()(out.clamp(0, 1))
            out_img.save(
                Path(output_dir) / (Path(path).stem + '_4x.png')
            )

Ограничения и типичные проблемы

• Галлюцинации текстур — Real-ESRGAN может добавить несуществующий текст на вывесках. На forensics-применениях это недопустимо
• OOM на больших изображениях — 12-мегапиксельное фото при 4x апскейл = 192Мп, не лезет в память целиком. Решение: tile_size=512 с tile_pad=10
• JPEG-артефакты — блочность артефактов JPEG усиливается SR. Предобработка: JPEG-aware денойзинг (nf_denoise из BasicSR)

Сроки