code/src/transforms.py

"""Clip-level augmentation and normalization.

Clips arrive as float32 arrays of shape (T, H, W) with values in [0, 1]
(grayscale). Transforms operate on the whole clip so that spatial ops are
consistent across time. Output is a tensor of shape (1, T, size, size).
"""
from __future__ import annotations

import numpy as np
import torch

# Mean/std computed over grayscale mouth ROIs; close enough for a baseline.
_MEAN = 0.421
_STD = 0.165


def _random_crop(clip: np.ndarray, size: int) -> np.ndarray:
    _, h, w = clip.shape
    top = np.random.randint(0, max(h - size, 0) + 1)
    left = np.random.randint(0, max(w - size, 0) + 1)
    return clip[:, top:top + size, left:left + size]


def _center_crop(clip: np.ndarray, size: int) -> np.ndarray:
    _, h, w = clip.shape
    top = max((h - size) // 2, 0)
    left = max((w - size) // 2, 0)
    return clip[:, top:top + size, left:left + size]


def _fix_length(clip: np.ndarray, num_frames: int) -> np.ndarray:
    """Pad (edge-repeat) or center-crop the time axis to num_frames."""
    t = clip.shape[0]
    if t == num_frames:
        return clip
    if t > num_frames:
        start = (t - num_frames) // 2
        return clip[start:start + num_frames]
    pad = num_frames - t
    front = pad // 2
    back = pad - front
    return np.concatenate(
        [np.repeat(clip[:1], front, axis=0), clip,
         np.repeat(clip[-1:], back, axis=0)], axis=0)


class ClipTransform:
    def __init__(self, image_size: int, num_frames: int, train: bool):
        self.image_size = image_size
        self.num_frames = num_frames
        self.train = train

    def __call__(self, clip: np.ndarray) -> torch.Tensor:
        clip = _fix_length(clip.astype(np.float32), self.num_frames)
        # Space: leave a margin so a random crop has something to move within.
        margin = self.image_size + 8
        clip = _center_crop(clip, min(margin, clip.shape[1]))
        if self.train:
            clip = _random_crop(clip, self.image_size)
            if np.random.rand() < 0.5:                      # horizontal flip
                clip = clip[:, :, ::-1]
        else:
            clip = _center_crop(clip, self.image_size)
        clip = (clip - _MEAN) / _STD
        clip = np.ascontiguousarray(clip)
        return torch.from_numpy(clip).unsqueeze(0)          # (1, T, H, W)