small simplification to the resize

huggingface · Aug 13, 2024 · 1c10a48 · 1c10a48
1 parent 8f01856
commit 1c10a48
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/src/transformers/models/idefics3/image_processing_idefics3.py b/src/transformers/models/idefics3/image_processing_idefics3.py
@@ -434,7 +434,6 @@ def resize(
             input_data_format (`ChannelDimension` or `str`, *optional*):
                 The channel dimension format of the input image. If not provided, it will be inferred.
         """
-        image_modes = {1: 'P', 3: 'RGB', 4: 'RGBA'}
         if "longest_edge" in size:
             size = get_resize_output_image_size(
                 image, resolution_max_side=size["longest_edge"], input_data_format=input_data_format
@@ -446,10 +445,10 @@ def resize(
         if isinstance(image, Image.Image):
             return image.resize((size[1], size[0]), resample=resample)
         else:
-            if image.ndim == 2:
-                image = np.expand_dims(image, axis=-1)
-            channel_count = image.shape[-1]
-            image = to_pil_image(image, input_data_format=input_data_format, image_mode=image_modes[channel_count])
+            image_mode = None
+            if image.ndim == 2 or image.shape[-1] == 1:
+                image_mode = 'P'
+            image = to_pil_image(image, input_data_format=input_data_format, image_mode=image_mode)
 
         resized_image = image.resize((size[1], size[0]), resample=resample)
         resized_array = np.array(resized_image)