| | import torch |
| | import open3d as o3d |
| |
|
| | class Node: |
| | |
| | def __init__(self, mask_list, visible_frame, contained_mask, point_ids, node_info, son_node_info): |
| | ''' |
| | mask_list: list of masks that is within this cluster |
| | visible_frame: one-hot vector, 1 if the node appears in the frame |
| | contained_mask: one-hot vector, 1 if the node is contained by the mask |
| | point_ids: the corresponding 3D point ids |
| | node_info: for debugging. The iteration and the index of the node in this iteration |
| | son_node_info: for debugging. Node infos from the last iteration that are merged into this node |
| | |
| | ''' |
| | self.mask_list = mask_list |
| | self.visible_frame = visible_frame |
| | self.contained_mask = contained_mask |
| | self.point_ids = point_ids |
| | self.node_info = node_info |
| | self.son_node_info = son_node_info |
| |
|
| |
|
| | @ staticmethod |
| | def create_node_from_list(node_list, node_info): |
| | mask_list = [] |
| | visible_frame = torch.zeros(len(node_list[0].visible_frame), dtype=bool).cuda() |
| | contained_mask = torch.zeros(len(node_list[0].contained_mask), dtype=bool).cuda() |
| | point_ids = set() |
| | son_node_info = set() |
| | for node in node_list: |
| | mask_list += node.mask_list |
| | visible_frame = visible_frame | (node.visible_frame).bool() |
| | contained_mask = contained_mask | (node.contained_mask).bool() |
| | point_ids = point_ids.union(node.point_ids) |
| | son_node_info.add(node.node_info) |
| | return Node(mask_list, visible_frame.float(), contained_mask.float(), point_ids, node_info, son_node_info) |
| | |
| | def get_point_cloud(self, scene_points): |
| | ''' |
| | return: |
| | pcld: open3d.geometry.PointCloud object, the point cloud of the node |
| | point_ids: list of int, the corresponding 3D point ids of the node |
| | ''' |
| | point_ids = list(self.point_ids) |
| | points = scene_points[point_ids] |
| | pcld = o3d.geometry.PointCloud() |
| | pcld.points = o3d.utility.Vector3dVector(points) |
| | return pcld, point_ids |
| |
|
