import torch
from PIL import Image
from torchvision import transforms
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature

MODEL_CONTENT_SIZE = 1444
MODEL_PADDED_SIZE = 1472
PAD = (MODEL_PADDED_SIZE - MODEL_CONTENT_SIZE) // 2


class Swinv2SegmenterImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def preprocess(self, images, return_tensors=None, **kwargs):
        """Resize to 1444x1444, ImageNet-normalize, zero-pad to 1472x1472."""
        if isinstance(images, Image.Image):
            images = [images]

        pixel_values = []
        for image in images:
            resized = image.resize((MODEL_CONTENT_SIZE, MODEL_CONTENT_SIZE))
            tensor = transforms.ToTensor()(resized)[:3]
            tensor = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(tensor)
            tensor = torch.nn.functional.pad(tensor, (PAD, PAD, PAD, PAD), value=0)
            pixel_values.append(tensor)

        return BatchFeature(data={"pixel_values": torch.stack(pixel_values)}, tensor_type=return_tensors)