Advertisement
Guest User

Untitled

a guest
Apr 6th, 2020
186
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 4.07 KB | None | 0 0
  1. def minimal_example():
  2.     #  Class containing all data for a specific frame
  3.     class Frame:
  4.         def __init__(self, sensor_data, frame_index, n_points):
  5.             self.color_img = sensor_data.get_rgb_image(frame_index=frame_index)
  6.             self.depth_img = sensor_data.get_depth(frame_index=frame_index)
  7.             self.pose      = sensor_data.get_pose(frame_index=frame_index)  # Camera to world (4x4)
  8.             self.R = self.pose[:3, :3]  # Pose rotation
  9.             self.t = self.pose[:3, 3]   # Pose translation
  10.  
  11.             self.p_color = np.empty((3, n_points))  # Points in color image
  12.             self.p_depth = np.empty((3, n_points))  # Points in depth image
  13.             self.X       = np.empty((3, n_points))  # Points in 3D-space, camera coordinate system
  14.  
  15.     # intrinsic camera matrix decomposition
  16.     class K:
  17.         def __init__(self, intrinsic_matrix):
  18.             self.fx = intrinsic_matrix[0, 0]
  19.             self.fy = intrinsic_matrix[1, 1]
  20.             self.f = np.array([self.fx, self.fy])
  21.             self.cx = intrinsic_matrix[0, 2]
  22.             self.cy = intrinsic_matrix[1, 2]
  23.             self.cz = intrinsic_matrix[2, 2]
  24.             self.c = np.array([self.cx, self.cy, self.cz])
  25.  
  26.     def to_3D(p, depth, K):
  27.         # 2D to 3D coordinates with depth
  28.         x = (p[0] - K.cx) / K.fx
  29.         y = (p[1] - K.cy) / K.fy
  30.         return np.array([depth * x, depth * y, depth]).squeeze()
  31.  
  32.     def to_2D(p, K):
  33.         x = (p[0] * K.fx) / p[2] + K.cx
  34.         y = (p[1] * K.fy) / p[2] + K.cy
  35.         return np.array([x, y, 1]).squeeze()
  36.         #return torch.Tensor([x, y, p[2]])
  37.  
  38.  
  39.     file_path = "/home/marda648/exjobb/datasets/ScanNet/sens/scene0000_00.sens"
  40.     sensor_data = SensorData(file_path, 20) # Load 20 frames from dataset
  41.     n_points = 6
  42.     margin = (10, 10)
  43.     rgb_dim = np.array((sensor_data.color_height, sensor_data.color_width))
  44.  
  45.     # Extract 2 frames
  46.     frame_indices = [5, 10]
  47.     f1 = Frame(sensor_data, frame_index=frame_indices[0], n_points=n_points)
  48.     f2 = Frame(sensor_data, frame_index=frame_indices[1], n_points=n_points)
  49.  
  50.     # Get camera matrices
  51.     A_color = sensor_data.intrinsic_color[:3, :3]
  52.     A_depth = sensor_data.intrinsic_depth[:3, :3]
  53.     K_color = K(A_color)
  54.     K_depth = K(A_depth)
  55.  
  56.     # Select points in frame 1
  57.     f1.p_color = np.array([[300.0, 680.0, 196.0, 390.0, 880.0, 808.0],
  58.                            [320.0, 900.0, 570.0, 870.0, 451.0, 174.0],
  59.                            [1,     1,     1,     1,     1,     1    ]])
  60.  
  61.     # Define points in world coordinates
  62.     X_world = np.empty((3, n_points))
  63.  
  64.     for i in range(n_points):
  65.         # Calculate all data for frame 1
  66.         f1.p_depth[:, i] = A_depth @ np.linalg.inv(A_color) @ f1.p_color[:, i]               # Get point in depth image
  67.         depth = interpolate_point(f1.depth_img, f1.p_depth[:, i]) / sensor_data.depth_shift  # Sample the depth
  68.         f1.X[:, i] = to_3D(f1.p_color[0:2, i], depth=depth, K=K_color)                       # Project to 3D
  69.  
  70.         # Calculate world coordinates
  71.         X_world[:, i] = f1.R @ f1.X[:, i] + f1.t
  72.  
  73.         # Calculate points in frame 2
  74.         f2.X[:, i] = f2.R.T @ (X_world[:, i] - f2.t)   # Project to camera coordinate system
  75.         f2.p_color[:, i] = to_2D(f2.X[:, i], K_color)  # Project into color image
  76.         f2.p_depth[:, i] = to_2D(f2.X[:, i], K_depth)  # Project into depth image
  77.  
  78.     # Plot result
  79.     fig, ax = plt.subplots(2, 2)
  80.     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")
  81.     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")
  82.     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")
  83.     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")
  84.  
  85.     plt.draw(), plt.pause(0.1), input(">>")
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement