official/vision/dataloaders/segmentation_input.py - master - githubmirror/models

official/vision/dataloaders/segmentation_input.py
355 строк13 KB
A. Unique TensorFlower
No public description
09 фев 2026, 19:00
09 фев 2026, 19:00799b0af
# Copyright 2026 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.

"""Data parser and processing for segmentation datasets."""

import tensorflow as tf, tf_keras
from official.vision.configs import semantic_segmentation as config_lib
from official.vision.dataloaders import decoder
from official.vision.dataloaders import parser
from official.vision.dataloaders import utils
from official.vision.ops import preprocess_ops


class Decoder(decoder.Decoder):
  """A tf.Example decoder for segmentation task."""

  def __init__(
      self,
      image_feature=config_lib.DenseFeatureConfig(),
      additional_dense_features=None,
  ):
    self._keys_to_features = {
        'image/encoded': tf.io.FixedLenFeature((), tf.string, default_value=''),
        'image/height': tf.io.FixedLenFeature((), tf.int64, default_value=0),
        'image/width': tf.io.FixedLenFeature((), tf.int64, default_value=0),
        'image/segmentation/class/encoded': tf.io.FixedLenFeature(
            (), tf.string, default_value=''
        ),
        image_feature.feature_name: tf.io.FixedLenFeature(
            (), tf.string, default_value=''
        ),
    }
    if additional_dense_features:
      for feature in additional_dense_features:
        self._keys_to_features[feature.feature_name] = tf.io.FixedLenFeature(
            (), tf.string, default_value=''
        )

  def decode(self, serialized_example):
    return tf.io.parse_single_example(
        serialized_example, self._keys_to_features
    )


class Parser(parser.Parser):
  """Parser to parse an image and its annotations into a dictionary of tensors."""

  def __init__(
      self,
      output_size,
      crop_size=None,
      resize_eval_groundtruth=True,
      gt_is_matting_map=False,
      groundtruth_padded_size=None,
      ignore_label=255,
      aug_rand_hflip=False,
      preserve_aspect_ratio=True,
      aug_scale_min=1.0,
      aug_scale_max=1.0,
      dtype='float32',
      image_feature=config_lib.DenseFeatureConfig(),
      additional_dense_features=None,
      centered_crop=False,
  ):
    """Initializes parameters for parsing annotations in the dataset.

    Args:
      output_size: `Tensor` or `list` for [height, width] of output image. The
        output_size should be divided by the largest feature stride 2^max_level.
      crop_size: `Tensor` or `list` for [height, width] of the crop. If
        specified a training crop of size crop_size is returned. This is useful
        for cropping original images during training while evaluating on
        original image sizes.
      resize_eval_groundtruth: `bool`, if True, eval ground-truth masks are
        resized to output_size.
      gt_is_matting_map: `bool`, if True, the expected mask is in the range
        between 0 and 255. The parser will normalize the value of the mask into
        the range between 0 and 1.
      groundtruth_padded_size: `Tensor` or `list` for [height, width]. When
        resize_eval_groundtruth is set to False, the ground-truth masks are
        padded to this size.
      ignore_label: `int` the pixel with ignore label will not used for training
        and evaluation.
      aug_rand_hflip: `bool`, if True, augment training with random horizontal
        flip.
      preserve_aspect_ratio: `bool`, if True, the aspect ratio is preserved,
        otherwise, the image is resized to output_size.
      aug_scale_min: `float`, the minimum scale applied to `output_size` for
        data augmentation during training.
      aug_scale_max: `float`, the maximum scale applied to `output_size` for
        data augmentation during training.
      dtype: `str`, data type. One of {`bfloat16`, `float32`, `float16`}.
      image_feature: the config for the image input (usually RGB). Defaults to
        the config for a 3-channel image with key = `image/encoded` and ImageNet
        dataset mean/stddev.
      additional_dense_features: `list` of DenseFeatureConfig for additional
        dense features.
      centered_crop: If `centered_crop` is set to True, then resized crop (if
        smaller than padded size) is place in the center of the image. Default
        behaviour is to place it at left top corner.
    """
    self._output_size = output_size
    self._crop_size = crop_size
    self._resize_eval_groundtruth = resize_eval_groundtruth
    if (not resize_eval_groundtruth) and (groundtruth_padded_size is None):
      raise ValueError(
          'groundtruth_padded_size ([height, width]) needs to be'
          'specified when resize_eval_groundtruth is False.'
      )
    self._gt_is_matting_map = gt_is_matting_map
    self._groundtruth_padded_size = groundtruth_padded_size
    self._ignore_label = ignore_label
    self._preserve_aspect_ratio = preserve_aspect_ratio

    # Data augmentation.
    self._aug_rand_hflip = aug_rand_hflip
    self._aug_scale_min = aug_scale_min
    self._aug_scale_max = aug_scale_max

    # dtype.
    self._dtype = dtype

    self._image_feature = image_feature
    self._additional_dense_features = additional_dense_features
    self._centered_crop = centered_crop
    if self._centered_crop and not self._resize_eval_groundtruth:
      raise ValueError(
          'centered_crop is only supported when resize_eval_groundtruth is'
          ' True.'
      )

  def _prepare_image_and_label(self, data):
    """Prepare normalized image and label."""
    height = data['image/height']
    width = data['image/width']

    label = tf.io.decode_image(
        data['image/segmentation/class/encoded'], channels=1
    )
    label = tf.reshape(label, (1, height, width))
    label = tf.cast(label, tf.float32)

    image = tf.io.decode_image(
        data[self._image_feature.feature_name],
        channels=self._image_feature.num_channels,
        dtype=tf.uint8,
    )
    image = tf.reshape(image, (height, width, self._image_feature.num_channels))
    # Normalizes the image feature.
    # The mean and stddev values are divided by 255 to ensure correct
    # normalization, as the input `uint8` image is automatically converted to
    # `float32` and rescaled to values in the range [0, 1] before the
    # normalization happens (as a pre-processing step). So, we re-scale the
    # mean and stddev values to the range [0, 1] beforehand.
    # See `preprocess_ops.normalize_image` for details on the expected ranges
    # for the image mean (`offset`) and stddev (`scale`).
    image = preprocess_ops.normalize_image(
        image,
        [mean / 255.0 for mean in self._image_feature.mean],
        [stddev / 255.0 for stddev in self._image_feature.stddev],
    )

    if self._additional_dense_features:
      input_list = [image]
      for feature_cfg in self._additional_dense_features:
        feature = tf.io.decode_image(
            data[feature_cfg.feature_name],
            channels=feature_cfg.num_channels,
            dtype=tf.uint8,
        )
        feature = tf.reshape(feature, (height, width, feature_cfg.num_channels))
        feature = preprocess_ops.normalize_image(
            feature,
            [mean / 255.0 for mean in feature_cfg.mean],
            [stddev / 255.0 for stddev in feature_cfg.stddev],
        )
        input_list.append(feature)
      concat_input = tf.concat(input_list, axis=2)
    else:
      concat_input = image

    if not self._preserve_aspect_ratio:
      label = tf.reshape(label, [data['image/height'], data['image/width'], 1])
      concat_input = tf.image.resize(
          concat_input, self._output_size, method='bilinear'
      )
      label = tf.image.resize(label, self._output_size, method='nearest')
      label = tf.reshape(label[:, :, -1], [1] + self._output_size)

    return concat_input, label

  def _parse_train_data(self, data):
    """Parses data for training and evaluation."""
    image, label = self._prepare_image_and_label(data)

    # Normalize the label into the range of 0 and 1 for matting ground-truth.
    # Note that the input ground-truth labels must be 0 to 255, and do not
    # contain ignore_label. For gt_is_matting_map case, ignore_label is only
    # used for padding the labels.
    if self._gt_is_matting_map:
      scale = tf.constant(255.0, dtype=tf.float32)
      scale = tf.expand_dims(scale, axis=0)
      scale = tf.expand_dims(scale, axis=0)
      label = tf.cast(label, tf.float32) / scale

    if self._crop_size:

      label = tf.reshape(label, [data['image/height'], data['image/width'], 1])
      # If output_size is specified, resize image, and label to desired
      # output_size.
      if self._output_size:
        image = tf.image.resize(image, self._output_size, method='bilinear')
        label = tf.image.resize(label, self._output_size, method='nearest')

      image_mask = tf.concat([image, label], axis=2)
      image_mask_crop = tf.image.random_crop(
          image_mask, self._crop_size + [tf.shape(image_mask)[-1]]
      )
      image = image_mask_crop[:, :, :-1]
      label = tf.reshape(image_mask_crop[:, :, -1], [1] + self._crop_size)

    # Flips image randomly during training.
    if self._aug_rand_hflip:
      image, _, label = preprocess_ops.random_horizontal_flip(
          image, masks=label
      )

    train_image_size = self._crop_size if self._crop_size else self._output_size
    # Resizes and crops image.
    image, image_info = preprocess_ops.resize_and_crop_image(
        image,
        train_image_size,
        train_image_size,
        aug_scale_min=self._aug_scale_min,
        aug_scale_max=self._aug_scale_max,
        centered_crop=self._centered_crop,
    )

    # Resizes and crops boxes.
    image_scale = image_info[2, :]
    offset = image_info[3, :]

    # Pad label and make sure the padded region assigned to the ignore label.
    # The label is first offset by +1 and then padded with 0.
    label += 1
    label = tf.expand_dims(label, axis=3)
    label = preprocess_ops.resize_and_crop_masks(
        label,
        image_scale,
        train_image_size,
        offset,
        centered_crop=self._centered_crop,
    )
    label -= 1
    label = tf.where(
        tf.equal(label, -1), self._ignore_label * tf.ones_like(label), label
    )
    label = tf.squeeze(label, axis=0)
    valid_mask = tf.not_equal(label, self._ignore_label)

    labels = {
        'masks': label,
        'valid_masks': valid_mask,
        'image_info': image_info,
    }

    # Cast image as self._dtype
    image = tf.cast(image, dtype=self._dtype)

    return image, labels

  def _parse_eval_data(self, data):
    """Parses data for training and evaluation."""
    image, label = self._prepare_image_and_label(data)

    # Binarize mask if ground-truth is a matting map
    if self._gt_is_matting_map:
      label = tf.divide(tf.cast(label, dtype=tf.float32), 255.0)
      label = utils.binarize_matting_map(label)

    # The label is first offset by +1 and then padded with 0.
    label += 1
    label = tf.expand_dims(label, axis=3)

    # Resizes and crops image.
    image, image_info = preprocess_ops.resize_and_crop_image(
        image,
        self._output_size,
        self._output_size,
        centered_crop=self._centered_crop,
    )

    if self._resize_eval_groundtruth:
      # Resizes eval masks to match input image sizes. In that case, mean IoU
      # is computed on output_size not the original size of the images.
      image_scale = image_info[2, :]
      offset = image_info[3, :]
      label = preprocess_ops.resize_and_crop_masks(
          label,
          image_scale,
          self._output_size,
          offset,
          centered_crop=self._centered_crop,
      )
    else:
      if self._centered_crop:
        label_size = tf.cast(tf.shape(label)[0:2], tf.int32)
        label = tf.image.pad_to_bounding_box(
            label,
            tf.maximum(
                (self._groundtruth_padded_size[0] - label_size[0]) // 2, 0
            ),
            tf.maximum(
                (self._groundtruth_padded_size[1] - label_size[1]) // 2, 0
            ),
            self._groundtruth_padded_size[0],
            self._groundtruth_padded_size[1],
        )
      else:
        label = tf.image.pad_to_bounding_box(
            label,
            0,
            0,
            self._groundtruth_padded_size[0],
            self._groundtruth_padded_size[1],
        )

    label -= 1
    label = tf.where(
        tf.equal(label, -1), self._ignore_label * tf.ones_like(label), label
    )
    label = tf.squeeze(label, axis=0)

    valid_mask = tf.not_equal(label, self._ignore_label)
    labels = {
        'masks': label,
        'valid_masks': valid_mask,
        'image_info': image_info,
    }

    # Cast image as self._dtype
    image = tf.cast(image, dtype=self._dtype)

    return image, labels