Create feature_extractor.py

feature_extractor.py

@@ -0,0 +1,36 @@
+from transformers import PreTrainedFeatureExtractor
+import tensorflow as tf
+
+class CustomImageProcessor(PreTrainedFeatureExtractor):
+    def __init__(self, image_size=(224, 224), normalize=True):
+        super().__init__()
+        self.image_size = image_size
+        self.normalize = normalize
+
+    def __call__(self, images, return_tensors="pt"):
+        """
+        Preprocesses the images (resize, normalize, and format for the model).
+        Args:
+            images: List or batch of images in PIL or Tensor format.
+            return_tensors: Specify the format to return ('pt' for PyTorch, 'tf' for TensorFlow).
+        Returns:
+            Processed image tensor ready for model inference.
+        """
+        # Resize images to the required shape
+        images = [tf.image.resize(image, self.image_size) for image in images]
+
+        # Normalize the image if required (scale pixel values to [0, 1])
+        if self.normalize:
+            images = [image / 255.0 for image in images]
+
+        # Convert to the required tensor format (tf.Tensor for TensorFlow models)
+        images = tf.stack(images)
+
+        if return_tensors == "tf":
+            return {"pixel_values": images}
+        elif return_tensors == "pt":
+            # Convert to PyTorch tensor if needed (for PyTorch models)
+            import torch
+            return {"pixel_values": torch.tensor(images.numpy())}
+        else:
+            return {"pixel_values": images}