utils/imageHandling.py

import torch, torchvision

# %% image loading
def hfImageToTensor(image, width:int=1024, height:int=512)->torch.Tensor:
    """
    Convert an input image (PIL.Image or numpy array) from Hugging Face to a torch tensor
    of shape (3, height, width) and type float32.

    Args:
        image: Input image (PIL.Image or numpy array).
        width (int): Target width.
        height (int): Target height.

    Returns:
        torch.Tensor: Image tensor of shape (3, height, width).
    """
    image = image if isinstance(image, torch.Tensor) else torchvision.transforms.functional.to_tensor(image)
    return torchvision.transforms.functional.resize(image, [height, width])

# %% preprocessing
def preprocessing(image_tensor: torch.Tensor) -> torch.Tensor:
    """
    Standardize the image tensor and add batch dimension.

    Args:
        image_tensor (torch.Tensor): Image tensor of shape (3, H, W).

    Returns:
        torch.Tensor: Preprocessed tensor of shape (1, 3, H, W).
    """
    return torchvision.transforms.functional.normalize(
        image_tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    ).unsqueeze(0)