def hierarchical_segmentation()

in gss/utils.py [0:0]

• 28 lines of code
• 10 McCabe index (conditional complexity)

def hierarchical_segmentation(pcd, pcd_color, scene_id, FLAGS,
    feature_mask=features3d.SimilarityMask(1, 1, 1, 1), seg=None):

    # load pre_computed
    with open(os.path.join(FLAGS.cgal_path, scene_id+'.pkl'), 'rb') as f:
        info = pickle.load(f)
    adj_mat = info['adj_mat']; A0 = info['A0']
    
    F0 = np.load(os.path.join(FLAGS.cgal_path, scene_id+'_shape.npy'))
    n_region = len(np.unique(F0))
    shapes = []
    for i in range(n_region):
        pcd_i = o3d.geometry.PointCloud()
        idx_i = np.where(F0 == i)
        pcd_i.points = o3d.utility.Vector3dVector(np.array(pcd.points)[idx_i])
        shapes += [pcd_i]

    feature_extractor = features3d.Features3D(pcd, pcd_color, shapes, F0, n_region, feature_mask, tau=FLAGS.tau, seg=seg)
    
    # stores list of regions sorted by their similarity
    S = _build_initial_similarity_set(A0, feature_extractor)

    # stores region label and its parent (empty if initial).
    R = {i : () for i in range(n_region)}

    A = [A0]    # stores adjacency relation for each step
    F = [F0]    # stores label pcd for each step

    # greedy hierarchical grouping loop
    while len(S):
        (s, (i, j)) = S.pop()
        t = feature_extractor.merge(i, j)
        # record merged region (larger region should come first)
        R[t] = (i, j) if feature_extractor.size[j] < feature_extractor.size[i] else (j, i)
        Ak = _new_adjacency_dict(A[-1], i, j, t)
        A.append(Ak)
        S = _merge_similarity_set(feature_extractor, Ak, S, i, j, t)
        F.append(_new_label_image(F[-1], i, j, t))

    # bounding boxes for each hierarchy
    L = feature_extractor.bbox
    return (R, F, L)