Реалізація AI-стабілізації відео

Розробка ШІ для стабілізації відео

Тремтіння камери — невідворотний артефакт зйомки з рук, дронів, спортивних камер. Класична стабілізація працює через optical flow: оцінити рух між кадрами, згладити траєкторію, компенсувати тремтіння. ШІ-методи додають семантичне розуміння: розрізняють рух оператора від руху об'єкту зйомки, краще обробляють динамічні сцени, можуть відновлювати «вирізані» пікселі через inpainting.

Класична стабілізація через optical flow

import cv2
import numpy as np
from scipy.signal import medfilt

class VideoStabilizer:
    def __init__(self, smoothing_window: int = 30,
                  crop_ratio: float = 0.1):
        self.smoothing_window = smoothing_window
        self.crop_ratio = crop_ratio  # обрізка границь після стабілізації

    def stabilize(self, input_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))

        # Крок 1: Обчислити траєкторію камери
        transforms = self._estimate_transforms(cap)
        cap.release()

        # Крок 2: Згладити траєкторію
        smoothed = self._smooth_trajectory(transforms)

        # Крок 3: Застосувати стабілізуючі трансформації
        cap = cv2.VideoCapture(input_path)
        out = cv2.VideoWriter(output_path,
                              cv2.VideoWriter_fourcc(*'mp4v'),
                              fps, (w, h))

        for i, (orig, smooth) in enumerate(zip(transforms, smoothed)):
            ret, frame = cap.read()
            if not ret:
                break
            stabilized = self._apply_transform(frame, orig, smooth, w, h)
            out.write(stabilized)

        cap.release()
        out.release()
        return {'frames': len(transforms), 'smoothing_window': self.smoothing_window}

    def _estimate_transforms(self, cap) -> list[np.ndarray]:
        """Оцінка аффінних трансформацій між сусідніми кадрами"""
        ret, prev = cap.read()
        prev_gray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
        transforms = []

        while True:
            ret, curr = cap.read()
            if not ret:
                break
            curr_gray = cv2.cvtColor(curr, cv2.COLOR_BGR2GRAY)

            # Детекція та трекінг точок
            prev_pts = cv2.goodFeaturesToTrack(
                prev_gray, maxCorners=200, qualityLevel=0.01,
                minDistance=30, blockSize=3
            )
            curr_pts, status, _ = cv2.calcOpticalFlowPyrLK(
                prev_gray, curr_gray, prev_pts, None
            )

            # Фільтрація надійних точок
            valid_prev = prev_pts[status == 1]
            valid_curr = curr_pts[status == 1]

            # Оцінка аффінної трансформації
            m, _ = cv2.estimateAffinePartial2D(valid_prev, valid_curr)
            if m is None:
                m = np.eye(2, 3, dtype=np.float64)

            transforms.append(m)
            prev_gray = curr_gray

        return transforms

    def _smooth_trajectory(self, transforms: list) -> list:
        """Ковзне середнє для згладжування траєкторії"""
        trajectory = np.cumsum([m[:, 2] for m in transforms], axis=0)
        smoothed = np.zeros_like(trajectory)

        for i in range(len(trajectory)):
            start = max(0, i - self.smoothing_window // 2)
            end = min(len(trajectory), i + self.smoothing_window // 2)
            smoothed[i] = trajectory[start:end].mean(axis=0)

        # Дельта трансформацій для застосування
        delta = smoothed - trajectory
        result = []
        for i, m in enumerate(transforms):
            m_smooth = m.copy()
            m_smooth[:, 2] += delta[i]
            result.append(m_smooth)
        return result

    def _apply_transform(self, frame: np.ndarray,
                          orig_m: np.ndarray,
                          smooth_m: np.ndarray,
                          w: int, h: int) -> np.ndarray:
        stabilized = cv2.warpAffine(frame, smooth_m, (w, h))
        # Кроп для приховування чорних границь
        crop = int(min(w, h) * self.crop_ratio)
        stabilized = stabilized[crop:h-crop, crop:w-crop]
        return cv2.resize(stabilized, (w, h))

DUT — Deep Unified Transformer для ШІ-стабілізації

class DeepVideoStabilizer:
    """
    ШІ-підхід: навчання стабілізувати відео на парах нестабільне/стабільне.
    Перевага перед класикою: краще обробляє
    rolling shutter, швидкий рух, розмиття.
    """
    def __init__(self, checkpoint_path: str, device: str = 'cuda'):
        import sys
        sys.path.append('/opt/DUT')
        from model import DUTStabilizer
        self.model = DUTStabilizer()
        self.model.load_state_dict(torch.load(checkpoint_path))
        self.model.eval().to(device)
        self.device = device

    @torch.no_grad()
    def stabilize_clip(self, frames: list[np.ndarray],
                        window_size: int = 16) -> list[np.ndarray]:
        """
        Обробляє відео вікнами по window_size кадрів.
        Ключова особливість DUT: використовує майбутні кадри
        для передбачення поточної стабілізації.
        """
        results = []
        for i in range(0, len(frames), window_size // 2):
            window = frames[i:i+window_size]
            if len(window) < window_size:
                # Дополнити останнім кадром
                window = window + [window[-1]] * (window_size - len(window))

            tensor = self._frames_to_tensor(window)
            stabilized_tensor = self.model(tensor.to(self.device))
            stabilized_frames = self._tensor_to_frames(stabilized_tensor)
            results.extend(stabilized_frames[:window_size//2])

        return results[:len(frames)]

Метрики якості стабілізації відео

def evaluate_stabilization(unstable_frames: list, stable_frames: list) -> dict:
    """
    Метрики:
    - Cropping Ratio: скільки пікселів збережено (вище = краще)
    - Distortion Value: спотворення геометрії (нижче = краще)
    - Stability Score: дисперсія руху між кадрами (нижче = краще)
    """
    # Stability: дисперсія optical flow у стабілізованому відео
    flows = []
    for i in range(1, len(stable_frames)):
        prev_gray = cv2.cvtColor(stable_frames[i-1], cv2.COLOR_BGR2GRAY)
        curr_gray = cv2.cvtColor(stable_frames[i], cv2.COLOR_BGR2GRAY)
        flow = cv2.calcOpticalFlowFarneback(prev_gray, curr_gray, None,
                                             0.5, 3, 15, 3, 5, 1.2, 0)
        flows.append(np.abs(flow).mean())

    return {
        'stability_score': float(np.std(flows)),
        'mean_motion': float(np.mean(flows)),
        'max_motion': float(np.max(flows))
    }