<commit_before>import io
from absl import logging
from functools import partial
import tensorflow as tf
import tensorflow_datasets as tfds


def scale_down(method="bicubic", dimension=1024, factor=4):
  
  def scale_fn(image, *args, **kwargs):
    high_resolution = image
    if not kwargs.get("no_random_crop", None):
      high_resolution = tf.image.random_crop(
          image, [dimension, dimension, image.shape[-1]])

    low_resolution = tf.image.resize(
        high_resolution,
        [dimension // factor, dimension // factor],
        method=method)
    low_resolution = tf.clip_by_value(low_resolution, 0, 255)
    high_resolution = tf.clip_by_value(high_resolution, 0, 255)
    return low_resolution, high_resolution
  scale_fn.dimension = dimension
  return scale_fn


def augment_image(
        low_res_map_fn,
        brightness_delta=0.05,
        contrast_factor=[0.7, 1.3],
        saturation=[0.6, 1.6]):
  
  def augment_fn(_, high_resolution, *args, **kwargs):
    high_resolution = tf.image.random_flip_left_right(high_resolution)
    high_resolution = tf.image.random_flip_up_down(high_resolution)
    high_resolution = tf.image.random_brightness(
        high_resolution, max_delta=brightness_delta)
    if contrast_factor:
      high_resolution = tf.image.random_contrast(
          high_resolution, *contrast_factor)
    if saturation:
      high_resolution = tf.image.random_saturation(
          high_resolution, *saturation)
    low_res, high_res = low_res_map_fn(high_resolution)
    return low_res, high_res
  return augment_fn


def reform_dataset(dataset, types):
  """ Helper function to convert the output_dtype of the dataset
      from (tf.float32, tf.uint8) to desired dtype
  """
  def generator_fn():
    for data in dataset:
      yield data[0], data[1]
  return tf.data.Dataset.from_generator(
      generator_fn, types, (tf.TensorShape([None, None, 3]), tf.TensorShape(None)))


def load_dataset_directory(
        name,
        directory,
        low_res_map_fn,
        batch_size=32,
        iterations=1,
        shuffle=True,
        augment=True,
        cache_dir="cache/",
        buffer_size=io.DEFAULT_BUFFER_SIZE):

  if not tf.io.gfile.exists(cache_dir):
    tf.io.gfile.mkdir(cache_dir)
  dl_config = tfds.download.DownloadConfig(manual_dir=directory)

  dataset = reform_dataset(
      tfds.load(
          "image_label_folder/dataset_name=%s" %
          name,
          split="train",
          as_supervised=True,
          download_and_prepare_kwargs={
              "download_config": dl_config}),
      (tf.float32, tf.float32))

  dataset = (dataset.filter(
      lambda image, *args: tf.greater_equal(
          image.shape[:-1],
          2 * [low_res_map_fn.dimension]).numpy().all())
      .map(low_res_map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      .batch(batch_size)
      .prefetch(buffer_size)
      .cache(cache_dir))

  if shuffle:
    dataset = dataset.shuffle(buffer_size, reshuffle_each_iteration=True)

  if augment:
    dataset = dataset.map(
        augment_image(
            partial(low_res_map_fn, no_random_crop=True),
            saturation=None),
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return dataset


def load_dataset(
        name,
        low_res_map_fn,
        split="train",
        batch_size=32,
        iterations=1,
        shuffle=True,
        augment=True,
        buffer_size=io.DEFAULT_BUFFER_SIZE,
        cache_dir="cache/",
        data_dir=None):

  if not tf.io.gfile.exists(cache_dir):
    tf.io.gfile.mkdir(cache_dir)
  dataset = reform_dataset(
      tfds.load(
          name,
          data_dir=data_dir,
          split=split,
          as_supervised=True),
      (tf.float32, tf.float32))
  
  dataset = (dataset.filter(
      lambda image, *args: tf.greater_equal(
          image.shape[:-1],
          2 * [low_res_map_fn.dimension]).numpy().all())
      .map(low_res_map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      .batch(batch_size)
      .prefetch(buffer_size)
      .cache(cache_dir))
  
  if shuffle:
    dataset = dataset.shuffle(buffer_size, reshuffle_each_iteration=True)
  
  if augment:<commit_msg>Added random 90 degree rotation of the images<commit_after>import io
from absl import logging
from functools import partial
import tensorflow as tf
import tensorflow_datasets as tfds


def scale_down(method="bicubic", dimension=256, factor=4):
  def scale_fn(image, *args, **kwargs):
    high_resolution = image
    if not kwargs.get("no_random_crop", None):
      high_resolution = tf.image.random_crop(
          image, [dimension, dimension, image.shape[-1]])

    low_resolution = tf.image.resize(
        high_resolution,
        [dimension // factor, dimension // factor],
        method=method)
    low_resolution = tf.clip_by_value(low_resolution, 0, 255)
    high_resolution = tf.clip_by_value(high_resolution, 0, 255)
    return low_resolution, high_resolution
  scale_fn.dimension = dimension
  return scale_fn


def augment_image(
        low_res_map_fn,
        brightness_delta=0.05,
        contrast_factor=[0.7, 1.3],
        saturation=[0.6, 1.6]):
  """ helper function used for augmentation of images in the dataset. """
  def augment_fn(_, high_resolution, *args, **kwargs):
    high_resolution = tf.image.rot90(
        high_resolution, tf.random.uniform(
            minval=1, maxval=5, dtype=tf.int32, shape=[]))
    high_resolution = tf.image.random_flip_left_right(high_resolution)
    high_resolution = tf.image.random_flip_up_down(high_resolution)
    high_resolution = tf.image.random_brightness(
        high_resolution, max_delta=brightness_delta)
    if contrast_factor:
      high_resolution = tf.image.random_contrast(
          high_resolution, *contrast_factor)
    if saturation:
      high_resolution = tf.image.random_saturation(
          high_resolution, *saturation)
    low_res, high_res = low_res_map_fn(high_resolution)
    return low_res, high_res
  return augment_fn


def reform_dataset(dataset, types):
  """ Helper function to convert the output_dtype of the dataset
      from (tf.float32, tf.uint8) to desired dtype
  """
  def generator_fn():
    for data in dataset:
      yield data[0], data[1]
  return tf.data.Dataset.from_generator(
      generator_fn, types, (tf.TensorShape([None, None, 3]), tf.TensorShape(None)))


def load_dataset_directory(
        name,
        directory,
        low_res_map_fn,
        batch_size=32,
        iterations=1,
        shuffle=True,
        augment=True,
        cache_dir="cache/",
        buffer_size=io.DEFAULT_BUFFER_SIZE):

  if not tf.io.gfile.exists(cache_dir):
    tf.io.gfile.mkdir(cache_dir)
  dl_config = tfds.download.DownloadConfig(manual_dir=directory)

  dataset = reform_dataset(
      tfds.load(
          "image_label_folder/dataset_name=%s" %
          name,
          split="train",
          as_supervised=True,
          download_and_prepare_kwargs={
              "download_config": dl_config}),
      (tf.float32, tf.float32))

  dataset = (dataset.filter(
      lambda image, *args: tf.greater_equal(
          image.shape[:-1],
          2 * [low_res_map_fn.dimension]).numpy().all())
      .map(low_res_map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      .batch(batch_size)
      .prefetch(buffer_size)
      .cache(cache_dir))

  if shuffle:
    dataset = dataset.shuffle(buffer_size, reshuffle_each_iteration=True)

  if augment:
    dataset = dataset.map(
        augment_image(
            partial(low_res_map_fn, no_random_crop=True),
            saturation=None),
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return dataset


def load_dataset(
        name,
        low_res_map_fn,
        split="train",
        batch_size=32,
        iterations=1,
        shuffle=True,
        augment=True,
        buffer_size=io.DEFAULT_BUFFER_SIZE,
        cache_dir="cache/",
        data_dir=None):

  if not tf.io.gfile.exists(cache_dir):
    tf.io.gfile.mkdir(cache_dir)
  dataset = reform_dataset(
      tfds.load(
          name,
          data_dir=data_dir,
          split=split,
          as_supervised=True),
      (tf.float32, tf.float32))

  dataset = (dataset.filter(
      lambda image, *args: tf.greater_equal(
          image.shape[:-1],
          2 * [low_res_map_fn.dimension]).numpy().all())
      .map(low_res_map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      .batch(batch_size)
      .prefetch(buffer_size)
      .cache(cache_dir))

  if shuffle:
    dataset = dataset.shuffle(buffer_size, reshuffle_each_iteration=True)

  if augment:
{'fraction_non_alphanumeric': 0.10344426343960902, 'fraction_numerical': 0.012217826390505003, 'mean_word_length': 2.9517241379310346, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1273113', 'n_tokens_mistral': 2916, 'n_tokens_neox': 2758, 'n_words': 508}