models/research/domain_adaptation/pixel_domain_adaptation/pixelda_utils.py

# Copyright 2017 Google Inc.
#
# 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.

"""Utilities for PixelDA model."""
import math

# Dependency imports

import tensorflow as tf

slim = tf.contrib.slim

flags = tf.app.flags
FLAGS = flags.FLAGS


def remove_depth(images):
  """Takes a batch of images and remove depth channel if present."""
  if images.shape.as_list()[-1] == 4:
    return images[:, :, :, 0:3]
  return images


def image_grid(images, max_grid_size=4):
  """Given images and N, return first N^2 images as an NxN image grid.

  Args:
    images: a `Tensor` of size [batch_size, height, width, channels]
    max_grid_size: Maximum image grid height/width

  Returns:
    Single image batch, of dim [1, h*n, w*n, c]
  """
  images = remove_depth(images)
  batch_size = images.shape.as_list()[0]
  grid_size = min(int(math.sqrt(batch_size)), max_grid_size)
  assert images.shape.as_list()[0] >= grid_size * grid_size

  # If we have a depth channel
  if images.shape.as_list()[-1] == 4:
    images = images[:grid_size * grid_size, :, :, 0:3]
    depth = tf.image.grayscale_to_rgb(images[:grid_size * grid_size, :, :, 3:4])

    images = tf.reshape(images, [-1, images.shape.as_list()[2], 3])
    split = tf.split(0, grid_size, images)
    depth = tf.reshape(depth, [-1, images.shape.as_list()[2], 3])
    depth_split = tf.split(0, grid_size, depth)
    grid = tf.concat(split + depth_split, 1)
    return tf.expand_dims(grid, 0)
  else:
    images = images[:grid_size * grid_size, :, :, :]
    images = tf.reshape(
        images, [-1, images.shape.as_list()[2],
                 images.shape.as_list()[3]])
    split = tf.split(images, grid_size, 0)
    grid = tf.concat(split, 1)
    return tf.expand_dims(grid, 0)


def source_and_output_image_grid(output_images,
                                 source_images=None,
                                 max_grid_size=4):
  """Create NxN image grid for output, concatenate source grid if given.

  Makes grid out of output_images and, if provided, source_images, and
  concatenates them.

  Args:
    output_images: [batch_size, h, w, c] tensor of images
    source_images: optional[batch_size, h, w, c] tensor of images
    max_grid_size: Image grid height/width

  Returns:
    Single image batch, of dim [1, h*n, w*n, c]


  """
  output_grid = image_grid(output_images, max_grid_size=max_grid_size)
  if source_images is not None:
    source_grid = image_grid(source_images, max_grid_size=max_grid_size)
    # Make sure they have the same # of channels before concat
    # Assumes either 1 or 3 channels
    if output_grid.shape.as_list()[-1] != source_grid.shape.as_list()[-1]:
      if output_grid.shape.as_list()[-1] == 1:
        output_grid = tf.tile(output_grid, [1, 1, 1, 3])
      if source_grid.shape.as_list()[-1] == 1:
        source_grid = tf.tile(source_grid, [1, 1, 1, 3])
    output_grid = tf.concat([output_grid, source_grid], 1)
  return output_grid


def summarize_model(end_points):
  """Summarizes the given model via its end_points.

  Args:
    end_points: A dictionary of end_point names to `Tensor`.
  """
  tf.summary.histogram('domain_logits_transferred',
                       tf.sigmoid(end_points['transferred_domain_logits']))

  tf.summary.histogram('domain_logits_target',
                       tf.sigmoid(end_points['target_domain_logits']))


def summarize_transferred_grid(transferred_images,
                               source_images=None,
                               name='Transferred'):
  """Produces a visual grid summarization of the image transferrence.

  Args:
    transferred_images: A `Tensor` of size [batch_size, height, width, c].
    source_images: A `Tensor` of size [batch_size, height, width, c].
    name: Name to use in summary name
  """
  if source_images is not None:
    grid = source_and_output_image_grid(transferred_images, source_images)
  else:
    grid = image_grid(transferred_images)
  tf.summary.image('%s_Images_Grid' % name, grid, max_outputs=1)


def summarize_transferred(source_images,
                          transferred_images,
                          max_images=20,
                          name='Transferred'):
  """Produces a visual summary of the image transferrence.

  This summary displays the source image, transferred image, and a grayscale
  difference image which highlights the differences between input and output.

  Args:
    source_images: A `Tensor` of size [batch_size, height, width, channels].
    transferred_images: A `Tensor` of size [batch_size, height, width, channels]
    max_images: The number of images to show.
    name: Name to use in summary name

  Raises:
    ValueError: If number of channels in source and target are incompatible
  """
  source_channels = source_images.shape.as_list()[-1]
  transferred_channels = transferred_images.shape.as_list()[-1]
  if source_channels < transferred_channels:
    if source_channels != 1:
      raise ValueError(
          'Source must be 1 channel or same # of channels as target')
    source_images = tf.tile(source_images, [1, 1, 1, transferred_channels])
  if transferred_channels < source_channels:
    if transferred_channels != 1:
      raise ValueError(
          'Target must be 1 channel or same # of channels as source')
    transferred_images = tf.tile(transferred_images, [1, 1, 1, source_channels])
  diffs = tf.abs(source_images - transferred_images)
  diffs = tf.reduce_max(diffs, reduction_indices=[3], keep_dims=True)
  diffs = tf.tile(diffs, [1, 1, 1, max(source_channels, transferred_channels)])

  transition_images = tf.concat([
      source_images,
      transferred_images,
      diffs,
  ], 2)

  tf.summary.image(
      '%s_difference' % name, transition_images, max_outputs=max_images)


def summaries_color_distributions(images, name):
  """Produces a histogram of the color distributions of the images.

  Args:
    images: A `Tensor` of size [batch_size, height, width, 3].
    name: The name of the images being summarized.
  """
  tf.summary.histogram('color_values/%s' % name, images)


def summarize_images(images, name):
  """Produces a visual summary of the given images.

  Args:
    images: A `Tensor` of size [batch_size, height, width, 3].
    name: The name of the images being summarized.
  """
  grid = image_grid(images)
  tf.summary.image('%s_Images' % name, grid, max_outputs=1)