Attempt to add undistortion

main.py

@@ -1,12 +1,69 @@
 from freenect2 import Device, FrameType
 import psutil
-import cv2
 
+from PIL import Image
 import open3d as o3d
 import numpy as np
 
-from threading import Thread
-from copy import deepcopy
+def depth_to_points(depth_map):
+    height, width = depth_map.shape
+    x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height))
+
+    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
 
@@ -17,82 +74,49 @@ reconstruction = o3d.geometry.PointCloud()
 reconstruction.points = o3d.utility.Vector3dVector(np.random.random((5, 3)))
 vis.add_geometry(reconstruction)
 
+import numpy as np
 
 threads = []
 
-def parse_frame(depth_frame, color_frame):
-
-    global reconstruction
-
-
-    undistorted, registered, big_depth = device.registration.apply(depth_frame,
-                                                                   color_frame,
-                                                                   with_big_depth=True)
-
-
-
-
-
-    cv2.imshow("1", np.array(registered.to_image()))
-    # points = device.registration.get_points_xyz_array(undistorted)
-    # points = points[~np.isnan(points).any(axis=2)]
-    # points = points.reshape(-1, 3) / 5
-
-    # print(points.shape)
-
-
-    # if np.isnan(points).any():
-    #     print("nan")
-    #     return
-    #
-    # reconstruction.points = o3d.utility.Vector3dVector(points)
-    # vis.update_geometry(reconstruction)
-
-
-    print("123")
-    # global threads
-    # threads.pop(thread_index)
 
 running = True
 device = Device()
 
-depth_frame = None
 color_frame = None
 
 n_cpus = psutil.cpu_count()
 
+
 with device.running():
 
     for type_, frame in device:
 
-        if type_ == FrameType.Depth:
-            depth_frame = frame
-        elif type_ == FrameType.Color:
-            color_frame = frame
+        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)
 
-        if depth_frame is not None and \
+        elif type_ == FrameType.Depth and \
             color_frame is not None:
-            #
-            if len(threads) > 100:
-                depth_frame = None
-                color_frame = None
-                continue
 
-            undistorted, registered, big_depth = device.registration.apply(depth_frame,
-                                                                           color_frame,
-                                                                           with_big_depth=True)
+            depth_map = frame.to_array()
+            depth_map /= 5
+            # ir_image /= ir_image.max()
+            # ir_image = np.sqrt(ir_image)
 
-            cv2.imshow('Grayscale', np.array(registered.to_image()))
 
-            # cv2.imshow("1", np.array(registered.to_image()))
-            # thread = Thread(target=parse_frame, args=(depth_frame, color_frame))
-            # threads.append(thread)
-            # thread.start()
-            #
-            # print(len(threads))
-            #
-            # depth_frame = None
-            # color_frame = None
+            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])
+            vis.update_geometry(reconstruction)
 
         if not vis.poll_events():
             break