Spaces:
Sleeping
Sleeping
| import numpy as np | |
| from image import load_image, ImageTyping | |
| __all__ = [ | |
| 'rgb_encode', | |
| ] | |
| _DEFAULT_ORDER = 'HWC' | |
| def _get_hwc_map(order_: str): | |
| return tuple(_DEFAULT_ORDER.index(c) for c in order_.upper()) | |
| def rgb_encode(image: ImageTyping, order_: str = 'CHW', use_float: bool = True) -> np.ndarray: | |
| """ | |
| Overview: | |
| Encode image as rgb channels. | |
| :param image: Image to be encoded. | |
| :param order_: Order of encoding, default is ``CHW``. | |
| :param use_float: Use float to represent the channels, default is ``True``. ``np.uint8`` will be used when false. | |
| :return: Encoded rgb image. | |
| Examples:: | |
| >>> from PIL import Image | |
| >>> from encode import rgb_encode | |
| >>> | |
| >>> image = Image.open('custom_image.jpg') | |
| >>> image | |
| <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=1606x1870 at 0x7F9EC37389D0> | |
| >>> | |
| >>> data = rgb_encode(image) | |
| >>> data.shape, data.dtype | |
| ((3, 1870, 1606), dtype('float32')) | |
| >>> data = rgb_encode(image, order_='CHW') | |
| >>> data.shape, data.dtype | |
| ((3, 1870, 1606), dtype('float32')) | |
| >>> data = rgb_encode(image, order_='WHC') | |
| >>> data.shape, data.dtype | |
| ((1606, 1870, 3), dtype('float32')) | |
| >>> data = rgb_encode(image, use_float=False) | |
| >>> data.shape, data.dtype | |
| ((3, 1870, 1606), dtype('uint8')) | |
| .. note:: | |
| The function :func:`rgb_encode`'s result is the same as \ | |
| ``torchvision.transforms.functional import to_tensor``'s result when the given ``image`` is in RGB mode. | |
| """ | |
| image = load_image(image, mode='RGB') | |
| array = np.asarray(image) | |
| array = np.transpose(array, _get_hwc_map(order_)) | |
| if use_float: | |
| array = (array / 255.0).astype(np.float32) | |
| assert array.dtype == np.float32 | |
| else: | |
| assert array.dtype == np.uint8 | |
| return array | |