tensorflow_probability/examples/sprites_dataset.py
def create_sprites_dataset(characters, actions, directions, channels=3,
                           length=8, shuffle=False, fake_data=False):
  """Creates a tf.data pipeline for the sprites dataset.

  Args:
    characters: A list of (skin, hair, top, pants) tuples containing
      relative paths to the sprite png image for each attribute.
    actions: A list of Actions.
    directions: A list of Directions.
    channels: Number of image channels to yield.
    length: Desired length of the sequences.
    shuffle: Whether or not to shuffle the characters and sequences
      start frame.
    fake_data: Boolean for whether or not to yield synthetic data.

  Returns:
    A tf.data.Dataset yielding (seq, skin label index, hair label index,
    top label index, pants label index, action label index, skin label
    name, hair label_name, top label name, pants label name, action
    label name) tuples.
  """
  if fake_data:
    dummy_image = tf.random.normal([HEIGHT, WIDTH, CHANNELS])
  else:
    basedir = download_sprites()

  action_names = [action.name for action in actions]
  action_metadata = [(action.start_row, action.frames) for action in actions]

  direction_rows = [direction.row_offset for direction in directions]

  chars = tf.data.Dataset.from_tensor_slices(characters)
  act_names = tf.data.Dataset.from_tensor_slices(action_names).repeat()
  acts_metadata = tf.data.Dataset.from_tensor_slices(action_metadata).repeat()
  dir_rows = tf.data.Dataset.from_tensor_slices(direction_rows).repeat()

  if shuffle:
    chars = chars.shuffle(len(characters))

  dataset = tf.data.Dataset.zip((chars, act_names, acts_metadata, dir_rows))

  skin_table = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))
  hair_table = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))
  top_table = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))
  pants_table = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))
  action_table = tf.contrib.lookup.index_table_from_tensor(sorted(action_names))

  def process_example(attrs, act_name, act_metadata, dir_row_offset):
    """Processes a dataset row."""
    skin_name = attrs[0]
    hair_name = attrs[1]
    top_name = attrs[2]
    pants_name = attrs[3]

    if fake_data:
      char = dummy_image
    else:
      skin = read_image(basedir + os.sep + skin_name)
      hair = read_image(basedir + os.sep + hair_name)
      top = read_image(basedir + os.sep + top_name)
      pants = read_image(basedir + os.sep + pants_name)
      char = create_character(skin, hair, top, pants)

    if shuffle:
      seq = create_random_seq(char, act_metadata, dir_row_offset, length)
    else:
      seq = create_seq(char, act_metadata, dir_row_offset, length)
    seq = seq[..., :channels]  # limit output channels

    skin_idx = skin_table.lookup(skin_name)
    hair_idx = hair_table.lookup(hair_name)
    top_idx = top_table.lookup(top_name)
    pants_idx = pants_table.lookup(pants_name)
    act_idx = action_table.lookup(act_name)

    return (seq, skin_idx, hair_idx, top_idx, pants_idx, act_idx,
            skin_name, hair_name, top_name, pants_name, act_name)

  dataset = dataset.map(process_example)
  return dataset
def create_sprites_dataset(characters, actions, directions, channels=3,
                           length=8, shuffle=False, fake_data=False):
  """Creates a tf.data pipeline for the sprites dataset.

  Args:
    characters: A list of (skin, hair, top, pants) tuples containing
      relative paths to the sprite png image for each attribute.
    actions: A list of Actions.
    directions: A list of Directions.
    channels: Number of image channels to yield.
    length: Desired length of the sequences.
    shuffle: Whether or not to shuffle the characters and sequences
      start frame.
    fake_data: Boolean for whether or not to yield synthetic data.

  Returns:
    A tf.data.Dataset yielding (seq, skin label index, hair label index,
    top label index, pants label index, action label index, skin label
    name, hair label_name, top label name, pants label name, action
    label name) tuples.
  """
  if fake_data:
    dummy_image = tf.random.normal([HEIGHT, WIDTH, CHANNELS])
  else:
    basedir = download_sprites()

  action_names = [action.name for action in actions]
  action_metadata = [(action.start_row, action.frames) for action in actions]

  direction_rows = [direction.row_offset for direction in directions]

  chars = tf.data.Dataset.from_tensor_slices(characters)
  act_names = tf.data.Dataset.from_tensor_slices(action_names).repeat()
  acts_metadata = tf.data.Dataset.from_tensor_slices(action_metadata).repeat()
  dir_rows = tf.data.Dataset.from_tensor_slices(direction_rows).repeat()

  if shuffle:
    chars = chars.shuffle(len(characters))

  dataset = tf.data.Dataset.zip((chars, act_names, acts_metadata, dir_rows))

  skin_table = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))
  hair_table = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))
  top_table = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))
  pants_table = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))
  action_table = tf.contrib.lookup.index_table_from_tensor(sorted(action_names))

  def process_example(attrs, act_name, act_metadata, dir_row_offset):
    """Processes a dataset row."""
    skin_name = attrs[0]
    hair_name = attrs[1]
    top_name = attrs[2]
    pants_name = attrs[3]

    if fake_data:
      char = dummy_image
    else:
      skin = read_image(basedir + os.sep + skin_name)
      hair = read_image(basedir + os.sep + hair_name)
      top = read_image(basedir + os.sep + top_name)
      pants = read_image(basedir + os.sep + pants_name)
      char = create_character(skin, hair, top, pants)

    if shuffle:
      seq = create_random_seq(char, act_metadata, dir_row_offset, length)
    else:
      seq = create_seq(char, act_metadata, dir_row_offset, length)
    seq = seq[..., :channels]  # limit output channels

    skin_idx = skin_table.lookup(skin_name)
    hair_idx = hair_table.lookup(hair_name)
    top_idx = top_table.lookup(top_name)
    pants_idx = pants_table.lookup(pants_name)
    act_idx = action_table.lookup(act_name)

    return (seq, skin_idx, hair_idx, top_idx, pants_idx, act_idx,
            skin_name, hair_name, top_name, pants_name, act_name)

  dataset = dataset.map(process_example)
  return dataset
Creates a tf.data pipeline for the sprites dataset.

  Args:
    characters: A list of (skin, hair, top, pants) tuples containing
      relative paths to the sprite png image for each attribute.
    actions: A list of Actions.
    directions: A list of Directions.
    channels: Number of image channels to yield.
    length: Desired length of the sequences.
    shuffle: Whether or not to shuffle the characters and sequences
      start frame.
    fake_data: Boolean for whether or not to yield synthetic data.

  Returns:
    A tf.data.Dataset yielding (seq, skin label index, hair label index,
    top label index, pants label index, action label index, skin label
    name, hair label_name, top label name, pants label name, action
    label name) tuples.