和上一版本区别：

• 读取了原始点云附上颜色属性，并进行分段。
• 增加了KDtree算法，通过点云的块内所有点匹配对应geojson所有中心点列表

from typing import List
import o3d_hdmap.open3d as o3d
import numpy as np
import glob
import pygal
import time
from loguru import logger
import ditu.topbind as tb
import os
import json
from tqdm import tqdm
import sys
import itertools
import argparse
# sys.path.insert(0, '/data/hongyuan/work/mpcv-point-cloud-layer-dectection/')


def intensity2rgb(intensities: np.ndarray):
    """
    default ColorScale in CloudCompare for intensity:
        Blue -> Green -> Yellow -> Red
    """
    Blue = [0, 0, 255]
    Green = [0, 255, 0]
    Yellow = [255, 255, 0]
    Red = [255, 0, 0]
    R_channel = [c[0] for c in [Blue, Green, Yellow, Red]]
    G_channel = [c[1] for c in [Blue, Green, Yellow, Red]]
    B_channel = [c[2] for c in [Blue, Green, Yellow, Red]]
    I_channel = [0, 255 / 3, 255 * 2 / 3, 255]
    R = np.interp(intensities, I_channel, R_channel)
    G = np.interp(intensities, I_channel, G_channel)
    B = np.interp(intensities, I_channel, B_channel)
    return np.column_stack((R, G, B)).astype(np.uint8)


def read_point_cloud(pcds_path: List):
    clouds = [tb.Pcd(p) for p in tqdm(pcds_path)]
    points = [c.points() for c in clouds]
    intensities = [c.intensities() for c in clouds]
    colors = intensity2rgb(np.array(list(itertools.chain(*intensities))))
    return np.vstack(points), colors


def pcd_intensity2rgb(input_pcd: str, output_pcd: str):
    pcd = tb.Pcd(input_pcd)
    points = pcd.points()
    colors = intensity2rgb(pcd.intensities())
    output = o3d.geometry.PointCloud()
    output.points = o3d.utility.Vector3dVector(points)
    output.colors = o3d.utility.Vector3dVector(colors / 255.0)
    o3d.io.write_point_cloud(output_pcd, output)
    logger.info(f'wrote to {output_pcd}')


if __name__ == '__main__':
    start = time.time()
    sys.argv = ' '.join([
        'cmd',
        f'--point-cloud-path /data/mpcv_lspo_download_data/prod/PLEF35196-2021-11-10-17-48-39/mapping_results/lidar_features_world_opt',
        f'--segment-output-path point_cloud_segment_test',
    ]).split()

    prog = 'python3 point_cloud_segment.py'
    description = (
        'run point_cloud_segment to obtain multiple segment point clouds')
    parser = argparse.ArgumentParser(prog=prog, description=description)
    parser.add_argument(
        '--point-cloud-path',
        type=str,
        help=f'point-cloud-path',
    )
    parser.add_argument(
        '--segment-output-path',
        type=str,
        help=f'segment-output-path',
    )
    args = parser.parse_args()
    point_cloud_path = args.point_cloud_path
    segment_output_path = args.segment_output_path
    os.makedirs(segment_output_path, exist_ok=True)

    paths = sorted(
        glob.glob(f'{point_cloud_path}/ground/*.pcd', recursive=True))
    points, colors = read_point_cloud(paths)
    logger.info(f'loaded {len(points):,} points from {len(paths):,} pcd files')

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    down_sample_voxel_size = 1
    pcd, _, down_sample_indexes = pcd.voxel_down_sample_and_trace(
        down_sample_voxel_size,
        min_bound=points.min(axis=0),
        max_bound=points.max(axis=0),
    )
    down_sample_points = np.asarray(pcd.points)
    logger.info(
        f'voxel down sample, #points: {len(points):,} -> {len(down_sample_points):,} (voxel_size: {down_sample_voxel_size}m)'
    )

    vq = tb.VectorQuery()
    vq.add(positions=down_sample_points)
    segs = vq.segmentize(
        radius=20,
        max_height_offset=3,
        max_cluster_capacity=5000,
    )
    logger.info(f'clustered to {len(segs)} segments')

    for segment_index, pt_indexes_of_seg in enumerate(segs):
        indexes = set()
        for idx in pt_indexes_of_seg:
            indexes.update(down_sample_indexes[idx])
        indexes = sorted(indexes)
        segment_points = points[indexes, :]
        segment_colors = colors[indexes, :] / 255.0
        segment_pcd = o3d.geometry.PointCloud()
        segment_pcd.points = o3d.utility.Vector3dVector(segment_points)
        segment_pcd.colors = o3d.utility.Vector3dVector(segment_colors)
        path = f'{segment_output_path}/segment_{segment_index:05d}.pcd'
        o3d.io.write_point_cloud(path, segment_pcd)
        logger.info(f'wrote to {path}')
    logger.info(
        f'point_cloud_segment.py 执行完毕, total time: {time.time()-start:.3f} (s)')

layer_dectection.py

from typing import List
import o3d_hdmap.open3d as o3d
import numpy as np
import glob
import time
import ditu.topbind as tb
import os
from loguru import logger
import json
import smpy.utils as smpy_utils
from scipy.spatial import cKDTree
from typing import Dict, List
import sys
import itertools
import argparse
# sys.path.insert(0, '/data/hongyuan/work/mpcv-point-cloud-layer-dectection/')


def get_road_surface_center_point(geojson_data: Dict):
    feature_center_point_list = []
    indices = []
    features = geojson_data['features']
    for idx in range(len(features)):
        feature = features[idx]
        if feature['properties']['type'] == 'road_surface':
            indices.append(idx)
            feature_center_point_list.append(
                feature['properties']['plane_parameters']['center'])
    return feature_center_point_list, indices


def gridify(points: np.ndarray, grid_size: int):
    box_min = points.min(axis=0)[:2]
    indices = (np.floor((points[:, :2] - box_min) / grid_size)).astype(int)
    return indices


if __name__ == '__main__':
    start = time.time()
    sys.argv = ' '.join([
        'cmd',
        f'--segment-input-path point_cloud_segment_test',
        f'--geojson-path /data/mpcv_lspo_download_data/prod/PLEF35196-2021-11-10-17-48-39/mapping_results/ground_pole_geojson_LC.json',
    ]).split()
    prog = 'python3 layer_dectection.py'
    description = (
        'run layer_dectection to dectect hot points from geojson by segment point cloud'
    )
    parser = argparse.ArgumentParser(prog=prog, description=description)
    parser.add_argument(
        '--segment-input-path',
        type=str,
        help=f'segment input path',
    )
    parser.add_argument(
        '--geojson-path',
        type=str,
        help=f'ground geojson path',
    )
    parser.add_argument(
        '--grid-size',
        type=float,
        default=2,
        help=f'grid size default value is 2m',
    )
    parser.add_argument(
        '--deltas',
        type=float,
        default=1,
        help=f'deltas value is z_max minus z_min for the grid,default 1m',
    )
    args = parser.parse_args()
    segment_input_path = args.segment_input_path
    geojson_path = args.geojson_path
    grid_size = args.grid_size
    deltas = args.deltas
    segment_pcds = glob.glob(f'{segment_input_path}/*.pcd')
    # segment_pcds = [
    #     p for p in segment_pcds
    #     if p.endswith('_27.pcd') or p.endswith('_28.pcd')
    # ]
    output_dir = f'output_grid_size_{grid_size}m_deltas_{deltas}m'
    os.makedirs(output_dir, exist_ok=True)

    for input_path in segment_pcds:
        logger.info(f'input segment pcd: {input_path}')
        output_npy = f'{output_dir}/{os.path.basename(input_path).replace(".pcd", ".npy")}'
        if os.path.isfile(output_npy):
            logger.warning(f'{output_npy} exists, skip detecting hot points')
            continue
        pcd = o3d.io.read_point_cloud(input_path)
        points = np.asarray(pcd.points)
        colors = np.asarray(pcd.colors)
        grid_indices = gridify(points, grid_size)
        hot_points = []
        for grid in sorted(set([tuple(t) for t in grid_indices])):
            indices_0 = np.where(
                np.logical_and(grid_indices[:, 0] == grid[0],
                               grid_indices[:, 1] == grid[1]))[0]
            grid_points = points[indices_0, :]
            z_min = np.min(grid_points, 0)[2]
            z_max = np.max(grid_points, 0)[2]
            if z_max - z_min < deltas:
                continue
            colors[indices_0, :] = [1, 0, 0]
            hot_points.extend(points[indices_0])
        if not hot_points:
            dummy_points = np.array([])
            dummy_points.reshape((-1, 3))
            np.save(output_npy, dummy_points)
            logger.info(f'generated dummy hot points: {output_npy}')
            continue
        hot_points = np.array(hot_points)
        np.save(output_npy, hot_points)
        logger.info(f'generated hot points: {output_npy} ({len(hot_points):,})')
        # for debugging
        output_path = f'{output_dir}/{os.path.basename(input_path)}'
        if os.path.isfile(output_path):
            logger.info(
                f'output segment pcd with bad points labeled: {output_path}')
            continue
        o3d.io.write_point_cloud(output_path, pcd)
        logger.info(f'wrote to {output_path}')

    center_point_enus = [
        np.load(path) for path in glob.glob(f'{output_dir}/*.npy')
    ]
    center_point_enus = np.vstack(center_point_enus)
    hot_points = tb.utils.voxel_down_sample(
        np.vstack(center_point_enus), voxel_size=1.0)
    logger.info(
        f'hot points: {len(hot_points):,} (#raw_points: {len(center_point_enus):,})'
    )
    logger.info(f'reading geojson from {geojson_path} ...')
    geojson_data = smpy_utils.read_json(geojson_path)
    # anchor_lla_ = tb.utils.ecef2lla( np.array(geojson_data['extrinsics'][0]['center']))
    # pcd = o3d.geometry.PointCloud()
    # pcd.points = o3d.utility.Vector3dVector(road_surface_center_point_list)
    # path = 'geojson.pcd'
    # o3d.io.write_point_cloud(path, pcd)
    # logger.info(f'wrote to {path}')

    # geojson center points
    logger.info(f'loading road surface cennter points from geojson...')
    road_surface_center_point_list, road_surface_indices = get_road_surface_center_point(
        geojson_data)
    logger.info(f'generating kdtree from geojson road surface centers...')
    tree = cKDTree(np.array(road_surface_center_point_list))
    search_radius = 4
    logger.info(
        f'query road surface centers by hot points... (radius: {search_radius})'
    )
    indic = tree.query_ball_point(hot_points, r=search_radius)

    issue_road_surface_list = set(itertools.chain(*indic))
    logger.info(f'#hit: {len(issue_road_surface_list):,}')
    for surface_idx in issue_road_surface_list:
        hot_surface = geojson_data['features'][surface_idx]
        if hot_surface['properties']['type'] != 'road_surface':
            continue
        hot_surface['properties']['stroke'] = [255, 0, 0]
        hot_surface['geometry']['coordinates'] = [[
            lla[0], lla[1], lla[2] + 0.1
        ] for lla in hot_surface['geometry']['coordinates']]
        hot_surface['properties']['is_hot'] = True

    geojson_data['features'] = [
        f for f in geojson_data['features']
        if f['properties'].get('is_hot', False)
    ]

    issue_road_surface_num = len(issue_road_surface_list)
    issue_box_percent = issue_road_surface_num / len(road_surface_indices) * 100
    logger.info(f'road_surface_num: {len(road_surface_indices):,}')
    logger.info(f'issue_road_surface_num: {issue_road_surface_num:,}')
    logger.info(f'issue_box_percent: {issue_box_percent:.2f}%')

    path = f"{output_dir}/output_geojson.json"
    with open(path, "w") as f:
        json.dump(geojson_data, f, indent=4)
    logger.info(f'wrote labelled geojson to {path}')

    logger.info(
        f'layer_dectection.py 执行完毕, total time：{time.time()-start:.2f} (s)')