stimulus_synthesis/scoring/robust_transform.py

from __future__ import annotations

import hashlib
from dataclasses import asdict, dataclass
from typing import Any

import torch
import torch.nn.functional as F


@dataclass(frozen=True)
class RobustTransformSpec:
    crop_scale: float = 0.80
    gaussian_sigma: float = 0.10
    num_draws: int = 4
    aggregate: str = "mean"
    seed: int = 0

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> "RobustTransformSpec | None":
        if data is None:
            return None
        enabled = bool(data.pop("enabled", True)) if "enabled" in data else True
        return cls(**data) if enabled else None

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


class RobustTransformScorer:
    """Apply deterministic robust scoring draws before delegating to a scorer."""

    def __init__(self, scorer: Any, spec: RobustTransformSpec | None = None) -> None:
        self.scorer = scorer
        self.spec = spec or RobustTransformSpec()

    def score(self, videos: torch.Tensor, target: Any, **kwargs) -> list[float]:
        transformed = apply_robust_transform(videos, self.spec)
        raw_scores = self.scorer.score(transformed, target, **kwargs)
        score_tensor = torch.as_tensor(raw_scores, dtype=torch.float32).reshape(videos.shape[0], self.spec.num_draws)
        if self.spec.aggregate != "mean":
            raise ValueError(f"Unsupported robust score aggregate: {self.spec.aggregate!r}")
        return score_tensor.mean(dim=1).tolist()


def apply_robust_transform(videos: torch.Tensor, spec: RobustTransformSpec | None = None) -> torch.Tensor:
    spec = spec or RobustTransformSpec()
    if videos.ndim != 5:
        raise ValueError(f"Expected videos shaped (B,T,C,H,W), got {tuple(videos.shape)}")
    if spec.num_draws <= 0:
        raise ValueError("num_draws must be positive.")
    if not (0.0 < spec.crop_scale <= 1.0):
        raise ValueError("crop_scale must be in (0, 1].")
    if spec.gaussian_sigma < 0.0:
        raise ValueError("gaussian_sigma must be non-negative.")

    videos = videos.float().clamp(0.0, 1.0)
    out = []
    for batch_idx in range(videos.shape[0]):
        base_seed = _content_seed(videos[batch_idx], spec.seed)
        for draw_idx in range(spec.num_draws):
            out.append(_transform_one(videos[batch_idx], spec, base_seed + draw_idx * 7919))
    return torch.stack(out, dim=0).contiguous()


def _transform_one(video: torch.Tensor, spec: RobustTransformSpec, seed: int) -> torch.Tensor:
    generator = torch.Generator(device=video.device).manual_seed(int(seed) % (2**63 - 1))
    transformed = _random_resized_crop(video, spec.crop_scale, generator)
    if spec.gaussian_sigma:
        noise = torch.randn(
            transformed.shape,
            generator=generator,
            device=transformed.device,
            dtype=transformed.dtype,
        )
        transformed = transformed + float(spec.gaussian_sigma) * noise
    return transformed.clamp(0.0, 1.0)


def _random_resized_crop(video: torch.Tensor, crop_scale: float, generator: torch.Generator) -> torch.Tensor:
    if crop_scale == 1.0:
        return video
    _t, _c, h, w = video.shape
    crop_h = max(1, int(round(h * crop_scale)))
    crop_w = max(1, int(round(w * crop_scale)))
    max_y = h - crop_h
    max_x = w - crop_w
    y0 = int(torch.randint(max_y + 1, (1,), generator=generator, device=video.device).item()) if max_y else 0
    x0 = int(torch.randint(max_x + 1, (1,), generator=generator, device=video.device).item()) if max_x else 0
    crop = video[:, :, y0 : y0 + crop_h, x0 : x0 + crop_w]
    return F.interpolate(crop, size=(h, w), mode="bilinear", align_corners=False)


def _content_seed(video: torch.Tensor, seed: int) -> int:
    quantized = (video.detach().cpu().clamp(0.0, 1.0) * 255).round().to(torch.uint8).numpy().tobytes()
    digest = hashlib.blake2b(quantized, digest_size=8, person=b"nevo-rbt").digest()
    return (int.from_bytes(digest, "little") + int(seed)) % (2**63 - 1)