Реалізація AI-збільшення роздільної здатності відео (Video Upscale)

ШІ Super-Resolution для відео — апскейл відеоконтенту

Апскейл відео складніше апскейлу зображень: потрібна часова узгодженість — сусідні кадри мають виглядати узгоджено, інакше результат миготить. Просто застосувати Real-ESRGAN до кожного кадру — неправильно: шум на однорідних поверхнях буде змінюватися від кадру до кадру.

Real-BasicVSR та BasicVSR++ — основні моделі

import torch
import numpy as np
import cv2
from basicsr.archs.basicvsrpp_arch import BasicVSRPlusPlus

def upscale_video_basicvsr(
    frames: list[np.ndarray],   # список кадрів (H, W, 3) BGR
    scale: int = 4,
    num_feat: int = 64,
    num_propagation_blocks: int = 7,
    cpu_cache_length: int = 100  # кадри в пам'яті GPU одночасно
) -> list[np.ndarray]:
    """
    BasicVSR++ використовує bidirectional propagation:
    інформацію з минулих І майбутніх кадрів.
    cpu_cache_length: для довгих відео виганяємо частину кадрів на CPU.
    """
    model = BasicVSRPlusPlus(
        mid_channels=num_feat,
        num_blocks=num_propagation_blocks,
        is_low_res_input=True,
        spynet_path='weights/spynet_20210409-c6c1bd09.pth'
    )
    state_dict = torch.load(
        f'weights/BasicVSR++_reds4_vimeo90k.pth'
    )['params']
    model.load_state_dict(state_dict, strict=True)
    model.eval().cuda()

    # Нормалізація та конвертація BGR→RGB
    tensor_frames = []
    for frame in frames:
        f_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        t = torch.from_numpy(f_rgb).float() / 255.0
        t = t.permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
        tensor_frames.append(t)

    # Batch усіх кадрів → (1, T, C, H, W)
    video_tensor = torch.stack(
        [f.squeeze(0) for f in tensor_frames], dim=0
    ).unsqueeze(0).cuda()

    with torch.no_grad(), torch.cuda.amp.autocast():
        output = model(video_tensor)  # (1, T, C, 4H, 4W)

    result = []
    for i in range(output.shape[1]):
        frame_t = output[0, i].float().cpu()
        frame_np = (frame_t.permute(1,2,0).numpy() * 255).clip(0,255)
        result.append(
            cv2.cvtColor(frame_np.astype(np.uint8), cv2.COLOR_RGB2BGR)
        )
    return result

Чанкована обробка довгих відео

Цілий фільм не вміститься в VRAM при BasicVSR++. Обробка чанками з перекриттям:

def upscale_long_video(
    input_path: str,
    output_path: str,
    chunk_frames: int = 50,     # кадрів у чанку
    overlap_frames: int = 5,    # перекриття для бесшовного сшивання
    scale: int = 4
) -> None:
    cap = cv2.VideoCapture(input_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    w   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    writer = cv2.VideoWriter(
        output_path,
        cv2.VideoWriter_fourcc(*'mp4v'),
        fps, (w * scale, h * scale)
    )

    frames_buffer = []
    processed_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            # Обробити залишилися кадри
            if frames_buffer:
                upscaled = upscale_video_basicvsr(frames_buffer)
                for upf in upscaled[overlap_frames:]:
                    writer.write(upf)
            break

        frames_buffer.append(frame)

        if len(frames_buffer) == chunk_frames:
            # Обробити чанк, зберегти перекриття для наступної ітерації
            upscaled = upscale_video_basicvsr(frames_buffer)
            for upf in upscaled[:-overlap_frames]:
                writer.write(upf)
            frames_buffer = frames_buffer[-overlap_frames:]

    cap.release()
    writer.release()

Спеціалізовані моделі для різних сценаріїв

Типові артефакти

• Часове миготіння — зазвичай спричинено невідповідною оцінкою шуму. Рішення: збільшити часове вікно (більше propagation blocks)
• Ghosting при швидких рухах — помилки оцінки flow. Рішення: використовувати більші feature dimensions
• Переповнення пам'яті на довгих послідовностях — чанкування з правильним перекриттям необхідне