diff --git a/main.py b/main.py index 91b2366..31d2a58 100644 --- a/main.py +++ b/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 diff --git a/output.jpg b/output.jpg new file mode 100644 index 0000000..f1ebab7 Binary files /dev/null and b/output.jpg differ