Open3D 可视化 （car_demo2版本）

example.py

import sys
import time
import traceback
# from multiprocessing import Queue, Process
from queue import Empty, Queue
from threading import Thread

import numpy as np
import open3d as o3d

from o3dvisualization import initVis, getVis, rerender, replace_geometries

def generatePCs(queue: Queue):
    # 生产者线程，不停产生点云
    try:
        while True:  
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(np.zeros(1000,3))
            pcd.colors = o3d.utility.Vector3dVector(np.zeros(1000,3))
            queue.put(pcd)
    except Exception:
        traceback.print_exc()
        sys.exit(1)

def main():
    initVis("./view-1646672208.json")
    vis = getVis()
    
    queue = Queue()  # 由一个生产者线程向queue中写入点云(o3d.visualization.PointCloud格式)
    pcd_process = Thread(target=generatePCs, args=(queue,))
    pcd_process.start()
    
    lastUpdatePCTime = 0.0
    while True:
        time.sleep(0.01)  # To release GIL
        now = time.time()
        if now - lastUpdatePCTime > 0.333: # 至少隔1/3s才允许更新一帧显示
            try:
                pcd = queue.get(block=False)
            except Empty:
                if pcd_process.is_alive():
                    continue
                else:
                    break

            replace_geometries(vis, [pcd])
            lastUpdatePCTime = now
        rerender(vis)

o3dvisualization.py

"""
注：vis.addGeometry/clearGeometry是会引起相机参数重置的，而updateGeometry则不会！
"""
import time
from typing import List, Union

import cv2
import open3d as o3d
import numpy as np
import threading

COLOR_LIST = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 1], [1, 0, 0], [0, 1, 1]], dtype=np.float32)

vis = o3d.visualization.VisualizerWithKeyCallback()
visQueue = []
visGeometrys = []
visCameraParam: o3d.camera.PinholeCameraParameters = None

saveImgBaseDir = None
saveCount = 0


def setPCSaveBaseDir(s):
    global saveImgBaseDir
    saveImgBaseDir = s


def initVis(camera_param_file: Union[str, None]):
    # vis.register_animation_callback(visAnimCallback)
    vis.register_key_callback(ord("S"), lambda vis: save_view_point(vis, "./view-" + str(int(time.time())) + ".json"))
    vis.create_window(width=1386, height=752)
    if camera_param_file is not None:
        vis.get_view_control().convert_from_pinhole_camera_parameters(
            o3d.io.read_pinhole_camera_parameters(camera_param_file))
    renderOption: o3d.visualization.RenderOption = vis.get_render_option()
    renderOption.background_color = np.array([1, 1, 1], dtype=np.float32)
    renderOption.show_coordinate_frame = True
    renderOption.point_size = 10.0


def getVis():
    return vis


def runVis():
    vis.run()


def add_geometry(vis: o3d.visualization.Visualizer, geometry: o3d.geometry.Geometry):
    saveCameraParam(vis)
    visGeometrys.append(geometry)
    vis.add_geometry(geometry)
    loadCameraParam(vis)


def clear_geometry(vis: o3d.visualization.Visualizer):
    saveCameraParam(vis)
    vis.clear_geometries()
    visGeometrys.clear()
    loadCameraParam(vis)


def replace_geometries(vis: o3d.visualization.Visualizer, geometries: List[o3d.geometry.Geometry]):
    saveCameraParam(vis)
    vis.clear_geometries()
    visGeometrys = geometries.copy()
    for g in visGeometrys:
        vis.add_geometry(g)
    loadCameraParam(vis)


def update_geometry(vis: o3d.visualization.Visualizer, geometry: o3d.geometry.Geometry):
    try:
        visGeometrys.index(geometry)
        vis.update_geometry(geometry)
    except:
        add_geometry(vis, geometry)


def saveCameraParam(vis: o3d.visualization.Visualizer):
    global visCameraParam
    visCameraParam = vis.get_view_control().convert_to_pinhole_camera_parameters()


def loadCameraParam(vis: o3d.visualization.Visualizer):
    if visCameraParam is not None:
        viewControl: o3d.visualization.ViewControl = vis.get_view_control()
        viewControl.convert_from_pinhole_camera_parameters(visCameraParam)


def rerender(vis: o3d.visualization.Visualizer):
    global saveCount
    vis.poll_events()
    vis.update_renderer()
    if saveImgBaseDir is not None:
        saveCount += 1
        toSaveImg = vis.capture_screen_float_buffer()
        toSaveImg = (np.array(toSaveImg) * 255).astype(np.uint8)
        filename = saveImgBaseDir + str(saveCount) + ".jpg"
        print("saved: " + filename)
        cv2.imwrite(filename, cv2.cvtColor(toSaveImg, cv2.COLOR_RGB2BGR))


def save_view_point(vis, filename):
    param = vis.get_view_control().convert_to_pinhole_camera_parameters()
    o3d.io.write_pinhole_camera_parameters(filename, param)
    # vis.destroy_window()


"""
以下的函数是第一版通过animationCallback更新点云时留下的，现在暂不使用
仅因为备份需求仍保留此段代码
"""


def visAnimCallback(vis: o3d.visualization.Visualizer):
    if len(visQueue) > 0:
        item = visQueue.pop(0)
        visQueue.clear()  # 避免累积延迟队列中填充很多以前的帧，强行清空列表
        saveCameraParam(vis)
        doDrawPC(vis, item)
        rerender(vis)


def doDrawPC(vis: o3d.visualization.Visualizer, points: np.ndarray):
    pc = o3d.geometry.PointCloud()
    pc.points = o3d.utility.Vector3dVector(points[:, :3])
    pc.colors = o3d.utility.Vector3dVector(COLOR_LIST[points[:, 3].astype(np.int)])
    clear_geometry(vis)
    add_geometry(vis, pc)


def addPC(points: np.ndarray):
    visQueue.append(points)


def sample_code():
    import torch

    res = torch.load("test.pth")


    def changeImg():
        for i in range(5):  # len(res["points"])):
            print(i)
            addPC(res["points"][i])
            time.sleep(5)


    initVis()
    threading.Thread(target=changeImg).start()
    # changeImg()
    vis.run()