#!/usr/bin/env python3
import argparse
import csv
import glob
import os
from random import random, randint
import utils
from collections import namedtuple
from data import ObjectCategories, ObjectData, House
from math_utils import *
from math_utils.OBB import OBB
from math_utils.Simulator import Simulator
class ObjectCollection:
"""Provides observation information for a collection of objects"""
def __init__(self, categorization_type='final', sim_mode='direct'):
self._object_data = ObjectData()
self._object_categories = ObjectCategories()
self._objects = {}
self._room = None
self._categorization_type = categorization_type
self._sim = Simulator(mode=sim_mode)
@property
def simulator(self):
return self._sim
@property
def room(self):
return self._room
@property
def objects(self):
return self._objects
def add_object(self, o):
if o.id in self._objects and self._objects[o.id] != o:
print(f'Warning: ignoring node with duplicate node id={o.id}')
return None
if hasattr(o, 'type') and o.type == 'Room': # room node
self._room = o
self._sim.add_room(o, wall=True, floor=True, ceiling=False)
else: # other nodes
self._sim.add_object(o)
self.update(o, update_sim=True)
return o.id
def _remove(self, obj_id):
if obj_id not in self._objects:
print(f'Warning: tried to remove not present object with id={obj_id}')
else:
del self._objects[obj_id]
self._sim.remove(obj_id)
def update(self, o, xform=None, update_sim=False):
if not hasattr(o, 'category'):
o.category = self.category(o.modelId, scheme=self._categorization_type)
model_id = o.modelId if hasattr(o, 'modelId') else None
o.model2world = self.semantic_frame_matrix(model_id)
o.xform = xform if xform else Transform.from_node(o)
if hasattr(o, 'transform'):
o.transform = o.xform.as_mat4_flat_row_major()
o.obb = OBB.from_local2world_transform(np.matmul(o.xform.as_mat4(), o.model2world))
o.frame = self.node_to_semantic_frame(o, o.obb)
self._objects[o.id] = o
# room geometries pre-transformed, so after obb computation above is done, set back to identity transform
if hasattr(o, 'type') and o.type == 'Room':
o.xform = Transform()
o.transform = o.xform.as_mat4_flat_row_major()
if update_sim:
self._sim.set_state(obj_id=o.id, position=o.xform.translation, rotation_q=o.xform.rotation)
def randomize_object_transforms(self):
for o in self._objects.values():
if o is self._room:
continue
t = self._room.obb.sample()
t[1] = o.xform.translation[1]
o.xform.set_translation(t)
r = random() * 2 * math.pi
o.xform.set_rotation(radians=r)
self.update(o, o.xform)
def init_from_room(self, house, room_id, only_architecture=False, update_sim=True):
self.reset()
room = next(r for r in house.rooms if r.id == room_id)
self.add_object(room)
if not only_architecture:
for o in room.nodes:
self.add_object(o)
# if update_sim:
# self._sim.add_house_room_only(house, room, only_architecture=only_architecture, no_ceil=True, no_floor=True)
def init_from_house(self, house, update_sim=True):
self.reset()
for o in house.nodes:
self.add_object(o)
# if update_sim:
# self._sim.add_house(house, no_ceil=True, no_floor=True)
def as_house(self):
room_nodes = [dict(n.__dict__) for n in self._objects.values() if n.type != 'Room']
for i, n in enumerate(room_nodes):
n['originalId'] = n.originalId if hasattr(n, 'originalId') else n['id']
n['id'] = f'0_{str(i + 1)}' # overwrite id with linearized id
room = {
'id': '0_0',
'originalId': self.room.originalId if hasattr(self.room, 'originalId') else self.room.id,
'type': 'Room',
'valid': 1,
'modelId': self.room.modelId,
'nodeIndices': list(range(1, len(room_nodes) + 1)),
'roomTypes': self.room.roomTypes,
'bbox': self.room.bbox,
}
house_dict = {
'version': 'suncg@1.0.2',
'id': self.room.house_id,
'up': [0, 1, 0],
'front': [0, 0, 1],
'scaleToMeters': 1,
'levels': [{'id': '0', 'nodes': [room] + room_nodes}]
}
return House(house_json=house_dict)
def reset(self):
self._objects = {}
self._sim.reset()
def reinit_simulator(self, wall=True, floor=True, ceiling=False):
self._sim.reset()
for o in self._objects.values():
if hasattr(o, 'type') and o.type == 'Room': # room node:
self._sim.add_room(self.room, wall=wall, floor=floor, ceiling=ceiling)
else: # other nodes
self._sim.add_object(o)
self._sim.set_state(obj_id=o.id, position=o.xform.translation, rotation_q=o.xform.rotation)
def get_relative_observations(self, room_id, filter_ref_obj, ignore_categories):
out = {}
ref_objects = [filter_ref_obj] if filter_ref_obj else self._objects.values()
for o_r in ref_objects:
if o_r.category in ignore_categories:
continue
for o_i in self._objects.values():
if o_i is o_r:
continue
if o_i.category in ignore_categories:
continue
out[(o_i.id, o_r.id)] = self.object_frames_to_relative_observation(o_i.frame, o_r.frame, room_id)
return out
def get_collisions(self, include_collision_with_static=True, obj_id_a=None):
# update sim state to match state of this ObjectCollection
for o_id, o in self._objects.items():
self._sim.set_state(obj_id=o.id, position=o.xform.translation, rotation_q=o.xform.rotation)
self._sim.step() # sim step needed to create contacts
return self._sim.get_contacts(obj_id_a=obj_id_a, include_collision_with_static=include_collision_with_static)
def get_observation_key(self, observation):
room_node = self._objects[observation.room_id]
room_types = '-'.join(room_node.roomTypes) if hasattr(room_node, 'roomTypes') else ''
obj_node = self._objects[observation.obj_id]
obj_cat = self.category(obj_node.modelId, scheme=self._categorization_type)
ref_node = self._objects[observation.ref_id]
ref_cat = self.category(ref_node.modelId, scheme=self._categorization_type)
key = ObservationCategory(room_types=room_types, obj_category=obj_cat, ref_obj_category=ref_cat)
return key
def category(self, model_id, scheme):
if 'rm' in model_id:
return 'room'
if scheme == 'coarse':
return self._object_categories.get_coarse_category(model_id)
elif scheme == 'fine':
return self._object_categories.get_fine_category(model_id)
elif scheme == 'final':
return self._object_categories.get_final_category(model_id)
else:
raise RuntimeError(f'Unknown categorization type: {scheme}')
def semantic_frame_matrix(self, model_id):
if model_id in self._object_data.model_to_data:
return self._object_data.get_model_semantic_frame_matrix(model_id)
else: # not a model, so assume identity semantic frame
return np.identity(4)
def object_frames_to_relative_observation(self, frame, ref_frame, room_id):
ref_dims = ref_frame['obb'].half_dimensions
rel_centroid = ref_frame['obb'].transform_point(frame['obb'].centroid)
rel_min = ref_frame['obb'].transform_point(frame['aabb']['min'])
rel_max = ref_frame['obb'].transform_point(frame['aabb']['max'])
rel_up = ref_frame['obb'].transform_direction(frame['obb'].rotation_matrix[:3, 1])
rel_front = ref_frame['obb'].transform_direction(-frame['obb'].rotation_matrix[:3, 2]) # note: +z = back
cp = self._sim.get_closest_point(obj_id_a=frame['obj_id'], obj_id_b=ref_frame['obj_id'])
rel_cp = ref_frame['obb'].transform_point(cp.positionOnAInWS)
# NOTE: below approximate closest point calls are for removing pybullet call and debugging memory leak
# cp = frame['obb'].closest_point(ref_frame['obb'].centroid)
# rel_cp = ref_frame['obb'].transform_point(cp)
out = RelativeObservation(room_id=room_id, obj_id=frame['obj_id'], ref_id=ref_frame['obj_id'],
ref_dims=ref_dims, centroid=rel_centroid, min=rel_min, max=rel_max, closest=rel_cp,
front=rel_front, up=rel_up)
return out
@staticmethod
def node_to_semantic_frame(node, obb):
aabb_min, aabb_max = obb.to_aabb()
out = {
'obj_id': node.id,
'obb': obb,
'aabb': {'min': aabb_min, 'max': aabb_max}
}
return out
# an observation of object id relative to a reference object ref_id
class RelativeObservation(namedtuple('RelativeObservation',
['room_id', 'obj_id', 'ref_id', 'ref_dims', 'centroid', 'min', 'max', 'closest',
'front', 'up'])):
__slots__ = () # prevent per-instance dict
def to_str_list(self):
return [self.room_id, self.obj_id, self.ref_id, nparr2str_compact(self.ref_dims),
nparr2str_compact(self.centroid), nparr2str_compact(self.min), nparr2str_compact(self.max),
nparr2str_compact(self.closest), nparr2str_compact(self.front), nparr2str_compact(self.up)]
@classmethod
def fromstring(cls, s):
return RelativeObservation(room_id=s['room_id'], obj_id=s['obj_id'], ref_id=s['ref_id'], ref_dims=s['ref_dims'],
centroid=str2nparr(s['centroid']), min=str2nparr(s['min']), max=str2nparr(s['max']),
closest=str2nparr(s['closest']), front=str2nparr(s['front']), up=str2nparr(s['up']))
def parameterize(self, scheme):
if scheme == 'dist_angles':
return self._to_distance_angles()
elif scheme == 'dist_cos_sin_angles':
return self._to_distance_cos_sin_angles()
elif scheme == 'offsets_cos_sin_angles':
return self._to_offsets_cos_sin_angles()
elif scheme == 'offsets_angles':
return self._to_offsets_angles()
else:
raise RuntimeError(f'Unknown RelativeObservation parameterization scheme {scheme}')
def _to_distance_angles(self):
d_center, a_center = relative_pos_to_xz_distance_angle(self.centroid)
d_closest, a_closest = relative_pos_to_xz_distance_angle(self.closest)
a_front = relative_dir_to_xz_angle(self.front)
return [d_center, d_closest, a_center, a_closest, a_front]
def _to_distance_cos_sin_angles(self):
d_center, a_center = relative_pos_to_xz_distance_angle(self.centroid)
d_closest, a_closest = relative_pos_to_xz_distance_angle(self.closest)
a_front = relative_dir_to_xz_angle(self.front)
return [d_center, d_closest, math.cos(a_center), math.sin(a_center), math.cos(a_closest), math.sin(a_closest),
math.cos(a_front), math.sin(a_front)]
def _to_offsets_angles(self):
a_front = relative_dir_to_xz_angle(self.front)
return [self.centroid[0], self.centroid[2], self.closest[0], self.closest[2], a_front]
def _to_offsets_cos_sin_angles(self):
a_front = relative_dir_to_xz_angle(self.front)
return [self.centroid[0], self.centroid[2], self.closest[0], self.closest[2],
math.cos(a_front), math.sin(a_front)]
# a record categorizing a relative observation by originating room type, object category and reference object category
ObservationCategory = namedtuple('ObservationCategory', ['room_types', 'obj_category', 'ref_obj_category'])
class RelativeObservationsDatabase:
"""Encapsulates a set of object arrangement priors"""
def __init__(self, name, priors_dir, verbose=False, categorization_type='final'):
self._name = name
self._priors_dir = priors_dir
utils.ensuredir(self._priors_dir)
self._verbose = verbose
self._objects = ObjectCollection(categorization_type=categorization_type)
self._semantic_frames = {} # house_id -> {obj_id: SemanticFrame}
self._observations = {} # house_id -> {(obj_id,ref_obj_id): RelativeObservation}
self._grouped_observations = {} # {ObservationCategory: [RelativeObservation]}
def clear(self):
self._semantic_frames = {}
self._observations = {}
self._grouped_observations = {}
@property
def num_houses(self):
return len(self._observations.keys())
@property
def grouped_observations(self):
return self._grouped_observations
def collect_observations(self, houses):
for h in houses:
self._observations[h.id] = self._get_observations_from_house(h)
def save_observations_by_house(self, prefix='', save_frames=False):
for (house_id, house_observations) in self._observations.items():
basename = prefix + '_' + house_id
csv_file = os.path.join(self._priors_dir, basename + '.relpos.csv')
f = csv.writer(open(csv_file, 'w'))
f.writerow(RelativeObservation._fields)
for o in house_observations.values():
f.writerow(o.to_str_list())
if save_frames:
for (house_id, house_frames) in self._semantic_frames.items():
basename = prefix + '_' + house_id
csv_file = os.path.join(self._priors_dir, basename + '.semframes.csv')
f = csv.writer(open(csv_file, 'w'))
f.writerow(['obj_id', 'local2world', 'world_aabb_min', 'world_aabb_max'])
for frame in house_frames.values():
local2world_flat = nparr2str_compact(np.squeeze(np.asarray(frame['obb'].local2world.flatten())))
f.writerow([frame['obj_id'], local2world_flat,
nparr2str_compact(frame['aabb']['min']), nparr2str_compact(frame['aabb']['max'])])
def load_observations(self, filenames, house_dir, load_frames=False):
counts_house_ids = list(map(lambda f: f.split('_'), filenames))
num_observations = len(counts_house_ids)
for i, (obs_id, house_id) in enumerate(counts_house_ids):
print(f'Loading observation obs_id={obs_id} house id={house_id}, {i}/{num_observations}')
house_observations = self._observations.get(house_id, {})
num_rows = 0
for row in csv.DictReader(open(os.path.join(self._priors_dir, obs_id+'_'+house_id + '.relpos.csv'))):
if len(row) == 0:
continue
obs = RelativeObservation.fromstring(row)
house_observations[(obs.obj_id, obs.ref_id)] = obs
num_rows += 1
if num_rows > 0:
self._observations[obs_id] = house_observations
if load_frames:
house_frames = self._semantic_frames.get(house_id, {})
for row in csv.DictReader(open(os.path.join(self._priors_dir, obs_id+'_'+house_id + '.semframes.csv'))):
local2world = np.matrix(str2nparr(row['local2world'])).reshape(4, 4)
frame = {'obj_id': row['obj_id'], 'obb': OBB.from_local2world_transform(local2world),
'aabb': {'min': row['world_aabb_min'], 'max': row['world_aabb_max']}}
house_frames[frame['obj_id']] = frame
self._semantic_frames[house_id] = house_frames
print('Grouping observations by categories...')
groups = {}
num_obs = len(self._observations)
for i, (obs_id, house_id) in enumerate(counts_house_ids):
print(f'Observation in house id={house_id} {i}/{num_obs}')
house = House(file_dir=os.path.join(house_dir, obs_id + '.json'), include_support_information=False)
self._objects.init_from_house(house, update_sim=False)
for observation in self._observations[obs_id].values():
key = self._objects.get_observation_key(observation)
key_bin = groups.get(key, [])
key_bin.append(observation)
groups[key] = key_bin
print('Done grouping')
self._grouped_observations = groups
def save(self):
pkl_file = os.path.join(self._priors_dir, self._name + '.priors.pkl.gz')
utils.pickle_dump_compressed(self._grouped_observations, pkl_file)
def load(self):
pkl_file = os.path.join(self._priors_dir, self._name + '.priors.pkl.gz')
self._grouped_observations = utils.pickle_load_compressed(pkl_file)
def _get_observations_from_house(self, house):
if self._verbose:
print(f'Observations for house id={house.id}...')
self._objects.reset()
self._objects.simulator.add_house(house, no_ceil=True, no_floor=True)
relative_observations = {} # (node_i.id,node_j.id) -> observation of node_i relative to node_j
if len(house.rooms) == 0:
print(f'House id={house.id} has no rooms, skipping.')
return relative_observations
for room in house.rooms:
if self._verbose:
print(f'Room id={room.id}, roomTypes={room.roomTypes}')
# prepare node and room metadata
self._objects.add_object(room)
nodes = [n for n in room.nodes if n.type == 'Object' and n.valid]
for o in nodes:
self._objects.add_object(o)
room_relative_observations = self._objects.get_relative_observations(room.id,
filter_ref_obj=None,
ignore_categories=list())
relative_observations.update(room_relative_observations)
return relative_observations
if __name__ == '__main__':
module_dir = os.path.dirname(os.path.abspath(__file__))
parser = argparse.ArgumentParser(description='Compute arrangement priors')
parser.add_argument('--task', type=str, required=True, help='<Required> task [collect|save_pkl]')
parser.add_argument('--input', type=str, help='house json or directory with house json files')
parser.add_argument('--priors_dir', type=str, required=True, help='priors data directory')
parser.add_argument('--house_dir', type=str, help='house data directory')
args = parser.parse_args()
rod = RelativeObservationsDatabase(name='suncg_priors', priors_dir=args.priors_dir, verbose=True)
if os.path.isdir(args.input):
house_files = sorted(glob.glob(os.path.join(args.input, '**/*.json'), recursive=True))
else:
house_files = [args.input]
if args.task == 'collect':
for f in house_files:
house = House(file_dir=f, include_support_information=False)
prefix = os.path.splitext(os.path.basename(f))[0]
priors_file = os.path.join(args.priors_dir, prefix + '_' + house.id + '.relpos.csv')
if os.path.exists(priors_file):
print(f'Priors already exist at {priors_file}')
else:
rod.collect_observations([house])
rod.save_observations_by_house(prefix)
rod.clear()
if args.task == 'save_pkl':
relpos_files = filter(lambda p: '.relpos.csv' in p, os.listdir(args.input))
relpos_files = list(map(lambda p: os.path.basename(p).split('.')[0], relpos_files))
print(relpos_files)
rod.load_observations(relpos_files, args.house_dir, load_frames=False)
# reduced = {k: rod.grouped_observations[k] for k in [
# ObservationCategory('Bedroom', 'chair', 'desk'), ObservationCategory('Bedroom', 'desk', 'chair'),
# ]}
# rod._grouped_observations = reduced
rod.save()
print('DONE')
