point cloud visuaize with open3d with multi windows

visualize.py

"""
Show three windows to visualize predictions, ground truth labels and real colors
press key t to synchronize the camera setting in win_follow
"""
import open3d as o3d
import numpy as np
import copy
import threading
import time
from easydict import EasyDict as edict
from matplotlib.cm import get_cmap
import traceback

gdata = edict()
gdata.room_name = 'Area_5_office_40.npy'
gdata.finished = False
gdata.cam_ex = None
gdata.changed = False

def vis_main_fn(pcd, win_idx):
    def callback(vis):
        print('callback execute')
        view_ctl = vis.get_view_control()
        cam = view_ctl.convert_to_pinhole_camera_parameters()
        gdata.cam_ex = cam.extrinsic
        gdata.changed = True
        print(gdata)
        return True

    vis = o3d.visualization.VisualizerWithKeyCallback()
    """The Key use GFLW keys mapping, https://www.glfw.org/docs/latest/group__keys.html
        T: 84
    """
    vis.register_key_callback(84, callback)
    vis.create_window(width=640, height=480, left=640*win_idx+50)
    vis.add_geometry(pcd)
    vis.run()
    gdata.finished = True
    vis.destroy_window()

def vis_follow_fn(pcd, win_idx):
    try:
        vis = o3d.visualization.VisualizerWithKeyCallback()
        vis.create_window(width=640, height=480, left=640*win_idx+50)
        vis.add_geometry(pcd)
        print('vis_gt ready...')
        while not gdata.finished:
            time.sleep(0.5)
            print(gdata)
            if not gdata.changed:
                continue
            print('CHANGE_FLAG Detect')
            view_ctl = vis.get_view_control()
            cam = view_ctl.convert_to_pinhole_camera_parameters()
            cam.extrinsic = gdata.cam_ex
            view_ctl.convert_from_pinhole_camera_parameters(cam)
            vis.poll_events()
            vis.update_renderer()
            gdata.changed = False    
        vis.destroy_window()
    except Exception:
        print(traceback.format_exc())

data = np.load('Area5/' + gdata.room_name)
print('data_shape: ', data.shape)
num = data.shape[0]
pt_idx = np.random.choice(range(num), int(1e6))

pcd1 = o3d.geometry.PointCloud()
pcd1.points = o3d.utility.Vector3dVector(data[pt_idx, :3])
pcd1.colors = o3d.utility.Vector3dVector(data[pt_idx, 3:6] / 255)

cmap = get_cmap('tab20')
colors = cmap(data[:, 6].astype(int))[:, :3]
pcd2 = o3d.geometry.PointCloud()
pcd2.points = o3d.utility.Vector3dVector(data[pt_idx, :3])
pcd2.colors = o3d.utility.Vector3dVector(colors[pt_idx])

# pred labels
import pickle
predictions = pickle.load(open('val5_0.5/pred_5.pickle', 'rb'))
pred_label = predictions[gdata.room_name]
assert pred_label.shape[0] == data.shape[0], f"pred_shape: {pred_label.shape[0]}, while gt_shape: {data.shape[0]}"
pred_colors = cmap(pred_label.astype(int))[:, :3]
pcd3 = o3d.geometry.PointCloud()
pcd3.points = o3d.utility.Vector3dVector(data[pt_idx, :3])
pcd3.colors = o3d.utility.Vector3dVector(pred_colors[pt_idx])


t0 = threading.Thread(target=vis_main_fn, args=(pcd1, 0))
t1 = threading.Thread(target=vis_follow_fn, args=(pcd2, 1))
t2 = threading.Thread(target=vis_follow_fn, args=(pcd3, 2))
t0.start()
t1.start()
t2.start()
t0.join()
t1.join()
t2.join()