Enable crop_center method to handle (W, H, C) images (#17626)

* enable crop_center method to handle (W, H, C) images

* minor style and comment edits

src/transformers/image_utils.py

@@ -163,6 +163,26 @@ class ImageFeatureExtractionMixin:
 
         return image
 
+    def expand_dims(self, image):
+        """
+        Expands 2-dimensional `image` to 3 dimensions.
+
+        Args:
+            image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`):
+                The image to expand.
+        """
+        self._ensure_format_supported(image)
+
+        # Do nothing if PIL image
+        if isinstance(image, PIL.Image.Image):
+            return image
+
+        if is_torch_tensor(image):
+            image = image.unsqueeze(0)
+        else:
+            image = np.expand_dims(image, axis=0)
+        return image
+
     def normalize(self, image, mean, std):
         """
         Normalizes `image` with `mean` and `std`. Note that this will trigger a conversion of `image` to a NumPy array
@@ -201,7 +221,7 @@ class ImageFeatureExtractionMixin:
 
     def resize(self, image, size, resample=PIL.Image.BILINEAR, default_to_square=True, max_size=None):
         """
-        Resizes `image`. Note that this will trigger a conversion of `image` to a PIL Image.
+        Resizes `image`. Enforces conversion of input to PIL.Image.
 
         Args:
             image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`):
@@ -225,6 +245,9 @@ class ImageFeatureExtractionMixin:
                 greater than `max_size` after being resized according to `size`, then the image is resized again so
                 that the longer edge is equal to `max_size`. As a result, `size` might be overruled, i.e the smaller
                 edge may be shorter than `size`. Only used if `default_to_square` is `False`.
+
+        Returns:
+            image: A resized `PIL.Image.Image`.
         """
         self._ensure_format_supported(image)
 
@@ -267,17 +290,28 @@ class ImageFeatureExtractionMixin:
         size given, it will be padded (so the returned result has the size asked).
 
         Args:
-            image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`):
+            image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor` of shape (n_channels, height, width) or (height, width, n_channels)):
                 The image to resize.
             size (`int` or `Tuple[int, int]`):
                 The size to which crop the image.
+
+        Returns:
+            new_image: A center cropped `PIL.Image.Image` or `np.ndarray` or `torch.Tensor` of shape: (n_channels,
+            height, width).
         """
         self._ensure_format_supported(image)
+
         if not isinstance(size, tuple):
             size = (size, size)
 
         # PIL Image.size is (width, height) but NumPy array and torch Tensors have (height, width)
-        image_shape = (image.size[1], image.size[0]) if isinstance(image, PIL.Image.Image) else image.shape[-2:]
+        if is_torch_tensor(image) or isinstance(image, np.ndarray):
+            if image.ndim == 2:
+                image = self.expand_dims(image)
+            image_shape = image.shape[1:] if image.shape[0] in [1, 3] else image.shape[:2]
+        else:
+            image_shape = (image.size[1], image.size[0])
+
         top = (image_shape[0] - size[0]) // 2
         bottom = top + size[0]  # In case size is odd, (image_shape[0] + size[0]) // 2 won't give the proper result.
         left = (image_shape[1] - size[1]) // 2
@@ -287,7 +321,17 @@ class ImageFeatureExtractionMixin:
         if isinstance(image, PIL.Image.Image):
             return image.crop((left, top, right, bottom))
 
-        # Check if all the dimensions are inside the image.
+        # Check if image is in (n_channels, height, width) or (height, width, n_channels) format
+        channel_first = True if image.shape[0] in [1, 3] else False
+
+        # Transpose (height, width, n_channels) format images
+        if not channel_first:
+            if isinstance(image, np.ndarray):
+                image = image.transpose(2, 0, 1)
+            if is_torch_tensor(image):
+                image = image.permute(2, 0, 1)
+
+        # Check if cropped area is within image boundaries
         if top >= 0 and bottom <= image_shape[0] and left >= 0 and right <= image_shape[1]:
             return image[..., top:bottom, left:right]
 
@@ -309,6 +353,8 @@ class ImageFeatureExtractionMixin:
         left += left_pad
         right += left_pad
 
-        return new_image[
+        new_image = new_image[
             ..., max(0, top) : min(new_image.shape[-2], bottom), max(0, left) : min(new_image.shape[-1], right)
         ]
+
+        return new_image