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