karolmajek/Axial-DeepLab-SWideRNet
0
1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""Contains common utility functions and classes for building dataset."""17 18import collections19import io20 21import numpy as np22from PIL import Image23from PIL import ImageOps24import tensorflow as tf25 26from deeplab2 import common27 28_PANOPTIC_LABEL_FORMAT = 'raw'29 30 31def read_image(image_data):32 """Decodes image from in-memory data.33 34 Args:35 image_data: Bytes data representing encoded image.36 37 Returns:38 Decoded PIL.Image object.39 """40 image = Image.open(io.BytesIO(image_data))41 42 try:43 image = ImageOps.exif_transpose(image)44 except TypeError:45 # capture and ignore this bug:46 # https://github.com/python-pillow/Pillow/issues/397347 pass48 49 return image50 51 52def get_image_dims(image_data, check_is_rgb=False):53 """Decodes image and return its height and width.54 55 Args:56 image_data: Bytes data representing encoded image.57 check_is_rgb: Whether to check encoded image is RGB.58 59 Returns:60 Decoded image size as a tuple of (height, width)61 62 Raises:63 ValueError: If check_is_rgb is set and input image has other format.64 """65 image = read_image(image_data)66 67 if check_is_rgb and image.mode != 'RGB':68 raise ValueError('Expects RGB image data, gets mode: %s' % image.mode)69 70 width, height = image.size71 return height, width72 73 74def _int64_list_feature(values):75 """Returns a TF-Feature of int64_list.76 77 Args:78 values: A scalar or an iterable of integer values.79 80 Returns:81 A TF-Feature.82 """83 if not isinstance(values, collections.Iterable):84 values = [values]85 86 return tf.train.Feature(int64_list=tf.train.Int64List(value=values))87 88 89def _bytes_list_feature(values):90 """Returns a TF-Feature of bytes.91 92 Args:93 values: A string.94 95 Returns:96 A TF-Feature.97 """98 if isinstance(values, str):99 values = values.encode()100 101 return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))102 103 104def create_features(image_data,105 image_format,106 filename,107 label_data=None,108 label_format=None):109 """Creates image/segmentation features.110 111 Args:112 image_data: String or byte stream of encoded image data.113 image_format: String, image data format, should be either 'jpeg' or 'png'.114 filename: String, image filename.115 label_data: String or byte stream of (potentially) encoded label data. If116 None, we skip to write it to tf.train.Example.117 label_format: String, label data format, should be either 'png' or 'raw'. If118 None, we skip to write it to tf.train.Example.119 120 Returns:121 A dictionary of feature name to tf.train.Feature maaping.122 """123 if image_format not in ('jpeg', 'png'):124 raise ValueError('Unsupported image format: %s' % image_format)125 126 # Check color mode, and convert grey image to rgb image.127 image = read_image(image_data)128 if image.mode != 'RGB':129 image = image.convert('RGB')130 image_data = io.BytesIO()131 image.save(image_data, format=image_format)132 image_data = image_data.getvalue()133 134 height, width = get_image_dims(image_data, check_is_rgb=True)135 136 feature_dict = {137 common.KEY_ENCODED_IMAGE: _bytes_list_feature(image_data),138 common.KEY_IMAGE_FILENAME: _bytes_list_feature(filename),139 common.KEY_IMAGE_FORMAT: _bytes_list_feature(image_format),140 common.KEY_IMAGE_HEIGHT: _int64_list_feature(height),141 common.KEY_IMAGE_WIDTH: _int64_list_feature(width),142 common.KEY_IMAGE_CHANNELS: _int64_list_feature(3),143 }144 145 if label_data is None:146 return feature_dict147 148 if label_format == 'png':149 label_height, label_width = get_image_dims(label_data)150 if (label_height, label_width) != (height, width):151 raise ValueError('Image (%s) and label (%s) shape mismatch' %152 ((height, width), (label_height, label_width)))153 elif label_format == 'raw':154 # Raw label encodes int32 array.155 expected_label_size = height * width * np.dtype(np.int32).itemsize156 if len(label_data) != expected_label_size:157 raise ValueError('Expects raw label data length %d, gets %d' %158 (expected_label_size, len(label_data)))159 else:160 raise ValueError('Unsupported label format: %s' % label_format)161 162 feature_dict.update({163 common.KEY_ENCODED_LABEL: _bytes_list_feature(label_data),164 common.KEY_LABEL_FORMAT: _bytes_list_feature(label_format)165 })166 167 return feature_dict168 169 170def create_tfexample(image_data,171 image_format,172 filename,173 label_data=None,174 label_format=None):175 """Converts one image/segmentation pair to TF example.176 177 Args:178 image_data: String or byte stream of encoded image data.179 image_format: String, image data format, should be either 'jpeg' or 'png'.180 filename: String, image filename.181 label_data: String or byte stream of (potentially) encoded label data. If182 None, we skip to write it to tf.train.Example.183 label_format: String, label data format, should be either 'png' or 'raw'. If184 None, we skip to write it to tf.train.Example.185 186 Returns:187 TF example proto.188 """189 feature_dict = create_features(image_data, image_format, filename, label_data,190 label_format)191 return tf.train.Example(features=tf.train.Features(feature=feature_dict))192 193 194def create_video_tfexample(image_data,195 image_format,196 filename,197 sequence_id,198 image_id,199 label_data=None,200 label_format=None,201 prev_image_data=None,202 prev_label_data=None):203 """Converts one video frame/panoptic segmentation pair to TF example.204 205 Args:206 image_data: String or byte stream of encoded image data.207 image_format: String, image data format, should be either 'jpeg' or 'png'.208 filename: String, image filename.209 sequence_id: ID of the video sequence as a string.210 image_id: ID of the image as a string.211 label_data: String or byte stream of (potentially) encoded label data. If212 None, we skip to write it to tf.train.Example.213 label_format: String, label data format, should be either 'png' or 'raw'. If214 None, we skip to write it to tf.train.Example.215 prev_image_data: An optional string or byte stream of encoded previous image216 data.217 prev_label_data: An optional string or byte stream of (potentially) encoded218 previous label data.219 220 Returns:221 TF example proto.222 """223 feature_dict = create_features(image_data, image_format, filename, label_data,224 label_format)225 feature_dict.update({226 common.KEY_SEQUENCE_ID: _bytes_list_feature(sequence_id),227 common.KEY_FRAME_ID: _bytes_list_feature(image_id)228 })229 if prev_image_data is not None:230 feature_dict[common.KEY_ENCODED_PREV_IMAGE] = _bytes_list_feature(231 prev_image_data)232 if prev_label_data is not None:233 feature_dict[common.KEY_ENCODED_PREV_LABEL] = _bytes_list_feature(234 prev_label_data)235 return tf.train.Example(features=tf.train.Features(feature=feature_dict))236 237 238def create_video_and_depth_tfexample(image_data,239 image_format,240 filename,241 sequence_id,242 image_id,243 label_data=None,244 label_format=None,245 next_image_data=None,246 next_label_data=None,247 depth_data=None,248 depth_format=None):249 """Converts an image/segmentation pair and depth of first frame to TF example.250 251 The image pair contains the current frame and the next frame with the252 current frame including depth label.253 254 Args:255 image_data: String or byte stream of encoded image data.256 image_format: String, image data format, should be either 'jpeg' or 'png'.257 filename: String, image filename.258 sequence_id: ID of the video sequence as a string.259 image_id: ID of the image as a string.260 label_data: String or byte stream of (potentially) encoded label data. If261 None, we skip to write it to tf.train.Example.262 label_format: String, label data format, should be either 'png' or 'raw'. If263 None, we skip to write it to tf.train.Example.264 next_image_data: An optional string or byte stream of encoded next image265 data.266 next_label_data: An optional string or byte stream of (potentially) encoded267 next label data.268 depth_data: An optional string or byte sream of encoded depth data.269 depth_format: String, depth data format, should be either 'png' or 'raw'.270 271 Returns:272 TF example proto.273 """274 feature_dict = create_features(image_data, image_format, filename, label_data,275 label_format)276 feature_dict.update({277 common.KEY_SEQUENCE_ID: _bytes_list_feature(sequence_id),278 common.KEY_FRAME_ID: _bytes_list_feature(image_id)279 })280 if next_image_data is not None:281 feature_dict[common.KEY_ENCODED_NEXT_IMAGE] = _bytes_list_feature(282 next_image_data)283 if next_label_data is not None:284 feature_dict[common.KEY_ENCODED_NEXT_LABEL] = _bytes_list_feature(285 next_label_data)286 if depth_data is not None:287 feature_dict[common.KEY_ENCODED_DEPTH] = _bytes_list_feature(288 depth_data)289 feature_dict[common.KEY_DEPTH_FORMAT] = _bytes_list_feature(290 depth_format)291 return tf.train.Example(features=tf.train.Features(feature=feature_dict))292 293 294class SegmentationDecoder(object):295 """Basic parser to decode serialized tf.Example."""296 297 def __init__(self,298 is_panoptic_dataset=True,299 is_video_dataset=False,300 use_two_frames=False,301 use_next_frame=False,302 decode_groundtruth_label=True):303 self._is_panoptic_dataset = is_panoptic_dataset304 self._is_video_dataset = is_video_dataset305 self._use_two_frames = use_two_frames306 self._use_next_frame = use_next_frame307 self._decode_groundtruth_label = decode_groundtruth_label308 string_feature = tf.io.FixedLenFeature((), tf.string)309 int_feature = tf.io.FixedLenFeature((), tf.int64)310 self._keys_to_features = {311 common.KEY_ENCODED_IMAGE: string_feature,312 common.KEY_IMAGE_FILENAME: string_feature,313 common.KEY_IMAGE_FORMAT: string_feature,314 common.KEY_IMAGE_HEIGHT: int_feature,315 common.KEY_IMAGE_WIDTH: int_feature,316 common.KEY_IMAGE_CHANNELS: int_feature,317 }318 if decode_groundtruth_label:319 self._keys_to_features[common.KEY_ENCODED_LABEL] = string_feature320 if self._is_video_dataset:321 self._keys_to_features[common.KEY_SEQUENCE_ID] = string_feature322 self._keys_to_features[common.KEY_FRAME_ID] = string_feature323 # Two-frame specific processing.324 if self._use_two_frames:325 self._keys_to_features[common.KEY_ENCODED_PREV_IMAGE] = string_feature326 if decode_groundtruth_label:327 self._keys_to_features[common.KEY_ENCODED_PREV_LABEL] = string_feature328 # Next-frame specific processing.329 if self._use_next_frame:330 self._keys_to_features[common.KEY_ENCODED_NEXT_IMAGE] = string_feature331 if decode_groundtruth_label:332 self._keys_to_features[common.KEY_ENCODED_NEXT_LABEL] = string_feature333 334 def _decode_image(self, parsed_tensors, key):335 """Decodes image udner key from parsed tensors."""336 image = tf.io.decode_image(337 parsed_tensors[key],338 channels=3,339 dtype=tf.dtypes.uint8,340 expand_animations=False)341 image.set_shape([None, None, 3])342 return image343 344 def _decode_label(self, parsed_tensors, label_key):345 """Decodes segmentation label under label_key from parsed tensors."""346 if self._is_panoptic_dataset:347 flattened_label = tf.io.decode_raw(348 parsed_tensors[label_key], out_type=tf.int32)349 label_shape = tf.stack([350 parsed_tensors[common.KEY_IMAGE_HEIGHT],351 parsed_tensors[common.KEY_IMAGE_WIDTH], 1352 ])353 label = tf.reshape(flattened_label, label_shape)354 return label355 356 label = tf.io.decode_image(parsed_tensors[label_key], channels=1)357 label.set_shape([None, None, 1])358 return label359 360 def __call__(self, serialized_example):361 parsed_tensors = tf.io.parse_single_example(362 serialized_example, features=self._keys_to_features)363 return_dict = {364 'image':365 self._decode_image(parsed_tensors, common.KEY_ENCODED_IMAGE),366 'image_name':367 parsed_tensors[common.KEY_IMAGE_FILENAME],368 'height':369 tf.cast(parsed_tensors[common.KEY_IMAGE_HEIGHT], dtype=tf.int32),370 'width':371 tf.cast(parsed_tensors[common.KEY_IMAGE_WIDTH], dtype=tf.int32),372 }373 return_dict['label'] = None374 if self._decode_groundtruth_label:375 return_dict['label'] = self._decode_label(parsed_tensors,376 common.KEY_ENCODED_LABEL)377 if self._is_video_dataset:378 return_dict['sequence'] = parsed_tensors[common.KEY_SEQUENCE_ID]379 if self._use_two_frames:380 return_dict['prev_image'] = self._decode_image(381 parsed_tensors, common.KEY_ENCODED_PREV_IMAGE)382 if self._decode_groundtruth_label:383 return_dict['prev_label'] = self._decode_label(384 parsed_tensors, common.KEY_ENCODED_PREV_LABEL)385 if self._use_next_frame:386 return_dict['next_image'] = self._decode_image(387 parsed_tensors, common.KEY_ENCODED_NEXT_IMAGE)388 if self._decode_groundtruth_label:389 return_dict['next_label'] = self._decode_label(390 parsed_tensors, common.KEY_ENCODED_NEXT_LABEL)391 return return_dict392