repo_description: Learning infinite-resolution image processing with GAN and RL from unpaired image datasets, using a differentiable photo editing model.
repo_topics: computational-photography deep-learning gan generative-adversarial-network image-processing reinforcement-learning
code: import tensorflow as tf
  import tensorflow.contrib.layers as ly
  from util import lrelu
  import cv2
  import math
  from pdf_sample_layer import pdf_sample
  from util import enrich_image_input
  from util import STATE_DROPOUT_BEGIN, STATE_REWARD_DIM, STATE_STEP_DIM, STATE_STOPPED_DIM
  
  
  def feature_extractor(net, output_dim, cfg):
    net = net - 0.5
    min_feature_map_size = 4
    assert output_dim % (
        min_feature_map_size**2) == 0, 'output dim=%d' % output_dim
    size = int(net.get_shape()[2])
    print('Agent CNN:')
    channels = cfg.base_channels
    print('    ', str(net.get_shape()))
    size /= 2
    net = ly.conv2d(
        net, num_outputs=channels, kernel_size=4, stride=2, activation_fn=lrelu)
    print('    ', str(net.get_shape()))
    while size > min_feature_map_size:
      if size == min_feature_map_size * 2:
        channels = output_dim / (min_feature_map_size**2)
      else:
        channels *= 2
      assert size % 2 == 0
      size /= 2
      net = ly.conv2d(
          net, num_outputs=channels, kernel_size=4, stride=2, activation_fn=lrelu)
      print('    ', str(net.get_shape()))
    print('before fc: ', net.get_shape()[1])
    net = tf.reshape(net, [-1, output_dim])
    net = tf.nn.dropout(net, cfg.dropout_keep_prob)
    return net
  
  
  # Output: float \in [0, 1]
  def agent_generator(inp, is_train, progress, cfg, high_res=None, alex_in=None):
    net, z, states = inp
    filters = cfg.filters
  
    filters = [x(net, cfg) for x in filters]
  
    selection_noise = z[:, 0:1]
    filtered_images = []
    filter_debug_info = []
    high_res_outputs = []
  
    if cfg.shared_feature_extractor:
      filter_features = feature_extractor(
          net=enrich_image_input(cfg, net, states),
          output_dim=cfg.feature_extractor_dims,
          cfg=cfg)
      # filter_features = ly.dropout(filter_features)
    for j, filter in enumerate(filters):
      with tf.variable_scope('filter_%d' % j):
        print('    creating filter:', j, 'name:', str(filter.__class__), 'abbr.',
              filter.get_short_name())
        if not cfg.shared_feature_extractor:
          filter_features = \
              feature_extractor(net=enrich_image_input(cfg, net),
                                output_dim=cfg.feature_extractor_dims, cfg=cfg)
        print('      filter_features:', filter_features.shape)
        filtered_image_batch, high_res_output, per_filter_debug_info = filter.apply(
            net, filter_features, high_res=high_res)
        high_res_outputs.append(high_res_output)
        filtered_images.append(filtered_image_batch)
        filter_debug_info.append(per_filter_debug_info)
        print('      output:', filtered_image_batch.shape)
  
    # [batch_size, #filters, H, W, C]
    for img in filtered_images:
      print('img', img.shape)
    filtered_images = tf.stack(values=filtered_images, axis=1)
    print('    filtered_images:', filtered_images.shape)
  
    with tf.variable_scope('action_selection'):
      selector_features = feature_extractor(
          net=enrich_image_input(cfg, net, states),
          output_dim=cfg.feature_extractor_dims,
          cfg=cfg)
      print('    selector features:', selector_features.shape)
  
      selector_features = ly.fully_connected(
          selector_features,
          num_outputs=cfg.fc1_size,
          scope='selector_fc1',
          activation_fn=lrelu)
  
      # selector_features = ly.dropout(selector_features)
  
      pdf = ly.fully_connected(
          selector_features,
          num_outputs=len(filters),
          activation_fn=None,
          scope='selector_fc2')
      pdf = tf.nn.softmax(pdf) + 1e-37
      print('    pdf_filter', pdf[:, 1:].shape)
      # print('    pdf_mask', states[:, STATE_DROPOUT_BEGIN:].shape)
  
      pdf = pdf * (1 - cfg.exploration) + cfg.exploration * 1.0 / len(filters)
      # pdf = tf.to_float(is_train) * tf.concat([pdf[:, :1], pdf[:, 1:] * states[:, STATE_DROPOUT_BEGIN:]], axis=1) \
      # + (1.0 - tf.to_float(is_train)) * pdf
      pdf = pdf / (tf.reduce_sum(pdf, axis=1, keep_dims=True) + 1e-30)
      entropy = -pdf * tf.log(pdf)
      entropy = tf.reduce_sum(entropy, axis=1)[:, None]
      print('    pdf:', pdf.shape)
      print('    entropy:', entropy.shape)
      print('    selection_noise:', selection_noise.shape)
      random_filter_id = pdf_sample(pdf, selection_noise)
      max_filter_id = tf.cast(tf.argmax(pdf, axis=1), tf.int32)
      selected_filter_id = is_train * random_filter_id + (
          1 - is_train) * max_filter_id
      print('    selected_filter_id:', selected_filter_id.shape)
      filter_one_hot = tf.one_hot(
          selected_filter_id, depth=len(filters), dtype=tf.float32)
      print('    filter one_hot', filter_one_hot.shape)
      surrogate = tf.reduce_sum(
          filter_one_hot * tf.log(pdf + 1e-10), axis=1, keep_dims=True)
  
    net = tf.reduce_sum(
        filtered_images * filter_one_hot[:, :, None, None, None], axis=1)
    if high_res is not None:
      high_res_outputs = tf.stack(values=high_res_outputs, axis=1)
      high_res_output = tf.reduce_sum(
          high_res_outputs * filter_one_hot[:, :, None, None, None], axis=1)
  
    # only the first image will get debug_info
    debug_info = {
        'state': states,
        'selected_filter_id': selected_filter_id[0],
        'filter_debug_info': filter_debug_info,
        'pdf': pdf[0]
    }
  
    # Combined: Three in one 64x64 ?
    #           otherwise returns pdf, detail, mask
    def debugger(debug_info, combined=True):
      size = 8
      img = None
      images = [None for i in range(3)]
      for i, filter in enumerate(filters):
        selected = i == debug_info['selected_filter_id']
        if selected:
          img = filter.visualize_mask(debug_info['filter_debug_info'][i],
                                      (64, 64)) * 0.8
      assert img is not None
      if not combined:
        # Mask
        images[2] = img.copy()
        # reset img
        img = img * 0 + 0.5
  
      c = 0
      for i, filter in enumerate(filters):
        pdf = debug_info['pdf'][i]
        if pdf < 1e-10:
          continue
        else:
          c += 1
        selected = i == debug_info['selected_filter_id']
        if selected:
          filter.visualize_filter(debug_info['filter_debug_info'][i], img)
      if not combined:
        # detail
        images[1] = img.copy()
        # reset img
        img = img * 0 + 0.5
      c = 0
      for i, filter in enumerate(filters):
        per_col = 4
        x = c // per_col * 30
        y = size * (c % per_col + 1)
        pdf = debug_info['pdf'][i]
        if pdf < 1e-10:
          continue
        else:
          c += 1
        cv2.putText(img,
                    filter.get_short_name(), (x + 6, y + 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.233, (255, 255, 255))
        selected = i == debug_info['selected_filter_id']
        color = 1.0 if selected else 0.3
        width = int(pdf * 20)
        height = 0.35
        corners = [(x + 16, int(y + (1 - height) * size // 2)),
                   (x + 16 + width, int(y + (1 + height) * size // 2))]
        cv2.rectangle(img, (corners[0][0] - 1, corners[0][1] - 1),
                      (corners[1][0] + 1, corners[1][1] + 1), (1, 1,
                                                               1), cv2.FILLED)
        cv2.rectangle(img, corners[0], corners[1], (color, 0.3, 0.3), cv2.FILLED)
      if not combined:
        # pdf
        images[0] = img.copy()
  
      if combined:
        return img
      else:
        return images
  
    debugger.width = int(net.shape[1])
    print('    surrogate: ', surrogate.shape)
  
    # Calculate new states
    new_states = [None for _ in range(STATE_DROPOUT_BEGIN + 1)]
    is_last_step = tf.cast(
        tf.abs(states[:, STATE_STEP_DIM:STATE_STEP_DIM + 1] + 1 - cfg.test_steps)
        < 1e-4,
        dtype=tf.float32)
    submitted = is_last_step
  
    new_states[STATE_REWARD_DIM] = submitted
    new_states[STATE_STOPPED_DIM] = submitted
    # Increment the step
    new_states[STATE_STEP_DIM] = (states[:, STATE_STEP_DIM] + 1)[:, None]
  
    # Update filter usage
    filter_usage = states[:, STATE_STEP_DIM + 1:]
    print('usage v.s. onehot', filter_usage.shape, filter_one_hot.shape)
    assert len(filter_usage.shape) == len(filter_one_hot.shape)
  
    regular_filter_start = 0
  
    # Penalize submission action that is not the final action.
    early_stop_penalty = (1 - is_last_step) * submitted * cfg.early_stop_penalty
  
    usage_penalty = tf.reduce_sum(
        filter_usage * filter_one_hot[:, regular_filter_start:],
        axis=1,
        keep_dims=True)
    new_filter_usage = tf.maximum(filter_usage,
                                  filter_one_hot[:, regular_filter_start:])
    new_states[STATE_STEP_DIM + 1] = new_filter_usage
  
    print(submitted.shape, new_states[STATE_STEP_DIM].shape)
    new_states = tf.concat(new_states, axis=1)
    print('new_states:', new_states.shape)
  
    if cfg.clamp:
      net = tf.clip_by_value(net, 0.0, 5.0)
  
    entropy_penalty = (1.0 - progress) * cfg.exploration_penalty * (
        -entropy + math.log(len(filters)))
  
    # Will be substracted from award
    penalty = tf.reduce_mean(
        tf.maximum(net - 1, 0)**2, axis=(1, 2, 3)
    )[:,
      None] + entropy_penalty + usage_penalty * cfg.filter_usage_penalty + early_stop_penalty
  
    print('states, new_states:', states.shape, new_states.shape)
    print('penalty:', penalty.shape)
  
    if high_res is None:
      return (net, new_states, surrogate, penalty), debug_info, debugger
    else:
      return (net, new_states, high_res_output), debug_info, debugger
  