Spaces:
Runtime error
Runtime error
File size: 3,503 Bytes
5672777 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""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)
# https://pillow.readthedocs.io/en/stable/reference/Image.html#image-attributes
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)
|