def gen_kinem_labels(actions):
state = lib_asm.Assembly()
action_segs = tuple(gen_segments(actions))
for start, end in action_segs:
segment = actions.loc[start:end]
for row in segment.itertuples(index=False):
# label = row.label
# i_start = row.start
# i_end = row.end
arg1 = row.arg1
arg2 = row.arg2
parent = lib_asm.Link(arg1)
child = lib_asm.Link(arg2)
joint = lib_asm.Joint((arg1, arg2), 'rigid', arg1, arg2)
state = state.add_joint(
joint,
parent,
child,
# directed=False,
in_place=False)
joint = lib_asm.Joint((arg2, arg1), 'rigid', arg2, arg1)
state = state.add_joint(
joint,
child,
parent,
# directed=False,
in_place=False)
start_idx = actions.loc[start]['start']
end_idx = actions.loc[end]['end']
state_idx = utils.getIndex(state, kinem_vocab)
yield start_idx, end_idx, state_idx
def assemblyLabels(labels_arr, num_samples, assembly_vocab=[], symmetries=[]):
def getIndex(assembly):
for i, a in enumerate(assembly_vocab):
if a == assembly:
return i
else:
assembly_vocab.append(assembly)
return len(assembly_vocab) - 1
label_seq = np.zeros(num_samples, dtype=int)
assembly = lib_assembly.Assembly()
assembly_index = getIndex(assembly)
prev_end_idx = 0
prev_start_idx = -1
for i, (start_idx, end_idx, action, part1, part2) in labels_arr.iterrows():
if start_idx != prev_start_idx or end_idx != prev_end_idx:
assembly_index = getIndex(assembly)
label_seq[prev_end_idx:end_idx] = assembly_index
if action == 'connect':
assembly = assembly.add_joint(part1, part2, in_place=False, directed=False)
elif action == 'disconnect':
assembly = assembly.remove_joint(part1, part2, in_place=False, directed=False)
prev_start_idx = start_idx
prev_end_idx = end_idx
assembly_index = getIndex(assembly)
label_seq[end_idx:] = assembly_index
return label_seq
def _assemblyLabels(labels_arr, num_samples, assembly_vocab=[], symmetries=[]):
def getIndex(assembly):
for i, a in enumerate(assembly_vocab):
if a == assembly:
return i
else:
assembly_vocab.append(assembly)
return len(assembly_vocab) - 1
label_seq = np.zeros(num_samples, dtype=int)
assembly = lib_assembly.Assembly()
getIndex(assembly)
paths = [[assembly]]
for i, (_, _, action, part1, part2) in labels_arr.iterrows():
new_paths = []
for i, path in enumerate(paths):
states = list(
iterateSymmetries(action, part1, part2, path[-1], symmetries=symmetries)
)
for assembly in states:
getIndex(assembly)
new_path = path + [assembly]
new_paths.append(new_path)
paths = new_paths
return label_seq
def renameLink(assembly, old_name, new_name):
link_dict = {}
for link in assembly.links.values():
new_link = lib_assembly.Link(link.name.replace(old_name, new_name),
pose=link.pose)
if new_link.name in link_dict:
if link_dict[new_link.name] != new_link:
raise AssertionError()
continue
else:
link_dict[new_link.name] = new_link
new_links = list(link_dict.values())
new_joints = [
lib_assembly.Joint(joint.name.replace(old_name, new_name),
joint.joint_type,
joint.parent_name.replace(old_name, new_name),
joint.child_name.replace(old_name, new_name),
transform=joint._transform)
for joint in assembly.joints.values()
]
new_assembly = lib_assembly.Assembly(links=new_links,
joints=new_joints)
return new_assembly
def make_goal_state(furn_name):
if furn_name == 'Kallax_Shelf_Drawer':
connections = (
('side panel 1', 'front panel 1'),
('side panel 2', 'front panel 1'),
('bottom panel 1', 'front panel 1'),
('bottom panel 1', 'side panel 1'),
('bottom panel 1', 'side panel 2'),
('back panel 1', 'side panel 1'),
('back panel 1', 'side panel 2'),
('back panel 1', 'bottom panel 1'),
)
elif furn_name == 'Lack_Coffee_Table':
connections = (('leg 1', 'table top 1'), ('leg 2', 'table top 1'),
('leg 3', 'table top 1'), ('leg 4', 'table top 1'),
('shelf 1', 'leg 1'), ('shelf 1', 'leg 2'),
('shelf 1', 'leg 3'), ('shelf 1', 'leg 4'))
elif furn_name == 'Lack_TV_Bench':
connections = (('leg 1', 'table top 1'), ('leg 2', 'table top 1'),
('leg 3', 'table top 1'), ('leg 4', 'table top 1'),
('shelf 1', 'leg 1'), ('shelf 1', 'leg 2'),
('shelf 1', 'leg 3'), ('shelf 1', 'leg 4'))
elif furn_name == 'Lack_Side_Table':
connections = (
('leg 1', 'table top 1'),
('leg 2', 'table top 1'),
('leg 3', 'table top 1'),
('leg 4', 'table top 1'),
)
else:
err_str = f"Unrecognized furniture name: {furn_name}"
raise ValueError(err_str)
goal_state = lib_asm.Assembly()
for arg1, arg2 in connections:
link1 = lib_asm.Link(arg1)
link2 = lib_asm.Link(arg2)
joint_12 = lib_asm.Joint((arg1, arg2), 'rigid', arg1, arg2)
joint_21 = lib_asm.Joint((arg2, arg1), 'rigid', arg2, arg1)
goal_state = goal_state.add_joint(joint_12,
link1,
link2,
in_place=False)
goal_state = goal_state.add_joint(joint_21,
link2,
link1,
in_place=False)
return goal_state
def main(out_dir=None, data_dir=None, annotation_dir=None):
out_dir = os.path.expanduser(out_dir)
data_dir = os.path.expanduser(data_dir)
annotation_dir = os.path.expanduser(annotation_dir)
annotation_dir = os.path.join(annotation_dir, 'action_annotations')
vocab_fn = os.path.join(data_dir, 'ANU_ikea_dataset', 'indexing_files',
'atomic_action_list.txt')
with open(vocab_fn, 'rt') as file_:
action_vocab = file_.read().split('\n')
part_names = (label.split('pick up ')[1] for label in action_vocab
if label.startswith('pick up'))
new_action_vocab = tuple(f"{part}" for part in part_names)
logger = utils.setupRootLogger(filename=os.path.join(out_dir, 'log.txt'))
fig_dir = os.path.join(out_dir, 'figures')
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
out_labels_dir = os.path.join(out_dir, 'labels')
if not os.path.exists(out_labels_dir):
os.makedirs(out_labels_dir)
# gt_action = np.load(os.path.join(annotation_dir, 'gt_action.npy'), allow_pickle=True)
with open(os.path.join(annotation_dir, 'gt_segments.json'), 'r') as _file:
gt_segments = json.load(_file)
ann_seqs = {
seq_name: [ann for ann in ann_seq['annotation']]
for seq_name, ann_seq in gt_segments['database'].items()
}
kinem_vocab = [lib_asm.Assembly()]
all_label_index_seqs = collections.defaultdict(list)
for seq_name, ann_seq in ann_seqs.items():
logger.info(f"Processing sequence {seq_name}...")
furn_name, other_name = seq_name.split('/')
goal_state = make_goal_state(furn_name)
label_seq = tuple(ann['label'] for ann in ann_seq)
segment_seq = tuple(ann['segment'] for ann in ann_seq)
start_seq, end_seq = tuple(zip(*segment_seq))
df = pd.DataFrame({
'start': start_seq,
'end': end_seq,
'label': label_seq
})
df = df.loc[df['label'] != 'NA']
if not df.any().any():
warn_str = f"No labels: {furn_name}, {other_name}"
logger.warning(warn_str)
continue
out_path = os.path.join(out_labels_dir, furn_name)
if not os.path.exists(out_path):
os.makedirs(out_path)
try:
new_df = convert_labels(df)
except AssertionError as e:
logger.warning(f" Skipping video: {e}")
continue
new_label_seq = tuple(' '.join(part_name.split()[:-1])
for part_name in new_df['arg1'])
label_index_seq = tuple(
new_action_vocab.index(label) for label in new_label_seq)
all_label_index_seqs[furn_name].append(label_index_seq)
kinem_df = parse_assembly_actions(new_df, kinem_vocab)
kinem_states = tuple(kinem_vocab[i] for i in kinem_df['state'])
if not kinem_states[-1] == goal_state:
warn_str = f" Final structure != goal structure:\n{kinem_states[-1]}"
logger.warning(warn_str)
lib_asm.writeAssemblies(
os.path.join(out_path, f"{other_name}_kinem-state.txt"),
kinem_states)
df.to_csv(os.path.join(out_path, f"{other_name}_human.csv"),
index=False)
new_df.to_csv(os.path.join(out_path, f"{other_name}_kinem-action.csv"),
index=False)
kinem_df.to_csv(os.path.join(out_path,
f"{other_name}_kinem-state.csv"),
index=False)
if not any(label_seq):
logger.warning(f"No labels: {seq_name}")
lib_asm.writeAssemblies(os.path.join(out_labels_dir, "kinem-vocab.txt"),
kinem_vocab)
symbol_table = fstutils.makeSymbolTable(new_action_vocab)
for furn_name, label_index_seqs in all_label_index_seqs.items():
label_fsts = tuple(
fstutils.fromSequence(label_index_seq, symbol_table=symbol_table)
for label_index_seq in label_index_seqs)
union_fst = libfst.determinize(fstutils.easyUnion(*label_fsts))
union_fst.minimize()
# for i, label_fst in enumerate(label_fsts):
# fn = os.path.join(fig_dir, f"{furn_name}-{i}")
# label_fst.draw(
# fn, isymbols=symbol_table, osymbols=symbol_table,
# # vertical=True,
# portrait=True,
# acceptor=True
# )
# gv.render('dot', 'pdf', fn)
fn = os.path.join(fig_dir, f"{furn_name}-union")
union_fst.draw(
fn,
isymbols=symbol_table,
osymbols=symbol_table,
# vertical=True,
portrait=True,
acceptor=True)
gv.render('dot', 'pdf', fn)