processing_egosieve.py

"""Image-sequence processor for EgoSieve checkpoints."""

from __future__ import annotations

from collections.abc import Sequence
from os import PathLike
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageOps
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature

DEFAULT_IMAGE_MEAN = (0.485, 0.456, 0.406)
DEFAULT_IMAGE_STD = (0.229, 0.224, 0.225)


def _is_frame(value: Any) -> bool:
    if isinstance(value, (Image.Image, str, PathLike)):
        return True
    if isinstance(value, np.ndarray):
        return value.ndim in (2, 3)
    try:
        import torch

        return isinstance(value, torch.Tensor) and value.ndim in (2, 3)
    except ImportError:
        return False


class EgoSieveProcessor(BaseImageProcessor):
    """Prepare one or more frame sequences for :class:`EgoSieveModel`.

    Video decoding intentionally lives outside this class. Passing decoded
    frames keeps timestamp selection explicit and makes the processor useful
    with PyAV, decord, camera streams, or the bundled ffmpeg sampler.
    """

    model_input_names = ["pixel_values", "frame_mask"]

    def __init__(
        self,
        size: int | dict[str, int] = 224,
        resize_shortest_edge: int | None = 256,
        num_frames: int = 12,
        image_mean: Sequence[float] = DEFAULT_IMAGE_MEAN,
        image_std: Sequence[float] = DEFAULT_IMAGE_STD,
        resample: int = Image.Resampling.BICUBIC,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        if isinstance(size, dict):
            height = int(size.get("height", size.get("shortest_edge", 224)))
            width = int(size.get("width", height))
        else:
            height = width = int(size)
        if height <= 0 or width <= 0:
            raise ValueError("size must be positive")
        if resize_shortest_edge is not None and (
            isinstance(resize_shortest_edge, bool)
            or not isinstance(resize_shortest_edge, int)
            or resize_shortest_edge <= 0
        ):
            raise ValueError("resize_shortest_edge must be a positive integer or None")
        if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0:
            raise ValueError("num_frames must be a positive integer")
        self.size = {"height": height, "width": width}
        self.resize_shortest_edge = resize_shortest_edge
        self.num_frames = int(num_frames)
        self.image_mean = [float(v) for v in image_mean]
        self.image_std = [float(v) for v in image_std]
        self.resample = int(resample)

    @staticmethod
    def _as_pil(frame: Any) -> Image.Image:
        if isinstance(frame, Image.Image):
            return frame.convert("RGB")
        if isinstance(frame, (str, PathLike)):
            with Image.open(Path(frame)) as image:
                return image.convert("RGB")
        try:
            import torch

            if isinstance(frame, torch.Tensor):
                frame = frame.detach().cpu().numpy()
        except ImportError:
            pass
        array = np.asarray(frame)
        if array.ndim == 2:
            array = np.repeat(array[..., None], 3, axis=-1)
        if array.ndim != 3:
            raise ValueError(f"each frame must have 2 or 3 dimensions, got {array.shape}")
        if array.shape[0] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4):
            array = np.moveaxis(array, 0, -1)
        if np.issubdtype(array.dtype, np.floating):
            if array.size and float(np.nanmax(array)) <= 1.0:
                array = array * 255.0
            array = np.nan_to_num(array, nan=0.0, posinf=255.0, neginf=0.0)
        array = np.clip(array, 0, 255).astype(np.uint8)
        if array.shape[-1] == 1:
            array = np.repeat(array, 3, axis=-1)
        if array.shape[-1] == 4:
            return Image.fromarray(array).convert("RGB")
        return Image.fromarray(array)

    def _prepare_frame(self, frame: Any) -> np.ndarray:
        image = self._as_pil(frame)
        if self.resize_shortest_edge is not None:
            shortest = min(image.size)
            scale = self.resize_shortest_edge / shortest
            resized = image.resize(
                (
                    max(1, round(image.width * scale)),
                    max(1, round(image.height * scale)),
                ),
                resample=self.resample,
            )
            crop_width = self.size["width"]
            crop_height = self.size["height"]
            left = (resized.width - crop_width) // 2
            top = (resized.height - crop_height) // 2
            fitted = resized.crop((left, top, left + crop_width, top + crop_height))
        else:
            resized = image
            fitted = ImageOps.fit(
                resized,
                (self.size["width"], self.size["height"]),
                method=self.resample,
                centering=(0.5, 0.5),
            )
        array = np.asarray(fitted, dtype=np.float32) / 255.0
        array = (array - np.asarray(self.image_mean, dtype=np.float32)) / np.asarray(
            self.image_std, dtype=np.float32
        )
        return np.moveaxis(array, -1, 0)

    @staticmethod
    def _uniform_indices(length: int, count: int) -> np.ndarray:
        if length <= count:
            return np.arange(length, dtype=np.int64)
        return np.rint(np.linspace(0, length - 1, count)).astype(np.int64)

    def preprocess(
        self,
        videos: Sequence[Any] | Any,
        return_tensors: str | None = None,
        **_: Any,
    ) -> BatchFeature:
        """Normalize decoded frames.

        `videos` may be a single sequence of frames or a batch of frame
        sequences. Long inputs are sampled uniformly. Short inputs are padded
        by repeating their final frame and marked false in `frame_mask`.
        """

        if videos is None:
            raise ValueError("videos is required")
        if _is_frame(videos):
            batch = [[videos]]
        else:
            values = list(videos)
            if not values:
                raise ValueError("videos cannot be empty")
            batch = [values] if _is_frame(values[0]) else [list(v) for v in values]

        pixel_batch: list[np.ndarray] = []
        mask_batch: list[np.ndarray] = []
        for frames in batch:
            if not frames:
                raise ValueError("a video cannot contain zero frames")
            indices = self._uniform_indices(len(frames), self.num_frames)
            selected = [frames[int(i)] for i in indices]
            valid = len(selected)
            if valid < self.num_frames:
                selected.extend([selected[-1]] * (self.num_frames - valid))
            pixel_batch.append(np.stack([self._prepare_frame(frame) for frame in selected]))
            mask = np.zeros(self.num_frames, dtype=np.int64)
            mask[:valid] = 1
            mask_batch.append(mask)

        return BatchFeature(
            {
                "pixel_values": np.stack(pixel_batch).astype(np.float32),
                "frame_mask": np.stack(mask_batch),
            },
            tensor_type=return_tensors,
        )

    __call__ = preprocess