import torch
from PIL import Image
import numpy as np
import torchvision.transforms as T

# Standard ImageNet normalization for ViT models if needed
PREPROCESS_TRANSFORMS = T.Compose([
    T.Resize((224, 224)),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def load_and_preprocess_image(image_path_or_pil) -> torch.Tensor:
    """
    Loads and resizes an image to 224x224, returning a normalized PyTorch tensor.
    Supports file paths or PIL Image objects.
    """
    if isinstance(image_path_or_pil, str):
        image = Image.open(image_path_or_pil).convert("RGB")
    else:
        image = image_path_or_pil.convert("RGB")
        
    tensor = PREPROCESS_TRANSFORMS(image)
    return tensor.unsqueeze(0) # Add batch dimension -> (1, 3, 224, 224)

def load_pil_image(image_path_or_pil) -> Image.Image:
    """
    Ensures a PIL Image is returned in RGB mode.
    """
    if isinstance(image_path_or_pil, str):
        return Image.open(image_path_or_pil).convert("RGB")
    return image_path_or_pil.convert("RGB")