def prune_df(df, max_num_images_per_source, summary_file_pruned=None):
  """Keep first N images from each source to create a pruned summary file.
  
    The motivation is to trim down samples from sources with too many images,
    to create more balanced datasets.

  Args:
    df: summary dataframe created by summarize_compression_features()
    max_num_images_per_source: int, keep this number of images (at max) from
      each source
    summary_file_pruned: if not None, save a copy of the pruned dataframe
      at this file path
  
  Returns:
    df_pruned: slimmer (pruned) dataset
  """
  df_pruned = df.groupby('source').head(max_num_images_per_source).reset_index(drop=True)
  if summary_file_pruned is not None:
    df_pruned.to_pickle(summary_file_pruned)  
  return df_pruned