Add: convert_to_rgb

huggingface · Nov 12, 2024 · ae573cf · ae573cf
1 parent e9503ec
commit ae573cf
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/transformers/models/vit/image_processing_vit_fast.py b/src/transformers/models/vit/image_processing_vit_fast.py
@@ -20,7 +20,7 @@
 from ...image_processing_base import BatchFeature
 from ...image_processing_utils import get_size_dict
 from ...image_processing_utils_fast import BaseImageProcessorFast, SizeDict
-from ...image_transforms import FusedRescaleNormalize, NumpyToTensor, Rescale
+from ...image_transforms import FusedRescaleNormalize, NumpyToTensor, Rescale, convert_to_rgb
 from ...image_utils import (
     IMAGENET_STANDARD_MEAN,
     IMAGENET_STANDARD_STD,
@@ -279,6 +279,9 @@ def preprocess(
             image_type=image_type,
         )
 
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
         transforms = self.get_transforms(
             do_resize=do_resize,
             do_rescale=do_rescale,