clevr_utils.py

"""Utilities for CLEVR-Dialog dataset generation.

Author: Satwik Kottur
"""

import copy


def pretty_print_templates(templates, verbosity=1):
  """Pretty prints templates.

  Args:
    templates: Templates to print
    verbosity: 1 to print name and type of the templates
  """

  # Verbosity 1: Name and type.
  print('-'*70)
  for ii in templates:
    print('[Name: %s] [Type: %s]' % (ii['name'], ii['type']))
  print('-'*70)
  print('Total of %s templates..' % len(templates))
  print('-'*70)


def pretty_print_scene_objects(scene):
  """Pretty prints scene objects.

  Args:
    scene: Scene graph containing list of objects
  """

  for index, ii in enumerate(scene['objects']):
    print_args = (index, ii['shape'], ii['color'], ii['size'], ii['material'])
    print('\t%d : %s-%s-%s-%s' % print_args)


def pretty_print_dialogs(dialogs):
  """Pretty prints generated dialogs.

  Args:
    dialogs: Generated dialogs to print
  """

  for scene_id, dialog_datum in enumerate(dialogs):
    for dialog in dialog_datum['dialogs']:
      print(dialog['caption'])
      for round_id, ii in enumerate(dialog['dialog']):
        coref_id = dialog['graph']['history'][round_id+1]['dependence']
        in_tuple = (round_id, ii['question'], str(ii['answer']),
                ii['template'], str(coref_id))
        print('\t[Q-%d: %s] [A: %s] [%s] [%s]' % in_tuple)


def merge_update_scene_graph(orig_graph, graph_item):
  """Merges two scene graphs into one.

  Args:
    orig_graph: Original scene graph
    graph_item: New graph item to add to the scene graph

  Returns:
    graph: Deep copy of the original scene graph after merging
  """

  graph = copy.deepcopy(orig_graph)
  # Local alias.
  objects = graph['objects']

  # If not mergeable, return the same scene graph.
  if not graph_item['mergeable']:
    return graph

  # 1. Go through each new object
  # 2. Find its batch in objects
  #   a. If found, assert for a clash of attributes, update
  #   b. If novel, just add the object as is
  for new_obj in graph_item['objects']:
    match_found = False
    obj = objects.get(new_obj['id'], None)

    if obj:
      # Assert for existing entries.
      for attr in new_obj:
        try:
          assert new_obj[attr] == obj.get(attr, new_obj[attr]),\
            'Some of the attributes do not match!'
        except: pdb.set_trace()

      # Add additional keys.
      objects[new_obj['id']].update(new_obj)
    else:
      # Add the new object.
      objects[new_obj['id']] = new_obj

  # if a relation, update it
  if 'relation' in graph_item:
    rel = graph_item['relation']
    ## update it with object 2 id
    id1 = graph_item['objects'][0]['id']
    id2 = graph_item['objects'][1]['id']
    rel_objs = graph['relationships'][rel][id1]
    rel_objs.append(id2)
    graph['relationships'][rel][id1] = rel_objs

  # update objects in graph
  graph['objects'] = objects
  return graph


def add_object_ids(scenes):
  """Adds object ids field for input scenes.

  Args:
    scenes: List of CLEVR scene graphs

  Returns:
    scenes: Adds object_id field for the objects in the scene graph inplace
  """

  for scene_id, scene in enumerate(scenes['scenes']):
    for obj_id, _ in enumerate(scene['objects']):
      scenes['scenes'][scene_id]['objects'][obj_id]['id'] = obj_id
  return scenes


def clean_object_attributes(scenes):
  """Cleans attributes for objects, keeping only attributes and id.

  Args:
    scenes: Scene graph to clean

  Returns:
    scenes: Cleaned up scene graphs inplace
  """

  keys = ['shape', 'size', 'material', 'color', 'id']
  for scene_id, scene in enumerate(scenes['scenes']):
    for obj_id, obj in enumerate(scene['objects']):
      new_obj = {key: obj[key] for key in keys}
      scenes['scenes'][scene_id]['objects'][obj_id] = new_obj
  return scenes


def pretty_print_corefs(dialog, coref_groups):
  """Prints coreferences for a dialog, higlighting different groups in colors.

  Args:
    dialog: Generated dialogs to print
    coref_groups: Coreference groups for dialogs
  """

  colorama.init()
  # Mapping of group_id -> color_ids for (foreground, background)
  color_map = {}
  groups = coref_groups.get(0, [])
  colored, color_map = pretty_print_coref_sentence(dialog['caption'], groups,
                                                   color_map)
  print('\n\nC: %s' % colored)
  for round_id, round_datum in enumerate(dialog['dialog']):
    question = round_datum['question']
    groups = coref_groups.get(round_id + 1, [])
    colored, color_map = pretty_print_coref_sentence(question, groups,
                                                     color_map)
    print('%d: %s' % (round_id, colored))


def pretty_print_coref_sentence(sentence, groups, color_map):
  """Prints a sentence containing difference coreference groups.

  Args:
    sentence: Text sentence
    groups: List of coreference groups with spans
    color_map: List of groups and associated color maps

  Returns:
    sentence: Text sentence with colors inserted
    color_map: Updated, if new groups in the current sentence
  """

  fore_colors = ['RED', 'GREEN', 'YELLOW', 'BLUE', 'MAGENTA']
  back_colors = ['BLACK', 'YELLOW', 'CYAN']
  insertions = []
  for group in groups:
    group_id = group['group_id']
    if group_id in color_map:
      forecolor_id, backcolor_id = color_map[group_id]
    else:
      num_groups = len(color_map)
      forecolor_id = num_groups % len(fore_colors)
      backcolor_id = num_groups // len(fore_colors)
      color_map[group_id] = (forecolor_id, backcolor_id)

    forecolor = fore_colors[forecolor_id]
    backcolor = back_colors[backcolor_id]
    insertions.append((group['span'][0], getattr(colorama.Fore, forecolor)))
    insertions.append((group['span'][0], getattr(colorama.Back, backcolor)))
    insertions.append((group['span'][1],
                       getattr(colorama.Style, 'RESET_ALL')))

  # Perform insertions.
  sentence = insert_into_sentence(sentence, insertions)
  return sentence, color_map


def insert_into_sentence(sentence, insertions):
  """Sorts and performs insertions from right.

  Args:
    sentence: Sentence to perform insertions into
    insertions: List of insertions, format: (position, text_insert)

  Returns:
    sentence: Inplace inserted sentence
  """

  insertions = sorted(insertions, key=lambda x: x[0], reverse=True)
  for position, text in insertions:
    sentence = sentence[:position] + text + sentence[position:]
  return sentence