Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def minimal_example():
- # Class containing all data for a specific frame
- class Frame:
- def __init__(self, sensor_data, frame_index, n_points):
- self.color_img = sensor_data.get_rgb_image(frame_index=frame_index)
- self.depth_img = sensor_data.get_depth(frame_index=frame_index)
- self.pose = sensor_data.get_pose(frame_index=frame_index) # Camera to world (4x4)
- self.R = self.pose[:3, :3] # Pose rotation
- self.t = self.pose[:3, 3] # Pose translation
- self.p_color = np.empty((3, n_points)) # Points in color image
- self.p_depth = np.empty((3, n_points)) # Points in depth image
- self.X = np.empty((3, n_points)) # Points in 3D-space, camera coordinate system
- # intrinsic camera matrix decomposition
- class K:
- def __init__(self, intrinsic_matrix):
- self.fx = intrinsic_matrix[0, 0]
- self.fy = intrinsic_matrix[1, 1]
- self.f = np.array([self.fx, self.fy])
- self.cx = intrinsic_matrix[0, 2]
- self.cy = intrinsic_matrix[1, 2]
- self.cz = intrinsic_matrix[2, 2]
- self.c = np.array([self.cx, self.cy, self.cz])
- def to_3D(p, depth, K):
- # 2D to 3D coordinates with depth
- x = (p[0] - K.cx) / K.fx
- y = (p[1] - K.cy) / K.fy
- return np.array([depth * x, depth * y, depth]).squeeze()
- def to_2D(p, K):
- x = (p[0] * K.fx) / p[2] + K.cx
- y = (p[1] * K.fy) / p[2] + K.cy
- return np.array([x, y, 1]).squeeze()
- #return torch.Tensor([x, y, p[2]])
- file_path = "/home/marda648/exjobb/datasets/ScanNet/sens/scene0000_00.sens"
- sensor_data = SensorData(file_path, 20) # Load 20 frames from dataset
- n_points = 6
- margin = (10, 10)
- rgb_dim = np.array((sensor_data.color_height, sensor_data.color_width))
- # Extract 2 frames
- frame_indices = [5, 10]
- f1 = Frame(sensor_data, frame_index=frame_indices[0], n_points=n_points)
- f2 = Frame(sensor_data, frame_index=frame_indices[1], n_points=n_points)
- # Get camera matrices
- A_color = sensor_data.intrinsic_color[:3, :3]
- A_depth = sensor_data.intrinsic_depth[:3, :3]
- K_color = K(A_color)
- K_depth = K(A_depth)
- # Select points in frame 1
- f1.p_color = np.array([[300.0, 680.0, 196.0, 390.0, 880.0, 808.0],
- [320.0, 900.0, 570.0, 870.0, 451.0, 174.0],
- [1, 1, 1, 1, 1, 1 ]])
- # Define points in world coordinates
- X_world = np.empty((3, n_points))
- for i in range(n_points):
- # Calculate all data for frame 1
- f1.p_depth[:, i] = A_depth @ np.linalg.inv(A_color) @ f1.p_color[:, i] # Get point in depth image
- depth = interpolate_point(f1.depth_img, f1.p_depth[:, i]) / sensor_data.depth_shift # Sample the depth
- f1.X[:, i] = to_3D(f1.p_color[0:2, i], depth=depth, K=K_color) # Project to 3D
- # Calculate world coordinates
- X_world[:, i] = f1.R @ f1.X[:, i] + f1.t
- # Calculate points in frame 2
- f2.X[:, i] = f2.R.T @ (X_world[:, i] - f2.t) # Project to camera coordinate system
- f2.p_color[:, i] = to_2D(f2.X[:, i], K_color) # Project into color image
- f2.p_depth[:, i] = to_2D(f2.X[:, i], K_depth) # Project into depth image
- # Plot result
- fig, ax = plt.subplots(2, 2)
- ax[0, 0].imshow(f1.color_img), ax[0, 0].scatter(f1.p_color[1, :], f1.p_color[0, :], s=11, c='r', marker='x'), ax[0, 0].set_title("Frame 1: RGB")
- ax[1, 0].imshow(f1.depth_img), ax[1, 0].scatter(f1.p_depth[1, :], f1.p_depth[0, :], s=11, c='r', marker='x'), ax[1, 0].set_title("Frame 1: Depth")
- ax[0, 1].imshow(f2.color_img), ax[0, 1].scatter(f2.p_color[1, :], f2.p_color[0, :], s=11, c='r', marker='x'), ax[0, 1].set_title("Frame 2: RGB")
- ax[1, 1].imshow(f2.depth_img), ax[1, 1].scatter(f2.p_depth[1, :], f2.p_depth[0, :], s=11, c='r', marker='x'), ax[1, 1].set_title("Frame 2: Depth")
- plt.draw(), plt.pause(0.1), input(">>")
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement