Upload processor

image_processing_efficientnet.py

@@ -1,3 +1,8 @@
+from PIL import Image
+from torch import Tensor, stack
+from numpy import ndarray
+from typing import Union, List
+
 from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
 from timm import create_model
 from timm.data import resolve_data_config
@@ -6,19 +11,54 @@ from timm.data.transforms_factory import create_transform
 class EfficientNetImageProcessor(BaseImageProcessor):
     model_input_names = ["pixel_values"]
 
-    def __init__(self, 
-                 model_name: str,
-                 **kwargs
-                 ):
-        super().__init__(**kwargs)
-
+    def __init__(
+        self, 
+        model_name: str,
+        **kwargs,
+    ):
         self.model_name = model_name
         self.config = resolve_data_config({}, model=create_model(model_name, pretrained=False))
+        super().__init__(**kwargs)
+
+    def preprocess(
+        self, 
+        images: Union[List[Union[Image.Image, Tensor]], Image.Image, Tensor],
+    ) -> BatchFeature:
+        """
+        Preprocesses input images by applying transformations and returning them as a BatchFeature.
+
+        Parameters
+        ----------
+        images : Union[List[PIL.Image.Image, torch.Tensor], PIL.Image.Image, torch.Tensor]
+            A single image or a list of images in one of the accepted formats.
+
+        Returns
+        -------
+        BatchFeature
+            A batch of transformed images 
+        """
+        images = [images] if not isinstance(images, list) else images
+
+        # TEST: empty list
+        if len(images) == 0:
+            raise ValueError("Received an empty list of images")
         
+        # TEST: validate input type
+        test_image = images[0]
+        if not isinstance(images[0], (Image.Image, Tensor)):
+            raise TypeError(
+                f"Expected image to be of type PIL.Image.Image, torch.Tensor, or numpy.ndarray, "
+                f"but got {type(test_image).__name__} instead."
+            )
 
-    def preprocess(self, image):
+        # Apply transformations
         transforms = create_transform(**self.config)
-        data = {'pixel_values': transforms(image).unsqueeze(0)}
+        transformed_images = [transforms(image) for image in images]
+
+        # Convert to batch tensor
+        transformed_image_tensors = stack(transformed_images)
+
+        data = {'pixel_values': transformed_image_tensors}
         return BatchFeature(data=data)
 
 __all__ = [