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):
  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 ) : 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.
Creates a tf . data pipeline for the sprites dataset .