952 lines
34 KiB
Python
952 lines
34 KiB
Python
r"""Generates CLEVR-Dialog dataset.
|
|
|
|
Needs access to the following files:
|
|
synonyms: Contains several synonyms for each word in the question/caption.
|
|
caption templates: List of caption templates.
|
|
question templates: List of question templates.
|
|
metainfo: Meta-information related to attributes and values of CLEVR objects.
|
|
|
|
Usage:
|
|
python -u generate_dataset.py \
|
|
--scene_path="data/scenes/CLEVR_train_scenes.json" \
|
|
--num_beams=100 \
|
|
--num_workers=12 \
|
|
--save_path="data/clevr_train_raw.json"
|
|
|
|
Author: Satwik Kottur
|
|
"""
|
|
|
|
|
|
import copy
|
|
import collections
|
|
import json
|
|
import multiprocessing
|
|
import os
|
|
import random
|
|
import re
|
|
import time
|
|
from absl import flags
|
|
from absl import app
|
|
import numpy as np
|
|
from tqdm import tqdm as progressbar
|
|
|
|
import clevr_utils as utils
|
|
import global_vars as gvars
|
|
# import constraints_splitB as constraints
|
|
import constraints
|
|
|
|
FLAGS = flags.FLAGS
|
|
flags.DEFINE_string('synonym_path', '/projects/abdessaied/clevr-dialog/templates/synonyms.json',
|
|
'Path to synonyms file')
|
|
flags.DEFINE_string('metainfo_path', '/projects/abdessaied/clevr-dialog/templates/metainfo.json',
|
|
'Path to meta information file')
|
|
flags.DEFINE_string('caption_template_root', '/projects/abdessaied/clevr-dialog/templates/captions/',
|
|
'Root to folder with caption templates')
|
|
flags.DEFINE_string('question_template_root', '/projects/abdessaied/clevr-dialog/templates/questions/',
|
|
'Root to folder with question templates')
|
|
flags.DEFINE_string('scene_path',
|
|
# '/projects/abdessaied/clevr-dialog/output/result_clevr_oroginal_test.json',
|
|
'/projects/abdessaied/clevr-dataset-gen/output_finetune_20_objs_with_masks_many_attr/CLEVR_scenes.json',
|
|
'Path to CLEVR scene path json file')
|
|
flags.DEFINE_string('scene_id_file', '',
|
|
'Path to specific CLEVR scene ids to generate dialogs')
|
|
flags.DEFINE_string('save_path', '/projects/abdessaied/clevr-dialog/output/raw_data_modified/dialogs_finetune_20_objects_10_rounds.json',
|
|
'Path to save the dataset json')
|
|
flags.DEFINE_integer('num_beams', 100, 'Number of beams in dialog search')
|
|
flags.DEFINE_integer('num_workers', 64, 'Number of workers to use in search')
|
|
flags.DEFINE_integer('captions_per_image', 5, 'Number of captions per image')
|
|
flags.DEFINE_integer('num_images', -1,
|
|
'Number of images to generate dialogs. -1 for all.')
|
|
flags.DEFINE_integer('num_rounds', 10, 'Number of rounds in each dialog')
|
|
|
|
|
|
# Number of beams and distribution of question types.
|
|
# Start cutting down beams after 5th round.
|
|
# Heuristics (for round 4):
|
|
# A. count <= 2 1 <= seek <= 3 exist <= 2
|
|
# B. count + exist <= 3
|
|
# C. Independent questions <= 1
|
|
# Heuristics (for round 5):
|
|
# A. count <= 2 2 <= seek <= 4 exist <= 2
|
|
# B. count + exist <= 3
|
|
# C. Independent questions <= 1
|
|
ranges = {3: {'indep': [0, 1], 'seek': [1, 4], 'exist': [0, 1],
|
|
'count': [0, 1], 'exist+count': [0, 2]},
|
|
4: {'indep': [0, 1], 'seek': [2, 4], 'exist': [0, 1],
|
|
'count': [0, 1], 'exist+count': [0, 2]},
|
|
5: {'indep': [0, 1], 'seek': [2, 5], 'exist': [0, 2],
|
|
'count': [0, 2], 'exist+count': [0, 3]},
|
|
6: {'indep': [0, 1], 'seek': [2, 5], 'exist': [0, 2],
|
|
'count': [0, 2], 'exist+count': [0, 3]},
|
|
7: {'indep': [0, 2], 'seek': [3, 5], 'exist': [0, 2],
|
|
'count': [0, 2], 'exist+count': [0, 3]},
|
|
8: {'indep': [0, 2], 'seek': [3, 6], 'exist': [0, 3],
|
|
'count': [0, 3], 'exist+count': [0, 3]},
|
|
9: {'indep': [0, 2], 'seek': [3, 6], 'exist': [0, 3],
|
|
'count': [0, 3], 'exist+count': [0, 4]}}
|
|
|
|
|
|
def mapping(tag):
|
|
"""Maps tag to attribute.
|
|
|
|
Args:
|
|
tag: An input tag
|
|
|
|
Returns:
|
|
tag_label: Label for the input tag
|
|
"""
|
|
|
|
return gvars.METAINFO['tag_map'][tag.replace('1', '')]
|
|
|
|
|
|
def inv_mapping(attribute, arg_id=0):
|
|
"""Inverse maps attribute to tag.
|
|
|
|
Args:
|
|
attribute: Name of the attribute
|
|
arg_id: Argument id to use. Append 1 if arg_id is 1, else nothing
|
|
|
|
Returns:
|
|
base_tag: The string for the tag
|
|
"""
|
|
|
|
base_tag = gvars.METAINFO['tag_inv_map'][attribute]
|
|
if arg_id > 0:
|
|
base_tag = base_tag[:-1] + str(arg_id) + base_tag[-1]
|
|
|
|
return base_tag
|
|
|
|
|
|
def get_tag_group(tag):
|
|
"""Gets the group id from tag string.
|
|
|
|
For example, tag string of <S> is 0, <S1> is 1.
|
|
Assumes single digit group id.
|
|
|
|
Args:
|
|
tag: Tag string
|
|
|
|
Returns:
|
|
group_id: Return extracted group id
|
|
"""
|
|
|
|
group_id = 0 if len(tag) <= 3 else int(tag[-2])
|
|
return group_id
|
|
|
|
|
|
def replace_attribute(text, tag, obj_group, eliminate=False):
|
|
"""Replaces the attribute tags in text using available object properties.
|
|
|
|
NOTE: If shape is to be replaced, we use 'thing' in its place.
|
|
|
|
Args:
|
|
text: The text template to perform replacement
|
|
tag: The tags to replace in the text
|
|
obj_group: Available object properties to replace with
|
|
eliminate: Eliminate the remaining attribute tags
|
|
|
|
Returns:
|
|
replaced_text: The replaced text
|
|
"""
|
|
|
|
group = get_tag_group(tag)
|
|
if mapping(tag) == 'relation':
|
|
# Actual relation tag, else position tag.
|
|
if tag == '<R>':
|
|
relation_list = gvars.METAINFO['relation_phrases'][obj_group['relation']]
|
|
relation_cand = random.choice(relation_list)
|
|
else:
|
|
relation_cand = obj_group['relation']
|
|
|
|
return text.replace(tag, relation_cand)
|
|
|
|
if mapping(tag) == 'shape':
|
|
if eliminate:
|
|
replacer = 'thing'
|
|
else:
|
|
replacer = str(obj_group['objects'][group][mapping(tag)])
|
|
|
|
# Plural forms for groups.
|
|
if obj_group.get('count', 1) > 1 or obj_group.get('use_plural', False):
|
|
replacer += 's'
|
|
elif mapping(tag) == 'count':
|
|
if eliminate:
|
|
replacer = ''
|
|
else:
|
|
replacer = str(obj_group['count'])
|
|
else:
|
|
if eliminate:
|
|
replacer = ''
|
|
else:
|
|
replacer = str(obj_group['objects'][group][mapping(tag)])
|
|
return text.replace(tag, replacer)
|
|
|
|
|
|
def realize_text_and_extract_scene(scene, template, filter_objs):
|
|
"""Samples attributes for template using filtered objects.
|
|
|
|
In addition, creates scene graph for the new information added.
|
|
|
|
Args:
|
|
scene: Current scene graph
|
|
template: Text template to use to generate questions
|
|
filter_objs: Set of objects satisfying constraints of current template
|
|
|
|
Returns:
|
|
sample: Contains the text realization and scene graph
|
|
"""
|
|
|
|
def default_list(): return collections.defaultdict(list)
|
|
graph = {'relationships': collections.defaultdict(default_list),
|
|
'counts': {}, 'exists': {}, 'history': [], 'objects': {}}
|
|
|
|
# number of inputs
|
|
n_inputs = template.get('inputs', 1)
|
|
# sample a text template
|
|
text_sample = random.choice(template['text'])
|
|
text_sample_index = template['text'].index(text_sample)
|
|
|
|
# extract attribute tags and get them into groups
|
|
tags = re.findall('(<[\d\w]*>)', text_sample)
|
|
|
|
tag_groups = collections.defaultdict(list)
|
|
for tag in tags:
|
|
group_id = get_tag_group(tag)
|
|
tag_groups[group_id].append(tag)
|
|
|
|
# sample a random element from filtered
|
|
arg_sample = random.choice(filter_objs)
|
|
# scene information obtained from the current round
|
|
graph_item = arg_sample['graph']
|
|
|
|
# remove tags from text not allowed by filter_objs
|
|
for arg_ind in range(n_inputs):
|
|
obj_sample = arg_sample['objects'][arg_ind]
|
|
avail_attrs = obj_sample['optional'] + obj_sample['required']
|
|
|
|
for ii in tag_groups[arg_ind][::-1]:
|
|
if mapping(ii) not in avail_attrs:
|
|
tag_groups[arg_ind].remove(ii)
|
|
text_sample = replace_attribute(
|
|
text_sample, ii, arg_sample, True)
|
|
|
|
# assert that all required attributes are present as tags
|
|
for attribute in obj_sample['required']:
|
|
required_tag = inv_mapping(attribute, arg_ind)
|
|
if required_tag not in tag_groups[arg_ind]:
|
|
print("required_tag: {}".format(required_tag))
|
|
print("template: {}".format(template))
|
|
assert required_tag in tag_groups[arg_ind], \
|
|
'A required attribute is missing in template!'
|
|
|
|
# start compiling tags to keep
|
|
tags_to_keep = [inv_mapping(ii, arg_ind)
|
|
for ii in obj_sample['required']]
|
|
|
|
# filter out those not present in text template
|
|
optional_tags = [inv_mapping(ii, arg_ind)
|
|
for ii in obj_sample['optional']]
|
|
optional_tags = [
|
|
ii for ii in optional_tags if ii in tag_groups[arg_ind]]
|
|
|
|
# if tags_to_keep is empty, sample from optional with 1:70 2:25 3:5
|
|
if len(optional_tags) > 0:
|
|
if len(tags_to_keep) > 0:
|
|
n_tags_sample = [0, 1, 2]
|
|
else:
|
|
n_tags_sample = [1, 2, 3]
|
|
n_sample = np.random.choice(n_tags_sample, 1,
|
|
p=gvars.METAINFO['probabilities'],
|
|
replace=False)
|
|
# lower cap at the length of optional
|
|
n_sample = min(n_sample[0], len(optional_tags))
|
|
if n_sample > 0:
|
|
tags_to_keep += random.sample(optional_tags, n_sample)
|
|
|
|
# now create a dictionary of placeholders with actual attribute values
|
|
for tag in tag_groups[arg_ind]:
|
|
remove = tag not in tags_to_keep
|
|
text_sample = replace_attribute(
|
|
text_sample, tag, arg_sample, remove)
|
|
|
|
# remove attributes from objects not included in tags_to_keep
|
|
if 'objects' in graph_item:
|
|
for ii in gvars.METAINFO['attributes']:
|
|
if inv_mapping(ii, arg_ind) not in tags_to_keep:
|
|
if ii in graph_item['objects'][arg_ind]:
|
|
del graph_item['objects'][arg_ind][ii]
|
|
|
|
# record the caption info
|
|
# Record info and merge scene graphs.
|
|
args = []
|
|
# if "unique-obj" == template['label']:
|
|
# print('yey')
|
|
for obj in arg_sample['objects']:
|
|
if obj is None:
|
|
continue
|
|
else:
|
|
for k in obj['required']:
|
|
arg = obj.get(k, None)
|
|
if arg is not None:
|
|
if arg not in args: # and type(arg) == str:
|
|
args.append(arg)
|
|
else:
|
|
arg = arg_sample.get(k, None)
|
|
if arg is not None and arg not in args and type(arg) == str:
|
|
args.append(arg)
|
|
arg = obj.get('attribute', None)
|
|
if arg is not None and arg not in args:
|
|
args.append(arg)
|
|
if template['label'] == 'obj-relation':
|
|
args.append(arg_sample['relation'])
|
|
|
|
if template['label'] == "count-att-no":
|
|
template['label'] = "count-att"
|
|
|
|
graph_item['round'] = 0
|
|
sample = {}
|
|
sample['template_info'] = [copy.deepcopy(template)]
|
|
sample['args'] = args
|
|
del sample['template_info'][-1]['text']
|
|
sample['template_info'][-1]['index'] = text_sample_index
|
|
sample['caption'] = text_sample
|
|
sample['template'] = template['label']
|
|
|
|
sample['dialog'] = []
|
|
|
|
# append history, update scene graph, and save the new scene graph
|
|
graph['history'].append(graph_item)
|
|
sample['graph'] = utils.merge_update_scene_graph(graph, graph_item)
|
|
return sample
|
|
|
|
|
|
def realize_question(dialog, template, filter_objs):
|
|
"""Samples attributes for template using filtered objects.
|
|
|
|
In addition, creates scene graph for the new information added.
|
|
|
|
Args:
|
|
scene: Current scene graph
|
|
template: Text template to use to generate questions
|
|
filter_objs: Set of objects satisfying constraints of current template
|
|
|
|
Returns:
|
|
sample: Contains the text realization and scene graph
|
|
"""
|
|
|
|
# Number of inputs.
|
|
n_inputs = template.get('inputs', 0)
|
|
# Sample a text template.
|
|
text_sample = random.choice(template['text'])
|
|
text_sample_index = template['text'].index(text_sample)
|
|
|
|
# Extract attribute tags and get them into groups.
|
|
tags = re.findall('(<[\d\w]*>)', text_sample)
|
|
tag_groups = collections.defaultdict(list)
|
|
for tag in tags:
|
|
group_id = get_tag_group(tag)
|
|
tag_groups[group_id].append(tag)
|
|
|
|
# Sample a random element from filtered.
|
|
arg_sample = random.choice(filter_objs)
|
|
|
|
# Remove tags from text not allowed by filter_objs.
|
|
for arg_ind in range(n_inputs):
|
|
obj_sample = arg_sample['objects'][arg_ind]
|
|
avail_attrs = obj_sample['optional'] + obj_sample['required']
|
|
|
|
for ii in tag_groups[arg_ind][::-1]:
|
|
if mapping(ii) not in avail_attrs:
|
|
tag_groups[arg_ind].remove(ii)
|
|
text_sample = replace_attribute(
|
|
text_sample, ii, arg_sample, True)
|
|
|
|
# Assert that all required attributes are present as tags.
|
|
for attribute in obj_sample['required']:
|
|
required_tag = inv_mapping(attribute, arg_ind)
|
|
# Make an exception for <R> and <P>
|
|
if required_tag == '<R>' and '<P>' in tag_groups[arg_ind]:
|
|
continue
|
|
assert required_tag in tag_groups[arg_ind], \
|
|
'A required attribute is missing in template!'
|
|
|
|
# Start compiling tags to keep.
|
|
tags_to_keep = [inv_mapping(ii, arg_ind)
|
|
for ii in obj_sample['required']]
|
|
# Filter out those not present in text template.
|
|
optional_tags = [inv_mapping(ii, arg_ind)
|
|
for ii in obj_sample['optional']]
|
|
optional_tags = [
|
|
ii for ii in optional_tags if ii in tag_groups[arg_ind]]
|
|
|
|
# If tags_to_keep is empty, sample from optional with (1:70, 2:25, 3:5).
|
|
if len(optional_tags) > 0:
|
|
if len(tags_to_keep) > 0:
|
|
n_tags_sample = [0, 1, 2]
|
|
else:
|
|
n_tags_sample = [1, 2, 3]
|
|
n_sample = np.random.choice(n_tags_sample, 1,
|
|
p=gvars.METAINFO['probabilities'],
|
|
replace=False)
|
|
# Lower cap at the length of optional.
|
|
n_sample = min(n_sample[0], len(optional_tags))
|
|
if n_sample > 0:
|
|
tags_to_keep += random.sample(optional_tags, n_sample)
|
|
|
|
# Now create a dictionary of placeholders with actual attribute values.
|
|
for tag in tag_groups[arg_ind]:
|
|
remove = tag not in tags_to_keep
|
|
text_sample = replace_attribute(
|
|
text_sample, tag, arg_sample, remove)
|
|
|
|
# Record info and merge scene graphs.
|
|
args = []
|
|
# if template['label'] == 'seek-attr-early':
|
|
# print('yey')
|
|
for obj in arg_sample['objects']:
|
|
if obj is None:
|
|
continue
|
|
else:
|
|
for k in obj['required']:
|
|
arg = obj.get(k, None)
|
|
if arg is not None:
|
|
if arg not in args:
|
|
args.append(arg)
|
|
else:
|
|
arg = arg_sample.get(k, None)
|
|
if arg is not None:
|
|
args.append(arg)
|
|
arg = obj.get('attribute', None)
|
|
if arg is not None and arg not in args:
|
|
args.append(arg)
|
|
|
|
# req_att_keys = [k for obj in arg_sample['objects'] for k in obj['required'] if obj is not None]
|
|
dialog_datum = {'question': text_sample, 'answer': arg_sample['answer'],
|
|
'template': template['label'], 'args': args}
|
|
dialog['template_info'].append(template.copy())
|
|
del dialog['template_info'][-1]['text']
|
|
dialog['template_info'][-1]['index'] = text_sample_index
|
|
if 'unique' in template['label']:
|
|
print('voila')
|
|
dialog['dialog'].append(dialog_datum)
|
|
graph_item = arg_sample['graph']
|
|
|
|
# If mergeable, add it to the objects list.
|
|
dialog['graph'] = utils.merge_update_scene_graph(
|
|
dialog['graph'], graph_item)
|
|
|
|
# If there are volatile objects in the graph item, remove them.
|
|
for obj in graph_item['objects'][::-1]:
|
|
if obj.get('volatile', False):
|
|
graph_item['objects'].remove(obj)
|
|
dialog['graph']['history'].append(graph_item)
|
|
return dialog
|
|
|
|
|
|
def clean_text_subroutine(text, thing, suffix):
|
|
"""Cleans the text and substitutes thing with object (subroutine).
|
|
|
|
Args:
|
|
text: Text string to be cleaned
|
|
thing: Whether to use 'thing' or 'object'
|
|
suffix: Either '?' (question) or '.' (caption)
|
|
|
|
Returns:
|
|
clean_text: Text string after cleaning procedure
|
|
"""
|
|
|
|
# Synonyms + skipping optional part of the sentence
|
|
clean_text = skip_and_replace_phrases(text)
|
|
|
|
# Remove full stop, empty spaces, capitalize the start letter.
|
|
clean_text = re.sub(' +', ' ', clean_text.replace(suffix, '').strip(' '))
|
|
# First replace 'a thing' -> 'an object'.
|
|
# Then perform remaining actions.
|
|
if thing == 'object':
|
|
clean_text = clean_text.replace('a thing', 'an object')
|
|
clean_text = clean_text.replace('thing', thing)
|
|
clean_text = clean_text[0].upper() + clean_text[1:] + suffix
|
|
return clean_text
|
|
|
|
|
|
def clean_dialog_text(dialogs):
|
|
"""Cleans the dialog texts.
|
|
|
|
Args:
|
|
dialogs: Generated dialogs to perform text cleaning
|
|
|
|
Returns:
|
|
dialogs: Return the dialogs after cleaning the text inplace
|
|
"""
|
|
|
|
# Replace thing with object throughout with probability 0.5.
|
|
thing = 'thing' if random.random() > 0.5 else 'object'
|
|
for index, dialog_datum in enumerate(dialogs):
|
|
# Clean the caption.
|
|
text = dialog_datum['caption']
|
|
dialogs[index]['caption'] = clean_text_subroutine(text, thing, '.')
|
|
|
|
for r_id, dialog in enumerate(dialog_datum['dialog']):
|
|
# Clean the question.
|
|
text = dialog['question']
|
|
text = clean_text_subroutine(text, thing, '?')
|
|
dialogs[index]['dialog'][r_id]['question'] = text
|
|
return dialogs
|
|
|
|
|
|
def skip_and_replace_phrases(text):
|
|
"""Substitutes synonyms and skips optional parts stochastically.
|
|
|
|
Args:
|
|
text: Text string
|
|
|
|
Returns:
|
|
text: Text string with synonyms replaced and optional parts skipped
|
|
"""
|
|
|
|
# For each text in [], replace it with '' with probability 0.5.
|
|
matches = re.findall('(\[[ \w]*\])', text)
|
|
for match in matches:
|
|
if random.uniform(0, 1) > 0.5:
|
|
text = text.replace(match, '')
|
|
else:
|
|
text = text.replace(match, match[1:-1])
|
|
|
|
# Remove empty spaces, if any.
|
|
text = re.sub(' +', ' ', text)
|
|
# Search for synonyms, replace at uniformly random.
|
|
text = text.lower()
|
|
for key, values in gvars.METAINFO['synonym_keys']:
|
|
if key in text:
|
|
text = text.replace(key, random.choice(values))
|
|
return text
|
|
|
|
|
|
def generate_captions(scenes, templates):
|
|
"""Wrapper generates captions.
|
|
|
|
Args:
|
|
scenes: List of scene graphs for which to generate captions
|
|
templates: List of available caption templates
|
|
|
|
Returns:
|
|
generated_content: Captions generated for the input scenes
|
|
"""
|
|
|
|
template_dictionary = {ii['label']: ii for ii in templates}
|
|
generated_content = []
|
|
for scene in scenes['scenes'][0:FLAGS.num_images]:
|
|
content = {}
|
|
# Copy over image_index, split, image_filename from scene.
|
|
for key in ['image_index', 'split', 'image_filename']:
|
|
content[key] = scene[key]
|
|
|
|
content['dialogs'] = []
|
|
# Filter objects based on constraints.
|
|
filter_objs = constraints.caption(scene, templates)
|
|
for filter_obj in filter_objs:
|
|
# Realize the text, and return the partial scene knowledge (q).
|
|
template = template_dictionary[filter_obj[0]['graph']['template']]
|
|
sample = realize_text_and_extract_scene(
|
|
scene, template, filter_obj)
|
|
# Add it to the list of dialogs.
|
|
content['dialogs'].append(sample)
|
|
generated_content.append(content)
|
|
return generated_content
|
|
|
|
|
|
def generate_questions(scenes, dialogs, templates, params):
|
|
"""Wrapper generates questions.
|
|
|
|
Args:
|
|
scenes: List of scene graphs to generate questions
|
|
dialogs: Contains already generated captions for scenes graphs
|
|
templates: List of available question templates
|
|
params: Beam search parameters for question generation
|
|
|
|
Returns:
|
|
new_dialogs: Generated raw dialogs with captions and questions
|
|
"""
|
|
|
|
new_dialogs = []
|
|
for scene_id, dialog_datum in enumerate(dialogs):
|
|
image_dialogs = copy.deepcopy(dialog_datum)
|
|
image_dialogs['dialogs'] = []
|
|
|
|
for dialog in dialog_datum['dialogs']:
|
|
# Pick a template at random.
|
|
flag = False
|
|
iter_count = 0
|
|
while not flag:
|
|
# Pick a template at random.
|
|
template = random.choice(templates)
|
|
|
|
# Filter objects based on constraints.
|
|
filter_objs = constraints.question(scenes['scenes'][scene_id],
|
|
dialog, template)
|
|
flag = len(filter_objs) != 0
|
|
|
|
# Extreme case -- exit
|
|
iter_count += 1
|
|
if iter_count > 10:
|
|
break
|
|
|
|
# Realize q question.
|
|
if flag:
|
|
deep_copy = copy.deepcopy(dialog)
|
|
gen_dialog = realize_question(deep_copy, template, filter_objs)
|
|
image_dialogs['dialogs'].append(copy.deepcopy(gen_dialog))
|
|
new_dialogs.append(image_dialogs)
|
|
|
|
return new_dialogs
|
|
|
|
|
|
def worker(scenes, cap_templates, ques_templates, worker_id, out_q):
|
|
"""Worker method generates dialogs (caption + questions) for pool of scenes.
|
|
|
|
Args:
|
|
scenes: List of CLEVR scenes to generate dialogs
|
|
cap_templates: Templates for caption generation
|
|
ques_templates: Templates for question generation
|
|
worker_id: Id for the current worker
|
|
out_q: Output queue to save generated dialogs from different sources
|
|
|
|
Returns:
|
|
Adds dialogs against the worker id in the output queue.
|
|
"""
|
|
|
|
dialogs = []
|
|
for index, scene in enumerate(scenes):
|
|
cur_time = time.strftime('%a-%d%b%y-%X', time.gmtime())
|
|
print('Generating [ %s ] [ Worker: %d, Progress: %d/%d Scene: %d ]' %
|
|
(cur_time, worker_id, index, len(scenes), scene['image_index']))
|
|
try:
|
|
gen_dialog = generate_dialog_bfs(
|
|
scene, cap_templates, ques_templates)
|
|
dialogs.append(json.loads(json.dumps(gen_dialog)))
|
|
except:
|
|
print('NOTE: Missing data for %d' % scene['image_index'])
|
|
out_q.put({worker_id: dialogs})
|
|
|
|
|
|
def generate_dialog_bfs(scene, cap_templates, ques_templates):
|
|
"""Perform approximate breadth-first-search (BFS) to generate dialogs.
|
|
|
|
Args:
|
|
scene: Scene graph for the CLEVR image
|
|
cap_templates: List of caption templates
|
|
ques_templates: List of question templates
|
|
|
|
Returns:
|
|
bundle: List of dialogs generated for the input scene graph
|
|
"""
|
|
|
|
bundle = {}
|
|
# Generate captions for the scene.
|
|
# Copy over image_index, split, image_filename from scene.
|
|
for key in ['image_index', 'split', 'image_filename']:
|
|
bundle[key] = scene[key]
|
|
|
|
template_dictionary = {ii['label']: ii for ii in cap_templates}
|
|
content = {}
|
|
|
|
# Filter objects based on constraints on captions.
|
|
filter_objs = constraints.caption(scene, cap_templates)
|
|
|
|
for filter_obj in filter_objs:
|
|
# Realize the text, and return the partial scene knowledge (q).
|
|
template = template_dictionary[filter_obj[0]['graph']['template']]
|
|
sample = realize_text_and_extract_scene(scene, template, filter_obj)
|
|
# Add it to the list of dialogs.
|
|
content[template['label']] = [sample]
|
|
|
|
# Now generate questions.
|
|
# Group templates, exist/count of similar type together.
|
|
ques_groups = collections.defaultdict(list)
|
|
|
|
labels = [ii['label'] for ii in ques_templates]
|
|
# print('\n'.join(labels))
|
|
for index, ii in enumerate(ques_templates):
|
|
if 'exist' in ii['label'] or 'count' in ii['label']:
|
|
ques_groups[labels[index][4:]].append(ii)
|
|
else:
|
|
ques_groups[labels[index]].append(ii)
|
|
|
|
for round_id in range(FLAGS.num_rounds):
|
|
new_content = {}
|
|
|
|
# For each group.
|
|
for cap_label, cap_dialogs in content.items():
|
|
cur_pool = []
|
|
for dialog_datum in cap_dialogs:
|
|
for _, group in ques_groups.items():
|
|
template = random.choice(group)
|
|
|
|
# Make a copy.
|
|
datum_copy = copy.deepcopy(dialog_datum)
|
|
|
|
# Filter objects based on constraints.
|
|
filter_objs = constraints.question(
|
|
scene, datum_copy, template)
|
|
|
|
if len(filter_objs) == 0:
|
|
continue
|
|
|
|
# Realize q question.
|
|
gen_dialog = realize_question(
|
|
datum_copy, template, filter_objs)
|
|
cur_pool.append(gen_dialog)
|
|
|
|
if round_id in ranges:
|
|
for d_id, dialog in enumerate(cur_pool):
|
|
n_types = {'indep': 0, 'seek': 0, 'exist': 0, 'count': 0}
|
|
keep_dialog = True
|
|
|
|
labels = [ii['label']
|
|
for ii in dialog['template_info'][1:]]
|
|
for label in labels:
|
|
if label in gvars.METAINFO['independent_questions']:
|
|
n_types['indep'] += 1
|
|
|
|
label_type = label.split('-')[0]
|
|
n_types[label_type] += 1
|
|
|
|
# Heuristic A, C
|
|
for q_type, count in n_types.items():
|
|
limit = ranges[round_id][q_type]
|
|
if limit[0] > count or count > limit[1]:
|
|
keep_dialog = False
|
|
break
|
|
|
|
# Heuristic B
|
|
limit = ranges[round_id]['exist+count']
|
|
if n_types['count'] + n_types['exist'] > limit[1]:
|
|
keep_dialog = False
|
|
if not keep_dialog:
|
|
cur_pool[d_id] = None
|
|
cur_pool = [ii for ii in cur_pool if ii is not None]
|
|
|
|
# Keep limited number of beams (for speed).
|
|
if len(cur_pool) > FLAGS.num_beams:
|
|
cur_pool = sample_beams(cur_pool)[:FLAGS.num_beams]
|
|
new_content[cap_label] = cur_pool
|
|
content = copy.deepcopy(new_content)
|
|
|
|
# Get dialogs with sim, imm2, early questions.
|
|
for cap_label, cap_dialogs in content.items():
|
|
# Sample beams.
|
|
content[cap_label] = sample_beams(cap_dialogs)
|
|
|
|
# Remove keys that are empty.
|
|
empty_keys = [key for key, val in content.items() if len(val) == 0]
|
|
for key in empty_keys:
|
|
del content[key]
|
|
|
|
# For each caption, sample one.
|
|
sampled_dialogs = []
|
|
for cap_label, cap_dialogs in content.items():
|
|
if len(cap_dialogs) > 0:
|
|
sampled_dialogs.append(cap_dialogs.pop())
|
|
|
|
# Get 5 per image, compensate by taking from other entries.
|
|
content_keys = [ii for ii in content.keys()]
|
|
while len(sampled_dialogs) < 5:
|
|
random_label = random.choice(content_keys)
|
|
sampled_dialogs.append(cap_dialogs.pop())
|
|
|
|
# Finally, make the dialog text readable.
|
|
sampled_dialogs = clean_dialog_text(sampled_dialogs)
|
|
|
|
# Generate the coreference chain.
|
|
for dialog_id, dialog in enumerate(sampled_dialogs):
|
|
sampled_dialogs[dialog_id] = identify_coref_chains(dialog)
|
|
bundle['dialogs'] = sampled_dialogs
|
|
return bundle
|
|
|
|
|
|
def sample_beams(dialogs):
|
|
"""Samples beams based on the number of constraints satisfied.
|
|
|
|
Args:
|
|
dialogs: Generated dialogs to sample beams
|
|
|
|
Returns:
|
|
sampled_dialogs: List of sampled dialogs based on the constraints
|
|
"""
|
|
|
|
num_constraints = []
|
|
for d_id, dialog in enumerate(dialogs):
|
|
satisfied = 0
|
|
labels = [ii['label'] for ii in dialog['template_info'][1:]]
|
|
|
|
# Have a imm2 for sure
|
|
satisfied += np.sum(['imm2' in ii for ii in labels])
|
|
# Have a imm2 for sure
|
|
satisfied += np.sum(['sim' in ii for ii in labels])
|
|
# Have 'early'
|
|
satisfied += min(4, np.sum(['early' in ii for ii in labels]))
|
|
|
|
# Add it with the number of constraints it satisfies.
|
|
num_constraints.append((satisfied, d_id))
|
|
|
|
# Then order.
|
|
def sort_key(x): return (x[0], random.random())
|
|
ids = sorted(num_constraints, key=sort_key, reverse=True)
|
|
sampled_dialogs = [dialogs[ii[1]] for ii in ids]
|
|
return sampled_dialogs
|
|
|
|
|
|
def identify_coref_chains(dialog):
|
|
"""Identifies the coreference chains in generated dialog.
|
|
|
|
Args:
|
|
dialog: Generated dialogs for which coreference chains to be identified
|
|
|
|
Returns:
|
|
dialog: A copy of dialog, with coreference chains annotated
|
|
"""
|
|
|
|
for r_id, datum in enumerate(dialog['dialog']):
|
|
label = datum['template']
|
|
if label in gvars.METAINFO['independent_questions']:
|
|
dialog['graph']['history'][r_id + 1]['dependence'] = None
|
|
continue
|
|
|
|
if (label == 'exist-attribute-group' or label == 'count-attribute-group' or
|
|
label == 'count-all-group'):
|
|
dialog['graph']['history'][r_id + 1]['dependence'] = r_id - 1
|
|
continue
|
|
|
|
if 'imm' in label:
|
|
dialog['graph']['history'][r_id + 1]['dependence'] = r_id - 1
|
|
continue
|
|
|
|
if 'early' in label:
|
|
# Go over previous history.
|
|
cur_history = dialog['graph']['history'][r_id + 1]
|
|
assert 'focus_id' in cur_history and 'focus_desc' in cur_history,\
|
|
'More focus objects than one, no focus objects!'
|
|
focus_id = cur_history['focus_id']
|
|
for attr in gvars.METAINFO['attributes']:
|
|
if attr in cur_history['focus_desc']:
|
|
break
|
|
|
|
history = dialog['graph']['history'][:r_id + 1]
|
|
for hist_id, hist_datum in enumerate(history):
|
|
for obj in hist_datum['objects']:
|
|
if obj['id'] == focus_id and attr in obj:
|
|
dialog['graph']['history'][r_id +
|
|
1]['dependence'] = hist_id - 1
|
|
break
|
|
return dialog
|
|
|
|
|
|
def main(unused_argv):
|
|
"""Main method generates the CLEVR-Dialog dataset.
|
|
"""
|
|
# Read the scene file.
|
|
with open(FLAGS.scene_path, 'r') as file_id:
|
|
scenes = json.load(file_id)
|
|
|
|
# Read the synonyms file.
|
|
with open(FLAGS.synonym_path, 'r') as file_id:
|
|
synonyms = json.load(file_id)
|
|
|
|
def sorter(x): return len(x[0].split(' '))
|
|
|
|
# Read the metainformation file.
|
|
with open(FLAGS.metainfo_path, 'r') as file_id:
|
|
gvars.METAINFO = json.load(file_id)
|
|
tag_inv_map = {attr: tag for tag, attr in gvars.METAINFO['tag_map'].items()
|
|
if tag != '<P>'}
|
|
gvars.METAINFO['tag_inv_map'] = tag_inv_map
|
|
gvars.METAINFO['synonym_keys'] = sorted(synonyms.items(),
|
|
key=sorter, reverse=True)
|
|
|
|
# Add ids to objects.
|
|
scenes = utils.add_object_ids(scenes)
|
|
scenes = utils.clean_object_attributes(scenes)
|
|
|
|
# Read the caption templates.
|
|
template_paths = os.listdir(FLAGS.caption_template_root)
|
|
cap_templates = []
|
|
for ii in template_paths:
|
|
with open(os.path.join(FLAGS.caption_template_root, ii), 'r') as file_id:
|
|
cur_templates = json.load(file_id)
|
|
cap_templates.extend(cur_templates)
|
|
# utils.pretty_print_templates(cap_templates, 1)
|
|
|
|
# Read the question templates.
|
|
template_paths = os.listdir(FLAGS.question_template_root)
|
|
ques_templates = []
|
|
for ii in template_paths:
|
|
with open(os.path.join(FLAGS.question_template_root, ii), 'r') as file_id:
|
|
cur_templates = json.load(file_id)
|
|
ques_templates.extend(cur_templates)
|
|
# utils.pretty_print_templates(ques_templates, 1)
|
|
|
|
# 1. Check if there a scene_id_file specified.
|
|
# 2. Check if num_images is -1
|
|
if FLAGS.scene_id_file != '':
|
|
with open(FLAGS.scene_id_file, 'r') as file_id:
|
|
missing_ids = [int(ii.strip('\n')) for ii in file_id.readlines()]
|
|
print('Dialogs missing for scenes: %d' % len(missing_ids))
|
|
|
|
# Create a image_index -> scenes list index dictionary
|
|
image_list_id_dict = {ii['image_index']: index
|
|
for index, ii in enumerate(scenes['scenes'])}
|
|
scenes_subset = [scenes['scenes'][image_list_id_dict[scene_id]]
|
|
for scene_id in missing_ids]
|
|
|
|
elif FLAGS.num_images == -1:
|
|
scenes_subset = scenes['scenes']
|
|
|
|
else:
|
|
scenes_subset = scenes['scenes'][0: FLAGS.num_images]
|
|
|
|
# BFS for each scene.
|
|
if FLAGS.num_workers == 1:
|
|
# Single thread version.
|
|
dialogs = []
|
|
for index, scene in enumerate(scenes_subset):
|
|
cur_time = time.strftime('%a-%d%b%y-%X', time.gmtime())
|
|
print('Generating [ %s ] [ Worker: %d, Progress: %d/%d Scene: %d ]' %
|
|
(cur_time, 0, index, len(scenes_subset), scene['image_index']))
|
|
gen_dialog = generate_dialog_bfs(
|
|
scene, cap_templates, ques_templates)
|
|
dialogs.append(gen_dialog)
|
|
|
|
else:
|
|
# Multithread version.
|
|
output_q = multiprocessing.Queue()
|
|
jobs = []
|
|
for worker_id in range(FLAGS.num_workers):
|
|
allotment = scenes_subset[worker_id::FLAGS.num_workers]
|
|
inputs = (allotment, cap_templates, ques_templates)
|
|
inputs += (worker_id, output_q)
|
|
|
|
process = multiprocessing.Process(target=worker, args=inputs)
|
|
jobs.append(process)
|
|
process.start()
|
|
|
|
# Wait for all the jobs to finish and collect the output.
|
|
final_results = {}
|
|
for _ in jobs:
|
|
final_results.update(output_q.get())
|
|
for job in jobs:
|
|
job.join()
|
|
|
|
# Flatten and sort.
|
|
final_results = [jj for _, ii in final_results.items() for jj in ii]
|
|
dialogs = sorted(final_results, key=lambda x: x['image_index'])
|
|
# utils.pretty_print_dialogs(dialogs)
|
|
|
|
# Save the dialogs.
|
|
print('Saving dialog at: %s' % FLAGS.save_path)
|
|
with open(FLAGS.save_path, 'w') as file_id:
|
|
json.dump(dialogs, file_id)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
gvars.initialize()
|
|
app.run(main)
|