def _fix_unassigned_points()

in wypr/dataset/s3dis/process.py [0:0]

• 31 lines of code
• 5 McCabe index (conditional complexity)

def _fix_unassigned_points(line):
    print(line)
    abs_out_path = os.path.abspath(PCD_PATH)
    area, room, _ = line.split('/')
    output_f = os.path.join(abs_out_path, area, room+'_shape.npy')
    if os.path.exists(output_f):
        return
    cgal_txt_f = os.path.join(abs_out_path, area, room+'.txt')
    assert os.path.exists(cgal_txt_f)
    with open(cgal_txt_f, 'r') as fin:
        cgal_txt = fin.readlines()
    pcd_path = os.path.join(abs_out_path, area, room+'.xyzn')
    assert os.path.exists(pcd_path)
    pcd = o3d.io.read_point_cloud(pcd_path)
    points = np.asarray(pcd.points)
    F0 = np.ones(points.shape[0]) * -1
    for i in range(len(cgal_txt[:-2])):
        try:
            row_list = [int(item) for item in cgal_txt[i].rstrip(' \n').split(' ')]
        except ValueError:
            print(i, cgal_txt_f)
            import pdb; pdb.set_trace()
        F0[np.array(row_list)] = i
    
    assigned_idx = np.where(F0!=-1)[0]
    unassigned_idx = np.where(F0==-1)[0]
    assigned_points = points[assigned_idx, :3]
    unassigned_points = points[unassigned_idx, :3]
    dist, idx = spatial.KDTree(assigned_points).query(unassigned_points)
    F0[unassigned_idx] = F0[assigned_idx][idx]
    assert np.sum(F0==-1) == 0
    np.save(output_f, F0)