utils/load_image.py

import torch
from PIL import Image
from torchvision import transforms

input_size = (224, 224)

img_transform = transforms.Compose([
    transforms.Resize(input_size),
    transforms.CenterCrop(input_size),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def load_image(image_path:str, input_size: tuple = input_size, img_transform = img_transform):

  """
  Load and preprocess an image for the model.

  Args:
      image_path (str): Path to the input image file.
      input_size (tuple): Target size for the image (default: (224, 224)).
      img_transform(transforms): Transforms that is to be applied to the image (default is defined above the function)

  Returns:
      torch.Tensor: Preprocessed image tensor with shape (1, 3, H, W).
  """

  image = Image.open(image_path).convert("RGB")
  input_tensor = img_transform(image).unsqueeze(0)
  input_tensor.requires_grad = True

  return input_tensor