stimulus_synthesis/scoring/encoder_preprocess.py

from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from stimulus_synthesis.media.normalize import video_to_t_c_h_w


@dataclass(frozen=True)
class EncoderPreprocessSpec:
    size: int | tuple[int, int] | None = 224
    num_frames: int | None = None
    frame_sampling: str = "uniform"
    normalize_mean: tuple[float, float, float] | None = None
    normalize_std: tuple[float, float, float] | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class EncoderPreparedInput:
    videos: torch.Tensor
    frame_indices: list[int]
    spec: dict[str, Any]


def prepare_image_for_encoder(image: Any, spec: EncoderPreprocessSpec | None = None) -> EncoderPreparedInput:
    spec = spec or EncoderPreprocessSpec(num_frames=1)
    frame = _image_to_c_h_w(image)
    num_frames = int(spec.num_frames or 1)
    video = frame.unsqueeze(0).repeat(num_frames, 1, 1, 1)
    video = _resize_video(video, spec.size)
    video = _normalize(video, spec)
    return EncoderPreparedInput(videos=video.unsqueeze(0).contiguous(), frame_indices=[0] * num_frames, spec=spec.to_dict())


def prepare_video_for_encoder(frames: Any, spec: EncoderPreprocessSpec | None = None) -> EncoderPreparedInput:
    spec = spec or EncoderPreprocessSpec()
    video = video_to_t_c_h_w(frames)
    indices = sample_frame_indices(video.shape[0], spec.num_frames, spec.frame_sampling)
    if indices:
        video = video[torch.as_tensor(indices, dtype=torch.long)]
    video = _resize_video(video, spec.size)
    video = _normalize(video, spec)
    return EncoderPreparedInput(videos=video.unsqueeze(0).contiguous(), frame_indices=indices, spec=spec.to_dict())


def sample_frame_indices(total_frames: int, num_frames: int | None, policy: str = "uniform") -> list[int]:
    if total_frames <= 0:
        raise ValueError("total_frames must be positive.")
    if num_frames is None:
        return list(range(total_frames))
    if num_frames <= 0:
        raise ValueError("num_frames must be positive when set.")
    if policy != "uniform":
        raise ValueError(f"Unsupported frame sampling policy: {policy!r}")
    if total_frames == num_frames:
        return list(range(total_frames))
    if total_frames > num_frames:
        return torch.linspace(0, total_frames - 1, steps=num_frames).round().long().tolist()
    reps = int(np.ceil(num_frames / total_frames))
    return (list(range(total_frames)) * reps)[:num_frames]


def _image_to_c_h_w(image: Any) -> torch.Tensor:
    if isinstance(image, Image.Image):
        arr = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
        return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
    if isinstance(image, np.ndarray):
        arr = image.astype(np.float32, copy=False)
        if arr.max() > 1.0:
            arr = arr / 255.0
        tensor = torch.from_numpy(arr)
        if tensor.ndim != 3:
            raise ValueError(f"Expected image array with 3 dims, got {arr.shape}")
        if tensor.shape[-1] == 3:
            tensor = tensor.permute(2, 0, 1)
        return tensor.float().contiguous()
    if torch.is_tensor(image):
        tensor = image.detach().float()
        if tensor.ndim == 4:
            if tensor.shape[0] != 1:
                raise ValueError(f"Expected single-frame image tensor, got {tuple(tensor.shape)}")
            tensor = tensor.squeeze(0)
        if tensor.ndim != 3:
            raise ValueError(f"Expected image tensor with 3 dims, got {tuple(tensor.shape)}")
        if tensor.shape[-1] == 3:
            tensor = tensor.permute(2, 0, 1)
        if tensor.max() > 1.0:
            tensor = tensor / 255.0
        return tensor.contiguous()
    raise TypeError(f"Unsupported image type: {type(image)!r}")


def _resize_video(video: torch.Tensor, size: int | tuple[int, int] | None) -> torch.Tensor:
    if size is None:
        return video.float().clamp(0.0, 1.0).contiguous()
    size_hw = (int(size), int(size)) if isinstance(size, int) else (int(size[0]), int(size[1]))
    if tuple(video.shape[-2:]) == size_hw:
        return video.float().clamp(0.0, 1.0).contiguous()
    return F.interpolate(video.float(), size=size_hw, mode="bilinear", align_corners=False).clamp(0.0, 1.0).contiguous()


def _normalize(video: torch.Tensor, spec: EncoderPreprocessSpec) -> torch.Tensor:
    if spec.normalize_mean is None and spec.normalize_std is None:
        return video
    mean = torch.tensor(spec.normalize_mean or (0.0, 0.0, 0.0), dtype=video.dtype, device=video.device).view(1, 3, 1, 1)
    std = torch.tensor(spec.normalize_std or (1.0, 1.0, 1.0), dtype=video.dtype, device=video.device).view(1, 3, 1, 1)
    return (video - mean) / std