training/iqa_backends/preprocessing.py

"""Shared image preprocessing helpers for IQA backends."""

from __future__ import annotations

from typing import Callable

import torch
from PIL import Image
from torchvision import transforms

from training.utils.utils_data import resize_crop


def random_resize_crop(crop_size: int, downscale_factor: int) -> Callable[[torch.Tensor], torch.Tensor]:
    """Random crop with optional downscale; input/output are ``[0, 1]`` RGB tensors."""

    def _transform(image: torch.Tensor) -> torch.Tensor:
        return resize_crop(image, crop_size=crop_size, downscale_factor=downscale_factor)

    return _transform


def center_crop_to_tensor(
    crop_size: int,
    *,
    downscale_factor: int = 1,
) -> Callable[[Image.Image], torch.Tensor]:
    """Deterministic center crop for inference and visualization."""

    def _downscale(img: Image.Image) -> Image.Image:
        width, height = img.size
        return img.resize((width // downscale_factor, height // downscale_factor))

    steps: list = []
    if downscale_factor > 1:
        steps.append(_downscale)
    steps.extend([transforms.CenterCrop(crop_size), transforms.ToTensor()])
    pipeline = transforms.Compose(steps)
    return pipeline