def write_batch_samples(batch: Batch, step: int,
                        executor: concurrent.futures.Executor,
                        result_dir: str, prefix: str,
                        label_names: Optional[Dict[int, str]],
                        config: ml_collections.ConfigDict):
  """Write a batch of samples to disk."""
  enable_condition = config.eval.get('results_with_condition', True)
  enable_input = config.eval.get('results_with_input', False)
  condition_alpha = config.eval.get('results_condition_alpha',
                                    0 if 'frame' not in prefix else 1)
  batch = apply_batch_mask(batch)
  if 'generated_video_ema' in batch:  # For VQ models
    videos = batch['generated_video_ema']
    suffix = '_ema'
    if enable_condition:
      videos = np.concatenate((videos, batch['original_video']), axis=2)
  else:  # For GVT models
    videos = batch['generated_video']
    suffix = ''
    if enable_input:
      videos = np.concatenate((videos, batch['original_video']), axis=2)
    if 'condition_video' in batch and enable_condition:
      cond_videos = batch['condition_video']
      if 'frame_interpolation' in prefix:
        # TODO(Lijun-Yu): refactor this visualization hack.
        l = cond_videos.shape[1] // 2
        cond_videos[:, 1:l] = cond_videos[:, :1]
        cond_videos[:, l:-1] = cond_videos[:, -1:]
      cond_videos = np.where(batch['condition_mask'], batch['condition_video'],
                             condition_alpha * cond_videos)
      videos = np.concatenate((videos, cond_videos), axis=2)
  label_ids = batch.get('label')
  if label_ids is not None and label_ids.shape[1] > 1:
    # Take the first one hot label
    label_ids = label_ids.argmax(axis=1, keepdims=True).astype(np.int32)
  idx_offset = step * config.eval_batch_size + jax.process_index(
  ) * config.eval_batch_size // jax.process_count()
  if 'frame_rate' in config.dataset_configs:
    fps = config.dataset_configs.frame_rate / config.dataset_configs.stride
  else:
    fps = 8
  for idx, video in enumerate(videos):
    label_str = ''
    if label_names is not None and label_ids is not None:
      label_id = int(label_ids[idx, 0])
      label_str = f'{label_id:04d}{label_names[label_id]}_'
    if config.eval.get('enable_ssim_psnr', False):
      label_str += f'ifssim={batch[f"interframe_ssim{suffix}"][idx]:.02f}_'
      # label_str += f'ssim={batch[f"ssim{suffix}"][idx]:.02f}_'
      # label_str += f'psnr={batch[f"psnr{suffix}"][idx]:02.0f}_'
    if config.eval.get('enable_lpips', False):
      # label_str += f'lpalex={batch[f"lpips_alexnet"][idx]:.02f}_'
      label_str += f'lpvgg={batch["lpips_vgg"][idx]:.02f}_'
    filename = f'{prefix}_{label_str}{(idx + idx_offset):09d}.mp4'
    path = os.path.join(result_dir, filename)
    executor.submit(mediapy.write_video, path, video, fps=fps)
Write a batch of samples to disk.
videogvt/train_lib/eval_utils.py