src/preprocess.py

from transformers import ViTImageProcessor
from .utils import MODEL_DIR

processor = ViTImageProcessor.from_pretrained(MODEL_DIR)


def preprocess_image(image):
    """Preprocesses a single image for ViT inference."""
    inputs = processor(images=image, return_tensors="pt")
    return inputs