Hooks to display videos at decode time.

def display_video_hooks(hook_args):
  """Hooks to display videos at decode time."""
  predictions = hook_args.predictions
  max_outputs = hook_args.decode_hparams.max_display_outputs
  max_decodes = hook_args.decode_hparams.max_display_decodes

  with tf.Graph().as_default():
    _, best_decodes = video_metrics.compute_video_metrics_from_predictions(
        predictions, decode_hparams=hook_args.decode_hparams)

  all_summaries = []
  # Displays decodes corresponding to the best/worst metric,
  for metric, metric_decode_inds in best_decodes.items():
    curr_metric_inds = metric_decode_inds[:max_outputs]
    best_inputs, best_outputs, best_targets = [], [], []
    for sample_ind, decode_ind in enumerate(curr_metric_inds):
      curr_decode = predictions[decode_ind][sample_ind]
      best_inputs.append(curr_decode["inputs"])
      best_outputs.append(curr_decode["outputs"])
      best_targets.append(curr_decode["targets"])
    best_inputs = np.array(best_inputs, dtype=np.uint8)
    best_outputs = np.array(best_outputs, dtype=np.uint8)
    best_targets = np.array(best_targets, dtype=np.uint8)
    summaries = convert_videos_to_summaries(
        best_inputs, best_outputs, best_targets,
        tag=metric, decode_hparams=hook_args.decode_hparams)
    all_summaries.extend(summaries)

  # Display random decodes for ten conditioning frames.
  for decode_ind, decode in enumerate(predictions[: max_decodes]):
    target_videos = video_metrics.stack_data_given_key(decode, "targets")
    output_videos = video_metrics.stack_data_given_key(decode, "outputs")
    input_videos = video_metrics.stack_data_given_key(decode, "inputs")
    target_videos = np.asarray(target_videos, dtype=np.uint8)
    output_videos = np.asarray(output_videos, dtype=np.uint8)
    input_videos = np.asarray(input_videos, dtype=np.uint8)
    summaries = convert_videos_to_summaries(
        input_videos, output_videos, target_videos,
        tag="decode_%d" % decode_ind, decode_hparams=hook_args.decode_hparams,
        display_ground_truth=decode_ind == 0)
    all_summaries.extend(summaries)
  return all_summaries