Реалізація AI-заміни фону на відео

Розробка ШІ для заміни фону у відео

Заміна фону у відео — завдання складніше, ніж у статичних зображеннях: потрібна часова узгодженість між кадрами (temporal coherence), інакше фон буде «миготіти». Застосування: віртуальний фон у відеоконференціях (Zoom, Teams), спортивні трансляції, новинні студії без хромакея, соціальні мережі. Реал-часні вимоги для конференцій — 15–30 FPS з затримкою <50 мс на CPU/GPU з низьким енергоспоживанням.

RVM — Robust Video Matting з часовою узгодженістю

import torch
import torchvision.transforms as T
from PIL import Image
import numpy as np
import cv2

class VideoBackgroundReplacer:
    def __init__(self, model_path: str, device: str = 'cuda'):
        self.device = device
        # RVM з recurrent state — ключ до часової узгодженості
        self.model = torch.jit.load(model_path).to(device)
        self.model.eval()
        self.transform = T.ToTensor()
        # Recurrent state збереженос між кадрами
        self.rec = [None] * 4

    def reset_state(self):
        """Скинути стан при зміні сцени/джерела"""
        self.rec = [None] * 4

    @torch.no_grad()
    def process_frame(self, frame_bgr: np.ndarray,
                       background_bgr: np.ndarray) -> np.ndarray:
        """
        Обробка одного кадру відео.
        Стан (rec) збережено між викликами для плавності.
        """
        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        bg_rgb = cv2.cvtColor(background_bgr, cv2.COLOR_BGR2RGB)

        # Змінюємо розмір до кратного 64 для моделі
        h, w = frame_rgb.shape[:2]
        src = self.transform(Image.fromarray(frame_rgb)).unsqueeze(0).to(self.device)
        bgr_tensor = self.transform(
            Image.fromarray(bg_rgb).resize((w, h))
        ).unsqueeze(0).to(self.device)

        # Основний висновок з передачею recurrent state
        fgr, pha, *self.rec = self.model(src, *self.rec, downsample_ratio=0.25)

        # Compositing
        composite = fgr * pha + bgr_tensor * (1 - pha)
        result = (composite.squeeze().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
        return cv2.cvtColor(result, cv2.COLOR_RGB2BGR)

    def replace_in_video(self, input_path: str,
                          background_path: str,
                          output_path: str) -> dict:
        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))

        # Завантаджуємо фон (може бути зображення або відео)
        if background_path.endswith(('.jpg', '.png')):
            bg = cv2.imread(background_path)
            bg = cv2.resize(bg, (w, h))
            bg_is_video = False
        else:
            bg_cap = cv2.VideoCapture(background_path)
            bg_is_video = True

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

        self.reset_state()
        frame_count = 0

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            if bg_is_video:
                ret_bg, bg = bg_cap.read()
                if not ret_bg:
                    bg_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                    _, bg = bg_cap.read()
                bg = cv2.resize(bg, (w, h))

            result = self.process_frame(frame, bg)
            out.write(result)
            frame_count += 1

        cap.release()
        out.release()

        return {'frames': frame_count, 'fps': fps, 'output': output_path}

Реал-час для відеоконференцій (WebRTC/ONNX)

import onnxruntime as ort

class RealtimeBackgroundProcessor:
    """
    ONNX Runtime для CPU-оптимізації на машинах без GPU.
    Мета: 30 FPS на ноутбуці, затримка <33 мс/кадр.
    """
    def __init__(self, onnx_model_path: str):
        # Налаштування для максимальної продуктивності на CPU
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = 4
        opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        self.session = ort.InferenceSession(
            onnx_model_path,
            sess_options=opts,
            providers=['TensorrtExecutionProvider',
                       'CUDAExecutionProvider',
                       'CPUExecutionProvider']
        )

        # Recurrent state як numpy масиви
        self.rec_states = [
            np.zeros((1, 1, 1, 1), dtype=np.float32) for _ in range(4)
        ]

    def process_frame_fast(self, frame_rgb: np.ndarray,
                             target_size: tuple = (256, 144)) -> np.ndarray:
        """
        Зменшення до 256x144 для CPU реал-часу.
        Апскейл маски назад через білінійну інтерполяцію.
        """
        orig_h, orig_w = frame_rgb.shape[:2]
        small = cv2.resize(frame_rgb, target_size)
        small_f = small.astype(np.float32) / 255.0
        src = small_f.transpose(2, 0, 1)[np.newaxis]  # [1, 3, H, W]

        outputs = self.session.run(
            None,
            {'src': src, 'r1i': self.rec_states[0], 'r2i': self.rec_states[1],
             'r3i': self.rec_states[2], 'r4i': self.rec_states[3],
             'downsample_ratio': np.array([0.25])}
        )
        fgr, pha = outputs[0], outputs[1]
        self.rec_states = list(outputs[2:6])

        # Апскейл маски назад до вихідного розміру
        alpha_small = pha[0, 0]
        alpha_full = cv2.resize(alpha_small, (orig_w, orig_h),
                                interpolation=cv2.INTER_LINEAR)
        return alpha_full

Віртуальний розмитий фон (Bokeh effect)

def apply_background_blur(frame: np.ndarray,
                            alpha: np.ndarray,
                            blur_radius: int = 25) -> np.ndarray:
    """
    Альтернатива заміні — розмиття фону (як у Google Meet/Teams).
    Не потребує завантаження зображення, працює швидше.
    """
    # Розмиваємо весь кадр
    blurred = cv2.GaussianBlur(frame, (blur_radius * 2 + 1, blur_radius * 2 + 1), 0)

    # Compositing з м'якими краями
    alpha_3ch = np.stack([alpha, alpha, alpha], axis=2)
    result = (frame * alpha_3ch + blurred * (1 - alpha_3ch)).astype(np.uint8)
    return result