Hugging Face's logo

Spaces:
akhaliq
/
deeplab2
Runtime error

App Files Community
deeplab2 / data /data_utils.py
akhaliq3
spaces demo
506da10
raw history blame contribute delete
No virus
14.4 kB
# coding=utf-8
# Copyright 2021 The Deeplab2 Authors.
#
# 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.
"""Contains common utility functions and classes for building dataset."""
import collections
import io
import numpy as np
from PIL import Image
from PIL import ImageOps
import tensorflow as tf
from deeplab2 import common
_PANOPTIC_LABEL_FORMAT = 'raw'
def read_image(image_data):
"""Decodes image from in-memory data.
Args:
image_data: Bytes data representing encoded image.
Returns:
Decoded PIL.Image object.
"""
image = Image.open(io.BytesIO(image_data))
try:
image = ImageOps.exif_transpose(image)
except TypeError:
# capture and ignore this bug:
# https://github.com/python-pillow/Pillow/issues/3973
pass
return image
def get_image_dims(image_data, check_is_rgb=False):
"""Decodes image and return its height and width.
Args:
image_data: Bytes data representing encoded image.
check_is_rgb: Whether to check encoded image is RGB.
Returns:
Decoded image size as a tuple of (height, width)
Raises:
ValueError: If check_is_rgb is set and input image has other format.
"""
image = read_image(image_data)
if check_is_rgb and image.mode != 'RGB':
raise ValueError('Expects RGB image data, gets mode: %s' % image.mode)
width, height = image.size
return height, width
def _int64_list_feature(values):
"""Returns a TF-Feature of int64_list.
Args:
values: A scalar or an iterable of integer values.
Returns:
A TF-Feature.
"""
if not isinstance(values, collections.Iterable):
values = [values]
return tf.train.Feature(int64_list=tf.train.Int64List(value=values))
def _bytes_list_feature(values):
"""Returns a TF-Feature of bytes.
Args:
values: A string.
Returns:
A TF-Feature.
"""
if isinstance(values, str):
values = values.encode()
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))
def create_features(image_data,
image_format,
filename,
label_data=None,
label_format=None):
"""Creates image/segmentation features.
Args:
image_data: String or byte stream of encoded image data.
image_format: String, image data format, should be either 'jpeg' or 'png'.
filename: String, image filename.
label_data: String or byte stream of (potentially) encoded label data. If
None, we skip to write it to tf.train.Example.
label_format: String, label data format, should be either 'png' or 'raw'. If
None, we skip to write it to tf.train.Example.
Returns:
A dictionary of feature name to tf.train.Feature maaping.
"""
if image_format not in ('jpeg', 'png'):
raise ValueError('Unsupported image format: %s' % image_format)
# Check color mode, and convert grey image to rgb image.
image = read_image(image_data)
if image.mode != 'RGB':
image = image.convert('RGB')
image_data = io.BytesIO()
image.save(image_data, format=image_format)
image_data = image_data.getvalue()
height, width = get_image_dims(image_data, check_is_rgb=True)
feature_dict = {
common.KEY_ENCODED_IMAGE: _bytes_list_feature(image_data),
common.KEY_IMAGE_FILENAME: _bytes_list_feature(filename),
common.KEY_IMAGE_FORMAT: _bytes_list_feature(image_format),
common.KEY_IMAGE_HEIGHT: _int64_list_feature(height),
common.KEY_IMAGE_WIDTH: _int64_list_feature(width),
common.KEY_IMAGE_CHANNELS: _int64_list_feature(3),
}
if label_data is None:
return feature_dict
if label_format == 'png':
label_height, label_width = get_image_dims(label_data)
if (label_height, label_width) != (height, width):
raise ValueError('Image (%s) and label (%s) shape mismatch' %
((height, width), (label_height, label_width)))
elif label_format == 'raw':
# Raw label encodes int32 array.
expected_label_size = height * width * np.dtype(np.int32).itemsize
if len(label_data) != expected_label_size:
raise ValueError('Expects raw label data length %d, gets %d' %
(expected_label_size, len(label_data)))
else:
raise ValueError('Unsupported label format: %s' % label_format)
feature_dict.update({
common.KEY_ENCODED_LABEL: _bytes_list_feature(label_data),
common.KEY_LABEL_FORMAT: _bytes_list_feature(label_format)
})
return feature_dict
def create_tfexample(image_data,
image_format,
filename,
label_data=None,
label_format=None):
"""Converts one image/segmentation pair to TF example.
Args:
image_data: String or byte stream of encoded image data.
image_format: String, image data format, should be either 'jpeg' or 'png'.
filename: String, image filename.
label_data: String or byte stream of (potentially) encoded label data. If
None, we skip to write it to tf.train.Example.
label_format: String, label data format, should be either 'png' or 'raw'. If
None, we skip to write it to tf.train.Example.
Returns:
TF example proto.
"""
feature_dict = create_features(image_data, image_format, filename, label_data,
label_format)
return tf.train.Example(features=tf.train.Features(feature=feature_dict))
def create_video_tfexample(image_data,
image_format,
filename,
sequence_id,
image_id,
label_data=None,
label_format=None,
prev_image_data=None,
prev_label_data=None):
"""Converts one video frame/panoptic segmentation pair to TF example.
Args:
image_data: String or byte stream of encoded image data.
image_format: String, image data format, should be either 'jpeg' or 'png'.
filename: String, image filename.
sequence_id: ID of the video sequence as a string.
image_id: ID of the image as a string.
label_data: String or byte stream of (potentially) encoded label data. If
None, we skip to write it to tf.train.Example.
label_format: String, label data format, should be either 'png' or 'raw'. If
None, we skip to write it to tf.train.Example.
prev_image_data: An optional string or byte stream of encoded previous image
data.
prev_label_data: An optional string or byte stream of (potentially) encoded
previous label data.
Returns:
TF example proto.
"""
feature_dict = create_features(image_data, image_format, filename, label_data,
label_format)
feature_dict.update({
common.KEY_SEQUENCE_ID: _bytes_list_feature(sequence_id),
common.KEY_FRAME_ID: _bytes_list_feature(image_id)
})
if prev_image_data is not None:
feature_dict[common.KEY_ENCODED_PREV_IMAGE] = _bytes_list_feature(
prev_image_data)
if prev_label_data is not None:
feature_dict[common.KEY_ENCODED_PREV_LABEL] = _bytes_list_feature(
prev_label_data)
return tf.train.Example(features=tf.train.Features(feature=feature_dict))
def create_video_and_depth_tfexample(image_data,
image_format,
filename,
sequence_id,
image_id,
label_data=None,
label_format=None,
next_image_data=None,
next_label_data=None,
depth_data=None,
depth_format=None):
"""Converts an image/segmentation pair and depth of first frame to TF example.
The image pair contains the current frame and the next frame with the
current frame including depth label.
Args:
image_data: String or byte stream of encoded image data.
image_format: String, image data format, should be either 'jpeg' or 'png'.
filename: String, image filename.
sequence_id: ID of the video sequence as a string.
image_id: ID of the image as a string.
label_data: String or byte stream of (potentially) encoded label data. If
None, we skip to write it to tf.train.Example.
label_format: String, label data format, should be either 'png' or 'raw'. If
None, we skip to write it to tf.train.Example.
next_image_data: An optional string or byte stream of encoded next image
data.
next_label_data: An optional string or byte stream of (potentially) encoded
next label data.
depth_data: An optional string or byte sream of encoded depth data.
depth_format: String, depth data format, should be either 'png' or 'raw'.
Returns:
TF example proto.
"""
feature_dict = create_features(image_data, image_format, filename, label_data,
label_format)
feature_dict.update({
common.KEY_SEQUENCE_ID: _bytes_list_feature(sequence_id),
common.KEY_FRAME_ID: _bytes_list_feature(image_id)
})
if next_image_data is not None:
feature_dict[common.KEY_ENCODED_NEXT_IMAGE] = _bytes_list_feature(
next_image_data)
if next_label_data is not None:
feature_dict[common.KEY_ENCODED_NEXT_LABEL] = _bytes_list_feature(
next_label_data)
if depth_data is not None:
feature_dict[common.KEY_ENCODED_DEPTH] = _bytes_list_feature(
depth_data)
feature_dict[common.KEY_DEPTH_FORMAT] = _bytes_list_feature(
depth_format)
return tf.train.Example(features=tf.train.Features(feature=feature_dict))
class SegmentationDecoder(object):
"""Basic parser to decode serialized tf.Example."""
def __init__(self,
is_panoptic_dataset=True,
is_video_dataset=False,
use_two_frames=False,
use_next_frame=False,
decode_groundtruth_label=True):
self._is_panoptic_dataset = is_panoptic_dataset
self._is_video_dataset = is_video_dataset
self._use_two_frames = use_two_frames
self._use_next_frame = use_next_frame
self._decode_groundtruth_label = decode_groundtruth_label
string_feature = tf.io.FixedLenFeature((), tf.string)
int_feature = tf.io.FixedLenFeature((), tf.int64)
self._keys_to_features = {
common.KEY_ENCODED_IMAGE: string_feature,
common.KEY_IMAGE_FILENAME: string_feature,
common.KEY_IMAGE_FORMAT: string_feature,
common.KEY_IMAGE_HEIGHT: int_feature,
common.KEY_IMAGE_WIDTH: int_feature,
common.KEY_IMAGE_CHANNELS: int_feature,
}
if decode_groundtruth_label:
self._keys_to_features[common.KEY_ENCODED_LABEL] = string_feature
if self._is_video_dataset:
self._keys_to_features[common.KEY_SEQUENCE_ID] = string_feature
self._keys_to_features[common.KEY_FRAME_ID] = string_feature
# Two-frame specific processing.
if self._use_two_frames:
self._keys_to_features[common.KEY_ENCODED_PREV_IMAGE] = string_feature
if decode_groundtruth_label:
self._keys_to_features[common.KEY_ENCODED_PREV_LABEL] = string_feature
# Next-frame specific processing.
if self._use_next_frame:
self._keys_to_features[common.KEY_ENCODED_NEXT_IMAGE] = string_feature
if decode_groundtruth_label:
self._keys_to_features[common.KEY_ENCODED_NEXT_LABEL] = string_feature
def _decode_image(self, parsed_tensors, key):
"""Decodes image udner key from parsed tensors."""
image = tf.io.decode_image(
parsed_tensors[key],
channels=3,
dtype=tf.dtypes.uint8,
expand_animations=False)
image.set_shape([None, None, 3])
return image
def _decode_label(self, parsed_tensors, label_key):
"""Decodes segmentation label under label_key from parsed tensors."""
if self._is_panoptic_dataset:
flattened_label = tf.io.decode_raw(
parsed_tensors[label_key], out_type=tf.int32)
label_shape = tf.stack([
parsed_tensors[common.KEY_IMAGE_HEIGHT],
parsed_tensors[common.KEY_IMAGE_WIDTH], 1
])
label = tf.reshape(flattened_label, label_shape)
return label
label = tf.io.decode_image(parsed_tensors[label_key], channels=1)
label.set_shape([None, None, 1])
return label
def __call__(self, serialized_example):
parsed_tensors = tf.io.parse_single_example(
serialized_example, features=self._keys_to_features)
return_dict = {
'image':
self._decode_image(parsed_tensors, common.KEY_ENCODED_IMAGE),
'image_name':
parsed_tensors[common.KEY_IMAGE_FILENAME],
'height':
tf.cast(parsed_tensors[common.KEY_IMAGE_HEIGHT], dtype=tf.int32),
'width':
tf.cast(parsed_tensors[common.KEY_IMAGE_WIDTH], dtype=tf.int32),
}
return_dict['label'] = None
if self._decode_groundtruth_label:
return_dict['label'] = self._decode_label(parsed_tensors,
common.KEY_ENCODED_LABEL)
if self._is_video_dataset:
return_dict['sequence'] = parsed_tensors[common.KEY_SEQUENCE_ID]
if self._use_two_frames:
return_dict['prev_image'] = self._decode_image(
parsed_tensors, common.KEY_ENCODED_PREV_IMAGE)
if self._decode_groundtruth_label:
return_dict['prev_label'] = self._decode_label(
parsed_tensors, common.KEY_ENCODED_PREV_LABEL)
if self._use_next_frame:
return_dict['next_image'] = self._decode_image(
parsed_tensors, common.KEY_ENCODED_NEXT_IMAGE)
if self._decode_groundtruth_label:
return_dict['next_label'] = self._decode_label(
parsed_tensors, common.KEY_ENCODED_NEXT_LABEL)
return return_dict