|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Image-related utilities that are useful to prepare dataset."""
|
|
|
| import dataclasses
|
| import imghdr
|
| import io
|
| from typing import Optional, Tuple
|
|
|
| import numpy as np
|
| from PIL import Image
|
|
|
|
|
| @dataclasses.dataclass
|
| class ImageFormat:
|
| """Supported image formats.
|
|
|
| For model development, this library should support the same image formats as
|
| `tf.io.decode_image`[1].
|
|
|
| [1]: https://www.tensorflow.org/api_docs/python/tf/io/decode_image
|
| """
|
| bmp: str = 'BMP'
|
| png: str = 'PNG'
|
| jpeg: str = 'JPEG'
|
| raw: str = 'RAW'
|
|
|
|
|
| def validate_image_format(format_str: str) -> str:
|
| """Validates `format_str` and returns canonical format.
|
|
|
| This function accepts image format in lower case and will returns the upper
|
| case string as canonical format.
|
|
|
| Args:
|
| format_str: Image format string.
|
|
|
| Returns:
|
| Canonical image format string.
|
|
|
| Raises:
|
| ValueError: If the canonical format is not listed in `ImageFormat`.
|
| """
|
| canonical_format = format_str.upper()
|
| if canonical_format in dataclasses.asdict(ImageFormat()).values():
|
| return canonical_format
|
| raise ValueError(f'Image format is invalid: {format_str}')
|
|
|
|
|
| def encode_image(image_np: np.ndarray, image_format: str) -> bytes:
|
| """Encodes `image_np` specified by `image_format`.
|
|
|
| Args:
|
| image_np: Numpy image array.
|
| image_format: An enum specifying the format of the generated image.
|
|
|
| Returns:
|
| Encoded image string.
|
| """
|
| if image_format == 'RAW':
|
| return image_np.tobytes()
|
|
|
| if len(image_np.shape) > 2 and image_np.shape[2] == 1:
|
| image_pil = Image.fromarray(np.squeeze(image_np), 'L')
|
| else:
|
| image_pil = Image.fromarray(image_np)
|
| with io.BytesIO() as output:
|
| image_pil.save(output, format=validate_image_format(image_format))
|
| return output.getvalue()
|
|
|
|
|
| def decode_image(image_bytes: bytes,
|
| image_format: Optional[str] = None,
|
| image_dtype: str = 'uint8') -> np.ndarray:
|
| """Decodes image_bytes into numpy array."""
|
| if image_format == 'RAW':
|
| return np.frombuffer(image_bytes, dtype=image_dtype)
|
| image_pil = Image.open(io.BytesIO(image_bytes))
|
| image_np = np.array(image_pil)
|
| if len(image_np.shape) < 3:
|
| image_np = image_np[..., np.newaxis]
|
| return image_np
|
|
|
|
|
| def decode_image_metadata(image_bytes: bytes) -> Tuple[int, int, int, str]:
|
| """Decodes image metadata from encoded image string.
|
|
|
| Note that if the image is encoded in RAW format, the metadata cannot be
|
| inferred from the image bytes.
|
|
|
| Args:
|
| image_bytes: Encoded image string.
|
|
|
| Returns:
|
| A tuple of height, width, number of channels, and encoding format.
|
| """
|
| image_np = decode_image(image_bytes)
|
|
|
| height, width, num_channels = image_np.shape
|
| image_format = imghdr.what(file=None, h=image_bytes)
|
| return height, width, num_channels, validate_image_format(image_format)
|
|
|
