class DatasetStats(DatasetAction):
"""
Gather stats of the houses
Useful to get an idea of what's in the dataset
And must be called to generate the model frequency information that's used
in the NN modules
"""
def __init__(self,
details=False,
model_details=False,
save_freq=True,
save_dest=""):
"""
Parameters
----------
details (bool, optional): If true, then frequency information will be shown on screen
model_details (bool, optional): since there are so many model ids, this additional bool
controls if those should be printed
save_freq (bool, optional): if true, then the category frequency information will be saved
save_dest (string, optional): directory to which frequency information is saved
"""
self.room_count = 0
self.object_count = 0
self.room_types_count = {}
self.fine_categories_count = {}
self.coarse_categories_count = {}
self.final_categories_count = {}
self.models_count = {}
self.object_category = ObjectCategories()
self.floor_node_only = False
self.details = details
self.model_details = model_details
self.save_freq = save_freq
data_dir = utils.get_data_root_dir()
self.save_dest = f"{data_dir}/{save_dest}"
def step(self, houses, num_threads=1):
for house in houses:
self.room_count += len(house.rooms)
for room in house.rooms:
room_types = room.roomTypes
for room_type in room_types:
self.room_types_count[room_type] = \
self.room_types_count.get(room_type, 0) + 1
filters = [floor_node_filter] if self.floor_node_only else []
nodes = list(set([node for nodes in [room.get_nodes(filters) \
for room in house.rooms] for node in nodes \
if node.type == "Object"]))
for node in nodes:
self.object_count += 1
fine_category = self.object_category.get_fine_category(
node.modelId)
coarse_category = self.object_category.get_coarse_category(
node.modelId)
final_category = self.object_category.get_final_category(
node.modelId)
self.fine_categories_count[fine_category] = \
self.fine_categories_count.get(fine_category, 0) + 1
self.coarse_categories_count[coarse_category] = \
self.coarse_categories_count.get(coarse_category, 0) + 1
self.final_categories_count[final_category] = \
self.final_categories_count.get(final_category, 0) + 1
self.models_count[node.modelId] = \
self.models_count.get(node.modelId, 0) + 1
yield house
def final(self):
print(f"\nPrinting Results...")
print(
f"\nThere are {self.room_count} non-empty rooms in the selection.")
print(f"There are {self.object_count} objects in the rooms.")
print(
f"On average, there are {self.object_count/self.room_count:.3f} objects for each room\n"
)
print(
f"There are {len(self.fine_categories_count)} fine categories among these objects."
)
if self.details:
print(f"\n{'Model Category':40s}{'Occurence'}")
for category in sorted(list((self.fine_categories_count.items())),
key=lambda x: -x[1]):
print(f"{category[0]:40s}{category[1]}")
print(
f"\nThere are {len(self.coarse_categories_count)} coarse categories among these objects."
)
if self.details:
print(f"\n{'Coarse Category':40s}{'Occurence'}")
for category in sorted(list(
(self.coarse_categories_count.items())),
key=lambda x: -x[1]):
print(f"{category[0]:40s}{category[1]}")
print(
f"\nThere are {len(self.final_categories_count)} final categories among these objects."
)
if self.details:
print(f"\n{'Final Category':40s}{'Occurence'}")
for category in sorted(list((self.final_categories_count.items())),
key=lambda x: -x[1]):
print(f"{category[0]:40s}{category[1]}")
print(
f"\nThere are {len(self.models_count)} unique models among these objects."
)
if self.details and self.model_details:
print(f"\n{'Model':40s}{'Occurence'}")
for category in sorted(list((self.models_count.items())),
key=lambda x: -x[1]):
print(f"{category[0]:40s}{category[1]}")
if self.save_freq:
with open(f"{self.save_dest}/fine_categories_frequency", "w") as f:
for cat in sorted(list((self.fine_categories_count.items())),
key=lambda x: -x[1]):
f.write(f"{cat[0]} {cat[1]}\n")
with open(f"{self.save_dest}/coarse_categories_frequency",
"w") as f:
for cat in sorted(list((self.coarse_categories_count.items())),
key=lambda x: -x[1]):
f.write(f"{cat[0]} {cat[1]}\n")
with open(f"{self.save_dest}/final_categories_frequency",
"w") as f:
for cat in sorted(list((self.final_categories_count.items())),
key=lambda x: -x[1]):
f.write(f"{cat[0]} {cat[1]}\n")
with open(f"{self.save_dest}/model_frequency", "w") as f:
for cat in sorted(list((self.models_count.items())),
key=lambda x: -x[1]):
f.write(f"{cat[0]} {cat[1]}\n")
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': '[email protected]',
'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