Skip to content

Instantly share code, notes, and snippets.

@chunibyo-wly

chunibyo-wly/share_pcd_projection.py

Created December 10, 2025 06:38
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save chunibyo-wly/c1a391006f3c3f41411b05c3a1d36a19 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save chunibyo-wly/c1a391006f3c3f41411b05c3a1d36a19 to your computer and use it in GitHub Desktop.
Download ZIP
SHARE S20 slam project point cloud to undistort images
Raw
share_pcd_projection.py
import cv2
import numpy as np
import trimesh
from tqdm import tqdm
import json
import open3d as o3d
from copy import deepcopy
from pathlib import Path
import xml.etree.ElementTree as ET
from shutil import rmtree
PLY_PATH = Path("data2/2025-12-09_15-13-54_colorized.ply")
DATA_ROOT = Path("data2")
OUTPUT = Path("output3")
if OUTPUT.exists():
rmtree(OUTPUT)
OUTPUT.mkdir(exist_ok=True)
def create_camera(fov_y=60, aspect=16 / 9, near=0.1, far=1.0):
# 角度转弧度
fov_y_rad = np.radians(fov_y)
# 近平面与远平面的宽高
h_near = np.tan(fov_y_rad / 2) * near
w_near = h_near * aspect
h_far = np.tan(fov_y_rad / 2) * far
w_far = h_far * aspect
# 八个顶点
vertices = np.array(
[
[-w_near, h_near, -near], # 0: nlt
[w_near, h_near, -near], # 1: nrt
[w_near, -h_near, -near], # 2: nrb
[-w_near, -h_near, -near], # 3: nlb
[-w_far, h_far, -far], # 4: flt
[w_far, h_far, -far], # 5: frt
[w_far, -h_far, -far], # 6: frb
[-w_far, -h_far, -far], # 7: flb
]
)
# 定义线的连接关系（每两个顶点形成一条边）
edges = np.array(
[
[0, 1],
[1, 2],
[2, 3],
[3, 0], # 近平面
[4, 5],
[5, 6],
[6, 7],
[7, 4], # 远平面
[0, 4],
[1, 5],
[2, 6],
[3, 7], # 连接线
]
)
# 利用 trimesh 生成一个 Path3D 线框
frustum = trimesh.load_path(vertices[edges] / 10)
return frustum
def parse_traj(tag="left"):
file_path = DATA_ROOT / "undistort" / "TransformedCam.json"
with open(file_path, "r") as f:
data = json.load(f)
traj = []
for entry in data["frames"]:
if tag not in entry["file_path"]:
continue
T = np.array(entry["transform_matrix"])
traj.append((entry["file_path"], T))
return traj
def parse_intrinsics(tag="left"):
fisheye_xml = DATA_ROOT / "images" / f"{tag.title()}.opt"
undistort_xml = DATA_ROOT / "undistort" / f"{tag.title()}_undistort.opt"
tree = ET.parse(fisheye_xml)
root = tree.getroot()
fx = float(root.find(".//M_00").text)
fy = float(root.find(".//M_11").text)
tree = ET.parse(undistort_xml)
root = tree.getroot()
cx = float(root.find(".//PrincipalPoint/X").text)
cy = float(root.find(".//PrincipalPoint/Y").text)
w = int(root.find(".//Width").text)
h = int(root.find(".//Height").text)
K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
reso = (h, w)
return K, reso
def main():
camera = create_camera()
camera.export("camera.ply")
pcd = o3d.io.read_point_cloud(PLY_PATH.as_posix())
kdtree = o3d.geometry.KDTreeFlann(pcd)
for tag in ["left", "right"]:
traj = parse_traj(tag)
K, reso = parse_intrinsics(tag)
output_meshes = []
for i, (file_path, T) in enumerate(tqdm(traj[::])):
basename = Path(file_path).stem
cam_instance = camera.copy()
cam_instance.apply_transform(T)
v = [0, 0, 0, 1]
v = T @ v
[k, idx, _] = kdtree.search_radius_vector_3d(o3d.utility.Vector3dVector([v[:3]])[0], 5.0)
if k == 0:
continue
tmp_pcd = deepcopy(pcd.select_by_index(idx))
tmp_pcd = tmp_pcd.transform(np.linalg.inv(T))
pts = np.asarray(tmp_pcd.points)
color = np.asarray(tmp_pcd.colors)[pts[:, 2] < 0]
r, g, b = color[:, 0], color[:, 1], color[:, 2]
color = np.stack([b, g, r], axis=-1)
pts = pts[pts[:, 2] < 0]
pts2 = pts / pts[:, 2:3]
uv = (K @ pts2.T).T
h, w = reso
_mask = (0 <= uv[:, 0]) & (uv[:, 0] < w) & (0 <= uv[:, 1]) & (uv[:, 1] < h)
color = color[_mask]
uv = uv[_mask].astype(np.int32)
uv2d = {}
pts[:, 2] = np.abs(pts[:, 2])
for _uv, _pt, _c in zip(uv, pts[_mask], color):
if (_uv[0], _uv[1]) in uv2d and uv2d[(_uv[0], _uv[1])][0] < _pt[2]:
continue
uv2d[(_uv[0], _uv[1])] = (_pt[2], _c)
sorted_uv = sorted(uv2d.items(), key=lambda x: x[1][0], reverse=True)
input_img = cv2.imread(DATA_ROOT / "undistort" / file_path)
origin = input_img.copy()
img = np.zeros((h, w, 3), dtype=np.uint8)
for (u, v), (_z, c) in sorted_uv:
r = int(np.clip(50 / (_z + 0.5), 10, 50))
cv2.circle(img, (u, v), r, (c * 255).astype(np.int32).tolist(), -1)
input_img[np.where(img != 0)] = img[np.where(img != 0)]
cv2.imwrite((OUTPUT / basename).with_suffix(".png").as_posix(), np.hstack([origin, input_img, img]))
output_meshes.append(cam_instance)
combined = trimesh.util.concatenate(output_meshes)
combined.export(f"traj_{tag}.ply")
if __name__ == "__main__":
main()
@chunibyo-wly
Copy link
Author

chunibyo-wly commented Dec 10, 2025

sam3 segmentation and pcd index select

import cv2
import numpy as np
import trimesh
from tqdm import tqdm
import json
import open3d as o3d
from copy import deepcopy
from pathlib import Path
import xml.etree.ElementTree as ET
from shutil import rmtree
import torch

PLY_PATH = Path("data2/2025-12-09_15-13-54_colorized.ply")
DATA_ROOT = Path("data2")
OUTPUT = Path("output2")
if OUTPUT.exists():
    rmtree(OUTPUT)
OUTPUT.mkdir(exist_ok=True)


def create_camera(fov_y=60, aspect=16 / 9, near=0.1, far=1.0):
    # 角度转弧度
    fov_y_rad = np.radians(fov_y)

    # 近平面与远平面的宽高
    h_near = np.tan(fov_y_rad / 2) * near
    w_near = h_near * aspect
    h_far = np.tan(fov_y_rad / 2) * far
    w_far = h_far * aspect

    # 八个顶点
    vertices = np.array(
        [
            [-w_near, h_near, -near],  # 0: nlt
            [w_near, h_near, -near],  # 1: nrt
            [w_near, -h_near, -near],  # 2: nrb
            [-w_near, -h_near, -near],  # 3: nlb
            [-w_far, h_far, -far],  # 4: flt
            [w_far, h_far, -far],  # 5: frt
            [w_far, -h_far, -far],  # 6: frb
            [-w_far, -h_far, -far],  # 7: flb
        ]
    )

    # 定义线的连接关系（每两个顶点形成一条边）
    edges = np.array(
        [
            [0, 1],
            [1, 2],
            [2, 3],
            [3, 0],  # 近平面
            [4, 5],
            [5, 6],
            [6, 7],
            [7, 4],  # 远平面
            [0, 4],
            [1, 5],
            [2, 6],
            [3, 7],  # 连接线
        ]
    )

    # 利用 trimesh 生成一个 Path3D 线框
    frustum = trimesh.load_path(vertices[edges] / 10)
    return frustum


def parse_traj(tag="left"):
    file_path = DATA_ROOT / "undistort" / "TransformedCam.json"

    with open(file_path, "r") as f:
        data = json.load(f)
    traj = []
    for entry in data["frames"]:
        if tag not in entry["file_path"]:
            continue
        T = np.array(entry["transform_matrix"])
        traj.append((entry["file_path"], T))
    return traj


def parse_intrinsics(tag="left"):
    fisheye_xml = DATA_ROOT / "images" / f"{tag.title()}.opt"
    undistort_xml = DATA_ROOT / "undistort" / f"{tag.title()}_undistort.opt"

    tree = ET.parse(fisheye_xml)
    root = tree.getroot()
    fx = float(root.find(".//M_00").text)
    fy = float(root.find(".//M_11").text)

    tree = ET.parse(undistort_xml)
    root = tree.getroot()
    cx = float(root.find(".//PrincipalPoint/X").text)
    cy = float(root.find(".//PrincipalPoint/Y").text)
    w = int(root.find(".//Width").text)
    h = int(root.find(".//Height").text)

    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
    reso = (h, w)
    return K, reso


def get_valid_img(folder_path):
    col = Path(folder_path).glob("*.pt")

    result = {}
    for c in col:
        mask = torch.load(c.as_posix())["masks"].detach().cpu().numpy()
        img = np.zeros((mask.shape[-2], mask.shape[-1]), dtype=np.uint8)
        for _mask in mask:
            img[_mask[0] > 0] = 1
        if img.sum() > 300:
            result[c.stem] = img
        else:
            print(f"Skip {c.stem} with too few valid pixels: {img.sum()}")
    return result


def main():
    sprinklers = get_valid_img("sprinkler")
    alarms = get_valid_img("alarm")

    camera = create_camera()
    camera.export("camera.ply")

    pcd = o3d.io.read_point_cloud(PLY_PATH.as_posix())
    kdtree = o3d.geometry.KDTreeFlann(pcd)

    for tag in ["left", "right"]:
        traj = parse_traj(tag)
        K, reso = parse_intrinsics(tag)
        output_meshes = []
        for i, (file_path, T) in enumerate(tqdm(traj[::])):
            basename = Path(file_path).stem
            for cls_name, cls_mask in [("alarm", alarms), ("sprinkler", sprinklers)]:
                cam_instance = camera.copy()
                cam_instance.apply_transform(T)
                output_meshes.append(cam_instance)

                if basename not in cls_mask:
                    continue

                v = [0, 0, 0, 1]
                v = T @ v
                [k, idx, _] = kdtree.search_radius_vector_3d(o3d.utility.Vector3dVector([v[:3]])[0], 5.0)

                if k == 0:
                    continue

                tmp_pcd = deepcopy(pcd.select_by_index(idx))
                tmp_pcd = tmp_pcd.transform(np.linalg.inv(T))

                pts = np.asarray(tmp_pcd.points)
                color = np.asarray(tmp_pcd.colors)[pts[:, 2] < 0]
                # r, g, b = color[:, 0], color[:, 1], color[:, 2]
                # color = np.stack([b, g, r], axis=-1)
                pts = pts[pts[:, 2] < 0]

                pts2 = pts / pts[:, 2:3]
                uv = (K @ pts2.T).T

                h, w = reso
                _mask = (0 <= uv[:, 0]) & (uv[:, 0] < w) & (0 <= uv[:, 1]) & (uv[:, 1] < h)
                color = color[_mask]
                uv = uv[_mask].astype(np.int32)

                valid_mask = cls_mask[basename]
                valid_pts_idx = np.where(valid_mask[uv[:, 1], uv[:, 0]] > 0)[0]
                valid_pts = pts[_mask][valid_pts_idx]

                if len(valid_pts) == 0:
                    continue

                out_pcd = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(valid_pts))
                out_pcd.colors = o3d.utility.Vector3dVector(color[valid_pts_idx])
                o3d.io.write_point_cloud(OUTPUT / f"{basename}_{cls_name}.ply", out_pcd.transform(T))

        combined = trimesh.util.concatenate(output_meshes)
        combined.export(f"traj_{tag}.ply")


if __name__ == "__main__":
    main()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment