data_pipeline.py 文件源码

def _augment_data(self, inout, nchan=6):
    """Flip, crop and rotate samples randomly."""

    with tf.name_scope('data_augmentation'):
      if self.fliplr:
        inout = tf.image.random_flip_left_right(inout, seed=1234)
      if self.flipud:
        inout = tf.image.random_flip_up_down(inout, seed=3456)
      if self.rotate:
        angle = tf.random_uniform((), minval=0, maxval=4, dtype=tf.int32, seed=4567)
        inout = tf.case([(tf.equal(angle, 1), lambda: tf.image.rot90(inout, k=1)),
                         (tf.equal(angle, 2), lambda: tf.image.rot90(inout, k=2)),
                         (tf.equal(angle, 3), lambda: tf.image.rot90(inout, k=3))],
                        lambda: inout)

      inout.set_shape([None, None, nchan])

      with tf.name_scope('crop'):
        shape = tf.shape(inout)
        new_height = tf.to_int32(self.output_resolution[0])
        new_width = tf.to_int32(self.output_resolution[1])
        height_ok = tf.assert_less_equal(new_height, shape[0])
        width_ok = tf.assert_less_equal(new_width, shape[1])
        with tf.control_dependencies([height_ok, width_ok]):
          if self.random_crop:
            inout = tf.random_crop(
                inout, tf.stack([new_height, new_width, nchan]))
          else:
            height_offset = tf.to_int32((shape[0]-new_height)/2)
            width_offset = tf.to_int32((shape[1]-new_width)/2)
            inout = tf.image.crop_to_bounding_box(
                inout, height_offset, width_offset,
                new_height, new_width)

      inout.set_shape([None, None, nchan])
      inout = tf.image.resize_images(
          inout, [self.output_resolution[0], self.output_resolution[1]])
      fullres = inout

      with tf.name_scope('resize'):
        new_size = 256
        inout = tf.image.resize_images(
            inout, [new_size, new_size],
            method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)

      return fullres, inout