Preprocess the image input. Accepted formats are PIL images, NumPy arrays or PyTorch tensors.
(
self,
image: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray],
height: Optional[int] = None,
width: Optional[int] = None,
)
| 263 | return image |
| 264 | |
| 265 | def preprocess( |
| 266 | self, |
| 267 | image: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray], |
| 268 | height: Optional[int] = None, |
| 269 | width: Optional[int] = None, |
| 270 | ) -> torch.Tensor: |
| 271 | """ |
| 272 | Preprocess the image input. Accepted formats are PIL images, NumPy arrays or PyTorch tensors. |
| 273 | """ |
| 274 | supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor) |
| 275 | |
| 276 | # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image |
| 277 | if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3: |
| 278 | if isinstance(image, torch.Tensor): |
| 279 | # if image is a pytorch tensor could have 2 possible shapes: |
| 280 | # 1. batch x height x width: we should insert the channel dimension at position 1 |
| 281 | # 2. channnel x height x width: we should insert batch dimension at position 0, |
| 282 | # however, since both channel and batch dimension has same size 1, it is same to insert at position 1 |
| 283 | # for simplicity, we insert a dimension of size 1 at position 1 for both cases |
| 284 | image = image.unsqueeze(1) |
| 285 | else: |
| 286 | # if it is a numpy array, it could have 2 possible shapes: |
| 287 | # 1. batch x height x width: insert channel dimension on last position |
| 288 | # 2. height x width x channel: insert batch dimension on first position |
| 289 | if image.shape[-1] == 1: |
| 290 | image = np.expand_dims(image, axis=0) |
| 291 | else: |
| 292 | image = np.expand_dims(image, axis=-1) |
| 293 | |
| 294 | if isinstance(image, supported_formats): |
| 295 | image = [image] |
| 296 | elif not (isinstance(image, list) and all(isinstance(i, supported_formats) for i in image)): |
| 297 | raise ValueError( |
| 298 | f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support {', '.join(supported_formats)}" |
| 299 | ) |
| 300 | |
| 301 | if isinstance(image[0], PIL.Image.Image): |
| 302 | if self.config.do_convert_rgb: |
| 303 | image = [self.convert_to_rgb(i) for i in image] |
| 304 | elif self.config.do_convert_grayscale: |
| 305 | image = [self.convert_to_grayscale(i) for i in image] |
| 306 | if self.config.do_resize: |
| 307 | height, width = self.get_default_height_width(image[0], height, width) |
| 308 | image = [self.resize(i, height, width) for i in image] |
| 309 | image = self.pil_to_numpy(image) # to np |
| 310 | image = self.numpy_to_pt(image) # to pt |
| 311 | |
| 312 | elif isinstance(image[0], np.ndarray): |
| 313 | image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) |
| 314 | |
| 315 | image = self.numpy_to_pt(image) |
| 316 | |
| 317 | height, width = self.get_default_height_width(image, height, width) |
| 318 | if self.config.do_resize: |
| 319 | image = self.resize(image, height, width) |
| 320 | |
| 321 | elif isinstance(image[0], torch.Tensor): |
| 322 | image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) |