Merge pull request #64 from rohitmenon86/main

Feature: Add capability for realsense to use depth stream as well

configs/live/realsense_rgbd.yaml

@@ -0,0 +1,80 @@
+Results:
+  save_results: False
+  save_dir: "results"
+  save_trj: False
+  save_trj_kf_intv: 0
+  use_gui: True
+  eval_rendering: False
+  use_wandb: False
+
+Dataset:
+  dataset_path: datasets/realsense/rgbd
+  sensor_type: 'depth'
+  pcd_downsample: 32
+  pcd_downsample_init: 32
+  adaptive_pointsize: True
+  point_size: 0.01
+  type: 'realsense'
+
+Training:
+  init_itr_num: 1050
+  init_gaussian_update: 100
+  init_gaussian_reset: 500
+  init_gaussian_th: 0.005
+  init_gaussian_extent: 30
+  tracking_itr_num: 50
+  mapping_itr_num: 50
+  gaussian_update_every: 150
+  gaussian_update_offset: 50
+  gaussian_th: 0.7
+  gaussian_extent: 1.0
+  gaussian_reset: 2001
+  size_threshold: 20
+  kf_interval: 5
+  window_size: 8
+  pose_window: 3
+  edge_threshold: 1.1
+  rgb_boundary_threshold: 0.01
+  single_thread: False
+  spherical_harmonics: True
+
+  kf_translation: 0.05
+  kf_min_translation: 0.02
+  kf_overlap: 0.9
+  kf_cutoff: 0.3
+  prune_mode: 'slam'
+
+  lr:
+    cam_rot_delta: 0.003
+    cam_trans_delta: 0.001
+
+
+opt_params:
+  iterations: 30000
+  position_lr_init: 0.0016
+  position_lr_final: 0.0000016
+  position_lr_delay_mult: 0.01
+  position_lr_max_steps: 30000
+  feature_lr: 0.0025
+  opacity_lr: 0.05
+  scaling_lr: 0.001
+  rotation_lr: 0.001
+  percent_dense: 0.01
+  lambda_dssim: 0.2
+  densification_interval: 100
+  opacity_reset_interval: 3000
+  densify_from_iter: 500
+  densify_until_iter: 15000
+  densify_grad_threshold: 0.0002
+
+model_params:
+  sh_degree: 0
+  source_path: ""
+  model_path: ""
+  resolution: -1
+  white_background: False
+  data_device: "cuda"
+
+pipeline_params:
+  convert_SHs_python: False
+  compute_cov3D_python: False

utils/dataset.py

@@ -431,9 +431,23 @@ def __init__(self, args, path, config):
         super().__init__(args, path, config)
         self.pipeline = rs.pipeline()
         self.h, self.w = 720, 1280
-        self.config = rs.config()
-        self.config.enable_stream(rs.stream.color, self.w, self.h, rs.format.bgr8, 30)
-        self.profile = self.pipeline.start(self.config)
+
+        self.depth_scale = 0
+        if self.config["Dataset"]["sensor_type"] == "depth":
+            self.has_depth = True 
+        else: 
+            self.has_depth = False
+
+        self.rs_config = rs.config()
+        self.rs_config.enable_stream(rs.stream.color, self.w, self.h, rs.format.bgr8, 30)
+        if self.has_depth:
+            self.rs_config.enable_stream(rs.stream.depth)
+
+        self.profile = self.pipeline.start(self.rs_config)
+
+        if self.has_depth:
+            self.align_to = rs.stream.color
+            self.align = rs.align(self.align_to)
 
         self.rgb_sensor = self.profile.get_device().query_sensors()[1]
         self.rgb_sensor.set_option(rs.option.enable_auto_exposure, False)
@@ -443,9 +457,8 @@ def __init__(self, args, path, config):
         self.rgb_profile = rs.video_stream_profile(
             self.profile.get_stream(rs.stream.color)
         )
-
         self.rgb_intrinsics = self.rgb_profile.get_intrinsics()
-
+        
         self.fx = self.rgb_intrinsics.fx
         self.fy = self.rgb_intrinsics.fy
         self.cx = self.rgb_intrinsics.ppx
@@ -464,15 +477,33 @@ def __init__(self, args, path, config):
             self.K, self.dist_coeffs, np.eye(3), self.K, (self.w, self.h), cv2.CV_32FC1
         )
 
-        # depth parameters
-        self.has_depth = False
-        self.depth_scale = None
+        if self.has_depth:
+            self.depth_sensor = self.profile.get_device().first_depth_sensor()
+            self.depth_scale  = self.depth_sensor.get_depth_scale()
+            self.depth_profile = rs.video_stream_profile(
+                self.profile.get_stream(rs.stream.depth)
+            )
+            self.depth_intrinsics = self.depth_profile.get_intrinsics()
+
+
+
 
     def __getitem__(self, idx):
         pose = torch.eye(4, device=self.device, dtype=self.dtype)
+        depth = None
 
         frameset = self.pipeline.wait_for_frames()
-        rgb_frame = frameset.get_color_frame()
+
+        if self.has_depth:
+            aligned_frames = self.align.process(frameset)
+            rgb_frame = aligned_frames.get_color_frame()
+            aligned_depth_frame = aligned_frames.get_depth_frame()
+            depth = np.array(aligned_depth_frame.get_data())*self.depth_scale
+            depth[depth < 0] = 0
+            np.nan_to_num(depth, nan=1000)
+        else:
+            rgb_frame = frameset.get_color_frame()
+
         image = np.asanyarray(rgb_frame.get_data())
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         if self.disorted:
@@ -484,7 +515,8 @@ def __getitem__(self, idx):
             .permute(2, 0, 1)
             .to(device=self.device, dtype=self.dtype)
         )
-        return image, None, pose
+
+        return image, depth, pose
 
 
 def load_dataset(args, path, config):