def get_dataset(
    config: ml_collections.ConfigDict,
    data_rng: PRNGKey,
    is_train: bool = True,
    *,
    dataset_service_address: Optional[str] = None,
) -> scenic_dataset_utils.Dataset:
  """Get dataset from scenic and transform image datasets."""
  if config.dataset_name == 'unused':
    return scenic_dataset_utils.Dataset(None, None, None, None)
  batch_size = config.batch_size
  dataset_name = config.dataset_name
  if not is_train:
    batch_size = config.get('eval_batch_size', batch_size)
    if config.get('loader_batch_size') is not None:
      batch_size = config.loader_batch_size
    if config.get('eval_dataset_name') is not None:
      # uses a different data_loader for evaluation.
      dataset_name = config.eval_dataset_name
      with config.unlocked():
        # overrides train logic. It's fine as the train_iter is not used here.
        config.dataset_configs.pp_train = config.dataset_configs.pp_eval

  dataset = _get_scenic_dataset(
      config,
      data_rng,
      batch_size,
      dataset_name=dataset_name,
      dataset_service_address=dataset_service_address)
  if dataset.valid_iter and config.dataset_configs.get('num_eval_clips', 1) > 1:
    dataset = dataset._replace(
        valid_iter=_split_batch_iter(dataset.valid_iter, batch_size))
  if dataset.test_iter and (
      config.dataset_configs.get('num_test_clips', 1) > 1 or
      config.dataset_configs.get('do_three_spatial_crops', False)):
    dataset = dataset._replace(
        test_iter=_split_batch_iter(dataset.test_iter, batch_size))
  dataset_type = config.get('dataset_type', 'video')
  if dataset_type == 'video':
    return dataset
  elif dataset_type == 'image':
    s = dataset.meta_data['input_shape']
    dataset.meta_data['input_shape'] = (s[0], 1, *s[1:])
    return dataset._replace(
        train_iter=map(_inputs_add_t_dim, dataset.train_iter)
        if dataset.train_iter else None,
        valid_iter=map(_inputs_add_t_dim, dataset.valid_iter)
        if dataset.valid_iter else None,
        test_iter=map(_inputs_add_t_dim, dataset.test_iter)
        if dataset.test_iter else None)
  else:
    raise NotImplementedError(
        f'Unrecognized dataset type: {dataset_type}')
Get dataset from scenic and transform image datasets.
videogvt/train_lib/train_utils.py