From 11ab9f0c46ea550aa63d51f2e22f1d6c63632825 Mon Sep 17 00:00:00 2001 From: Astatin3 <77305074+Astatin3@users.noreply.github.com> Date: Wed, 19 Jun 2024 16:32:25 -0600 Subject: [PATCH] Swap to pylibfreenect2 --- main.py | 197 +++++++++++++++++++++++++------------------------------- 1 file changed, 88 insertions(+), 109 deletions(-) diff --git a/main.py b/main.py index 31d2a58..1ce276c 100644 --- a/main.py +++ b/main.py @@ -1,133 +1,112 @@ -from freenect2 import Device, FrameType -import psutil - -from PIL import Image -import open3d as o3d +from threading import Thread +import cv2 import numpy as np +import ktb -def depth_to_points(depth_map): - height, width = depth_map.shape - x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height)) +import open3d as o3d - points = np.stack((x_coords, y_coords, depth_map), axis=-1) - points = points.reshape(-1, 3) - - return points - -def remove_invalid_points(points): - mask = points[:, 2] != 0 - points = points[mask] - - return points, mask - - -def undistort_points(points, params): - fx = params.fx - fy = params.fy - cx = params.cx - cy = params.cy - k1 = params.k1 - k2 = params.k2 - k3 = params.k3 - p1 = params.p1 - p2 = params.p2 - - # Normalize points - x = (points[:, 0] - cx) / fx - y = (points[:, 1] - cy) / fy - - # Compute radial distances - r2 = x ** 2 + y ** 2 - r4 = r2 ** 2 - r6 = r2 ** 3 - - # Compute radial distortion - radial_distortion = (1 + k1 * r2 + k2 * r4 + k3 * r6) - - # Compute tangential distortion - xy = x * y - xy2 = 2 * xy - x2 = x ** 2 - y2 = y ** 2 - tangential_distortion_x = p1 * xy2 + p2 * (r2 + 2 * x2) - tangential_distortion_y = p1 * (r2 + 2 * y2) + p2 * xy2 - - # Undistort points - undistorted_x = (x - tangential_distortion_x) / radial_distortion - undistorted_y = (y - tangential_distortion_y) / radial_distortion - - # Reproject points - undistorted_points = np.zeros_like(points) - undistorted_points[:, 0] = undistorted_x * fx + cx - undistorted_points[:, 1] = undistorted_y * fy + cy - undistorted_points[:, 2] = points[:, 2] # Preserve depth values - - return undistorted_points - - -voxel_size = 0.005 +voxel_size = 0.1 +fitness_minimum = 0.9 vis = o3d.visualization.Visualizer() -vis.create_window(height=480, width=640) +vis.create_window() reconstruction = o3d.geometry.PointCloud() -reconstruction.points = o3d.utility.Vector3dVector(np.random.random((5, 3))) +reconstruction.points = o3d.utility.Vector3dVector(np.random.rand(2, 3)) + + +# vis.add_geometry(pcd) vis.add_geometry(reconstruction) -import numpy as np +def get_pcd(k): + points, colors = k.get_ptcld(colorized=True, scale=1000) -threads = [] + points = points.reshape((-1, 3)) + colors = colors.reshape((-1, 3)) + colors[:, [0, 2]] = colors[:, [2, 0]] + + points = o3d.utility.Vector3dVector(points) + colors = o3d.utility.Vector3dVector(colors) + + pcd = o3d.geometry.PointCloud() + + pcd.points = points + pcd.colors = colors + + pcd.points = pcd.voxel_down_sample(voxel_size=voxel_size).points + + return pcd + +def calc_transformation(cur_pcd, prev_pcd, cur_fpfh, prev_fpfh): + result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching( + cur_pcd, prev_pcd, cur_fpfh, prev_fpfh, + mutual_filter=True, + max_correspondence_distance=voxel_size, + estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False), + ransac_n=4, checkers=[], + criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(max_iteration=100000, confidence=0.999) + ) + print(result.fitness) running = True -device = Device() -color_frame = None +def run_loop(): + k = ktb.Kinect() -n_cpus = psutil.cpu_count() + prev_pcd = None + prev_fpfh = None + prev_normals = None + + while running: + cur_pcd = get_pcd(k) + + cur_pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) + + cur_fpfh = o3d.pipelines.registration.compute_fpfh_feature(cur_pcd, + o3d.geometry.KDTreeSearchParamHybrid(radius=0.25, + max_nn=100)) + if prev_pcd is None: + reconstruction.points = cur_pcd.points + reconstruction.paint_uniform_color([0.8, 0.8, 0]) + vis.update_geometry(reconstruction) + else: + # thread = Thread(target=calc_transformation, args=(cur_pcd, prev_pcd, cur_fpfh, prev_fpfh)) + # thread.start() + + result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching( + cur_pcd, prev_pcd, cur_fpfh, prev_fpfh, + mutual_filter=True, + max_correspondence_distance=voxel_size, + estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False), + ransac_n=4, checkers=[], + criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(max_iteration=100000, confidence=0.999) + ) -with device.running(): - - for type_, frame in device: - - if type_ == FrameType.Color: - img = frame.to_image() - img = img.resize((424,512)) - color_frame = np.array(img) - color_frame = color_frame[:, :, :3] - width, height = color_frame.shape[:2] - color_frame = color_frame.reshape(width * height, 3) - - elif type_ == FrameType.Depth and \ - color_frame is not None: - - depth_map = frame.to_array() - depth_map /= 5 - # ir_image /= ir_image.max() - # ir_image = np.sqrt(ir_image) - - - points = depth_to_points(depth_map) - points, mask = remove_invalid_points(points) - points = undistort_points(points, device.ir_camera_params) - - points /= 512 - - reconstruction.points = o3d.utility.Vector3dVector(points) - # reconstruction.colors = o3d.utility.Vector3dVector(color_frame[mask]) + # print(result.transformation) + # eg_p2p = o3d.pipelines.registration.registration_icp( + # reconstruction, pcd, threshold, trans_init, + # o3d.pipelines.registration.TransformationEstimationPointToPoint()) + # if result.fitness > fitness_minimum: + reconstruction.points.extend(cur_pcd.transform(-result.transformation).points) vis.update_geometry(reconstruction) - if not vis.poll_events(): - break + prev_pcd = cur_pcd + prev_fpfh = cur_fpfh - vis.update_renderer() + print("update") -vis.close() -running = False +t = Thread(target=run_loop) +t.start() -print("joining threads...") +while running: + # vis.update_geometry(reconstruction) + # print("E") + if cv2.waitKey(1) & 0xFF == ord('q'): + running = False + running = vis.poll_events() + vis.update_renderer() -for thread in threads: - thread.join() \ No newline at end of file +t.join() \ No newline at end of file