utils/transforms.py

"""Differentiable image transforms used by the EOT layer are defined in this module. Forward passes apply realistic degradations while backward passes use straight-through estimators where exact gradients are unavailable."""

from __future__ import annotations

import io
import random
from dataclasses import dataclass
from typing import List

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

from config import EOTConfig


@dataclass(slots=True)
class TransformSample:
    """A sampled transform stores the transformed tensor and a readable label. This makes validation logs easier to interpret while keeping the EOT API compact."""

    image: torch.Tensor
    name: str


def _tensor_to_pil(image: torch.Tensor) -> Image.Image:
    """Convert a single CHW tensor in the 0..1 range into a PIL image. This helper is used for non-differentiable transforms such as JPEG round-trips."""

    image = image.detach().clamp(0.0, 1.0).cpu()
    array = (image.permute(1, 2, 0).numpy() * 255.0).round().astype(np.uint8)
    return Image.fromarray(array)


def _pil_to_tensor(image: Image.Image, device: torch.device | str) -> torch.Tensor:
    """Convert a PIL image into a float CHW tensor on the requested device. Output pixels are normalized to the 0..1 interval."""

    array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
    tensor = torch.from_numpy(array).permute(2, 0, 1)
    return tensor.to(device=device)


def _ste_replace(original: torch.Tensor, transformed: torch.Tensor) -> torch.Tensor:
    """Apply a straight-through estimator by keeping the transformed forward value and the original backward path. This is used for JPEG and screenshot simulations whose exact gradients are not practical here."""

    return original + (transformed - original).detach()


def jpeg_roundtrip(image: torch.Tensor, quality: int) -> torch.Tensor:
    """Apply an in-memory JPEG encode-decode round-trip to a BCHW tensor. The backward path is approximated with the identity function so optimization remains stable."""

    outputs = []
    for sample in image:
        pil_image = _tensor_to_pil(sample)
        buffer = io.BytesIO()
        pil_image.save(buffer, format="JPEG", quality=int(quality), subsampling=0)
        buffer.seek(0)
        decoded = Image.open(buffer).convert("RGB")
        outputs.append(_pil_to_tensor(decoded, sample.device))
    transformed = torch.stack(outputs, dim=0)
    return _ste_replace(image, transformed)


def resize_roundtrip(image: torch.Tensor, scale_factor: float) -> torch.Tensor:
    """Resize a BCHW tensor down or up and then restore it to the original shape. The second interpolation step makes the distortion compatible with downstream models."""

    height, width = image.shape[-2:]
    resized = F.interpolate(
        image,
        scale_factor=scale_factor,
        mode="bilinear",
        align_corners=False,
        recompute_scale_factor=True,
    )
    return F.interpolate(resized, size=(height, width), mode="bilinear", align_corners=False)


def gaussian_blur(image: torch.Tensor, kernel_size: int, sigma: float) -> torch.Tensor:
    """Apply a Gaussian blur to a BCHW tensor. The kernel size is expected to be odd and is provided from config to avoid hard-coded values."""

    return TF.gaussian_blur(image, kernel_size=[kernel_size, kernel_size], sigma=[sigma, sigma])


def add_gaussian_noise(image: torch.Tensor, std: float) -> torch.Tensor:
    """Add bounded Gaussian noise to a BCHW tensor. The result is clamped to the valid image range immediately after perturbation."""

    noisy = image + torch.randn_like(image) * std
    return noisy.clamp(0.0, 1.0)


def brightness_shift(image: torch.Tensor, delta: float) -> torch.Tensor:
    """Apply a multiplicative brightness shift to a BCHW tensor. Positive and negative deltas are sampled by the caller and passed through directly."""

    factor = 1.0 + delta
    return (image * factor).clamp(0.0, 1.0)


def screenshot_simulation(
    image: torch.Tensor,
    scale_factor: float,
    blur_sigma: float,
    noise_std: float,
    jpeg_quality: int,
    blur_kernel_size: int,
) -> torch.Tensor:
    """Approximate a screenshot or screen recapture process with resize, blur, noise, and JPEG artifacts. This is a practical open-source substitute for a literal render-recature pipeline in Phase 1."""

    transformed = resize_roundtrip(image, scale_factor=scale_factor)
    transformed = gaussian_blur(transformed, kernel_size=blur_kernel_size, sigma=blur_sigma)
    transformed = add_gaussian_noise(transformed, std=noise_std)
    transformed = jpeg_roundtrip(transformed, quality=jpeg_quality)
    return transformed


def sample_eot_transforms(image: torch.Tensor, config: EOTConfig) -> List[TransformSample]:
    """Sample the configured transformation pool for one EOT step. Each returned sample preserves the input tensor shape so losses can be averaged directly."""

    samples: List[TransformSample] = []
    for _ in range(config.sample_size):
        choice = random.choice(("jpeg", "resize", "blur", "noise", "brightness", "screenshot"))
        if choice == "jpeg":
            quality = random.randint(config.jpeg_quality_min, config.jpeg_quality_max)
            samples.append(TransformSample(jpeg_roundtrip(image, quality=quality), f"jpeg_q{quality}"))
        elif choice == "resize":
            scale = random.uniform(config.resize_min_scale, config.resize_max_scale)
            samples.append(TransformSample(resize_roundtrip(image, scale_factor=scale), f"resize_{scale:.2f}"))
        elif choice == "blur":
            sigma = random.uniform(config.blur_sigma_min, config.blur_sigma_max)
            samples.append(
                TransformSample(
                    gaussian_blur(image, kernel_size=config.blur_kernel_size, sigma=sigma),
                    f"blur_{sigma:.2f}",
                )
            )
        elif choice == "noise":
            std = random.uniform(config.noise_std_min, config.noise_std_max)
            samples.append(TransformSample(add_gaussian_noise(image, std=std), f"noise_{std:.3f}"))
        elif choice == "brightness":
            delta = random.uniform(-config.brightness_delta, config.brightness_delta)
            samples.append(TransformSample(brightness_shift(image, delta=delta), f"brightness_{delta:.3f}"))
        else:
            quality = random.randint(config.jpeg_quality_min, config.jpeg_quality_max)
            samples.append(
                TransformSample(
                    screenshot_simulation(
                        image,
                        scale_factor=config.screenshot_scale,
                        blur_sigma=config.screenshot_blur_sigma,
                        noise_std=config.screenshot_noise_std,
                        jpeg_quality=quality,
                        blur_kernel_size=config.blur_kernel_size,
                    ),
                    f"screenshot_q{quality}",
                )
            )
    return samples


def pil_image_to_bchw(image: Image.Image, device: torch.device | str) -> torch.Tensor:
    """Convert a PIL image into a 1xCxHxW tensor. This is used by the CLI and validation code when moving between file IO and the optimization pipeline."""

    return _pil_to_tensor(image, device=device).unsqueeze(0)


def bchw_to_pil_image(image: torch.Tensor) -> Image.Image:
    """Convert a 1xCxHxW or CxHxW tensor into a PIL image. The helper clamps values so saved images remain valid even after optimization."""

    if image.ndim == 4:
        image = image[0]
    return _tensor_to_pil(image)