ROS implementation of Foundation Pose along with Grounded SAM for first-frame mask

LICENSE → FoundationPose_ROS/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2022-Present, NVIDIA Corporation & affiliates. All rights reserved.
+Copyright (c) 2022-2023, NVIDIA Corporation & affiliates. All rights reserved.
 
 
 =======================================================================
@@ -91,4 +91,4 @@ LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
 COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
 THE POSSIBILITY OF SUCH DAMAGES.
 
-=======================================================================
+=======================================================================

Utils.py → FoundationPose_ROS/Utils.py

@@ -158,7 +158,7 @@ def nvdiffrast_render(K=None, H=None, W=None, ob_in_cams=None, glctx=None, conte
 
   ob_in_glcams = torch.tensor(glcam_in_cvcam, device='cuda', dtype=torch.float)[None]@ob_in_cams
   if projection_mat is None:
-    projection_mat = projection_matrix_from_intrinsics(K, height=H, width=W, znear=0.001, zfar=100)
+    projection_mat = projection_matrix_from_intrinsics(K, height=H, width=W, znear=0.1, zfar=100)
   projection_mat = torch.as_tensor(projection_mat.reshape(-1,4,4), device='cuda', dtype=torch.float)
   mtx = projection_mat@ob_in_glcams
 
@@ -318,7 +318,7 @@ def bilateral_filter_depth_kernel(depth:wp.array(dtype=float, ndim=2), out:wp.ar
         if v<0 or v>=H:
           continue
         cur_depth = depth[v,u]
-        if cur_depth>=0.001 and cur_depth<zfar:
+        if cur_depth>=0.1 and cur_depth<zfar:
           num_valid += 1
           mean_depth += cur_depth
     if num_valid==0:
@@ -335,7 +335,7 @@ def bilateral_filter_depth_kernel(depth:wp.array(dtype=float, ndim=2), out:wp.ar
         if v<0 or v>=H:
           continue
         cur_depth = depth[v,u]
-        if cur_depth>=0.001 and cur_depth<zfar and abs(cur_depth-mean_depth)<0.01:
+        if cur_depth>=0.1 and cur_depth<zfar and abs(cur_depth-mean_depth)<0.01:
           weight = wp.exp( -float((u-w)*(u-w) + (h-v)*(h-v)) / (2.0*sigmaD*sigmaD) - (depthCenter-cur_depth)*(depthCenter-cur_depth)/(2.0*sigmaR*sigmaR) )
           sum_weight += weight
           sum += weight*cur_depth
@@ -364,7 +364,7 @@ def erode_depth_kernel(depth:wp.array(dtype=float, ndim=2), out:wp.array(dtype=f
     if w>=W or h>=H:
       return
     d_ori = depth[h,w]
-    if d_ori<0.001 or d_ori>=zfar:
+    if d_ori<0.1 or d_ori>=zfar:
       out[h,w] = 0.0
     bad_cnt = float(0)
     total = float(0)
@@ -376,7 +376,7 @@ def erode_depth_kernel(depth:wp.array(dtype=float, ndim=2), out:wp.array(dtype=f
           continue
         cur_depth = depth[v,u]
         total += 1.0
-        if cur_depth<0.001 or cur_depth>=zfar or abs(cur_depth-d_ori)>depth_diff_thres:
+        if cur_depth<0.1 or cur_depth>=zfar or abs(cur_depth-d_ori)>depth_diff_thres:
           bad_cnt += 1.0
     if bad_cnt/total>ratio_thres:
       out[h,w] = 0.0
@@ -397,7 +397,7 @@ def erode_depth(depth, radius=2, depth_diff_thres=0.001, ratio_thres=0.8, zfar=1
 
 
 def depth2xyzmap(depth, K, uvs=None):
-  invalid_mask = (depth<0.001)
+  invalid_mask = (depth<0.1)
   H,W = depth.shape[:2]
   if uvs is None:
     vs,us = np.meshgrid(np.arange(0,H),np.arange(0,W), sparse=False, indexing='ij')
@@ -423,7 +423,7 @@ def depth2xyzmap_batch(depths, Ks, zfar):
   @Ks: torch tensor (B,3,3)
   '''
   bs = depths.shape[0]
-  invalid_mask = (depths<0.001) | (depths>zfar)
+  invalid_mask = (depths<0.1) | (depths>zfar)
   H,W = depths.shape[-2:]
   vs,us = torch.meshgrid(torch.arange(0,H),torch.arange(0,W), indexing='ij')
   vs = vs.reshape(-1).float().cuda()[None].expand(bs,-1)
@@ -460,7 +460,7 @@ def depth_to_vis(depth, zmin=None, zmax=None, mode='rgb', inverse=True):
     zmax = depth.max()
 
   if inverse:
-    invalid = depth<0.001
+    invalid = depth<0.1
     vis = zmin/(depth+1e-8)
     vis[invalid] = 0
   else:

build_all.sh → FoundationPose_ROS/build_all.sh

@@ -1,7 +1,7 @@
 DIR=$(pwd)
 
-cd $DIR/mycpp/ && mkdir -p build && cd build && cmake .. -DPYTHON_EXECUTABLE=$(which python) && make -j11
+cd $DIR/mycpp/ && mkdir -p build && cd build && cmake .. && make -j11
 cd /kaolin && rm -rf build *egg* && pip install -e .
 cd $DIR/bundlesdf/mycuda && rm -rf build *egg* && pip install -e .
 
-cd ${DIR}
+cd ${DIR}

FoundationPose_ROS/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/.ninja_log

@@ -0,0 +1,5 @@
+# ninja log v5
+0	19355	1717606098904099100	/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/bindings.o	c45a0e0392988c7c
+0	43728	1717606123284114914	/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/common.o	609ffa27f9e26a4e
+2	19532	1717606143228126221	/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/bindings.o	1b38b18bc6c45b52
+2	93048	1717606216744158515	/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/gridencoder.o	848beea3ff82f97f

FoundationPose_ROS/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/build.ninja

@@ -0,0 +1,34 @@
+ninja_required_version = 1.3
+cxx = c++
+nvcc = /usr/local/cuda/bin/nvcc
+
+cflags = -pthread -B /opt/conda/envs/my/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include/TH -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/local/include/eigen3 -I/usr/include/eigen3 -I/opt/conda/envs/my/include/python3.8 -I/opt/conda/envs/my/include/python3.8 -c
+post_cflags = -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=gridencoder -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++17
+cuda_cflags = -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include/TH -I/opt/conda/envs/my/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/local/include/eigen3 -I/usr/include/eigen3 -I/opt/conda/envs/my/include/python3.8 -I/opt/conda/envs/my/include/python3.8 -c
+cuda_post_cflags = -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''"'"'-fPIC'"'"'' -Xcompiler -O3 -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=gridencoder -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_75,code=compute_75 -gencode=arch=compute_75,code=sm_75
+cuda_dlink_post_cflags = 
+ldflags = 
+
+rule compile
+  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags
+  depfile = $out.d
+  deps = gcc
+
+rule cuda_compile
+  depfile = $out.d
+  deps = gcc
+  command = $nvcc  $cuda_cflags -c $in -o $out $cuda_post_cflags
+
+
+
+
+
+build /home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/bindings.o: compile /home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/bindings.cpp
+build /home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/build/temp.linux-x86_64-cpython-38/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/gridencoder.o: cuda_compile /home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/gridencoder.cu
+
+
+
+
+
+
+

FoundationPose_ROS/bundlesdf/mycuda/common.egg-info/PKG-INFO

@@ -0,0 +1,3 @@
+Metadata-Version: 2.1
+Name: common
+Version: 0.0.0

FoundationPose_ROS/bundlesdf/mycuda/common.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+bindings.cpp
+common.cu
+setup.py
+/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/bindings.cpp
+/home/shubhodeep/Figueroa_Lab/Foundation_Pose/FoundationPose/bundlesdf/mycuda/torch_ngp_grid_encoder/gridencoder.cu
+common.egg-info/PKG-INFO
+common.egg-info/SOURCES.txt
+common.egg-info/dependency_links.txt
+common.egg-info/top_level.txt

FoundationPose_ROS/bundlesdf/mycuda/common.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

FoundationPose_ROS/bundlesdf/mycuda/common.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+common
+gridencoder

bundlesdf/nerf_helpers.py → FoundationPose_ROS/bundlesdf/nerf_helpers.py

@@ -257,7 +257,7 @@ def preprocess_data(rgbs,depths,masks,normal_maps,poses,sc_factor,translation):
     @normal_maps: (N,H,W,3)
     @poses: (N,4,4)
     '''
-    depths[depths<0.001] = BAD_DEPTH
+    depths[depths<0.1] = BAD_DEPTH
     if masks is not None:
         rgbs[masks==0] = BAD_COLOR
         depths[masks==0] = BAD_DEPTH

FoundationPose_ROS/docker/run_container_new_env.sh

@@ -0,0 +1,12 @@
+docker rm -f ros_fp_new_env
+DIR=$(pwd)/../
+xhost + && docker run --gpus all --env NVIDIA_DISABLE_REQUIRE=1 -it --network=host --name ros_fp_new_env --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
+  -v $DIR:$DIR \
+  -v /home:/home \
+  -v /mnt:/mnt \
+  -v /tmp/.X11-unix:/tmp/.X11-unix \
+  -v /tmp:/tmp \
+  --ipc=host \
+  -e DISPLAY=${DISPLAY} \
+  -e GIT_INDEX_FILE \
+  ghcr.io/shubho-upenn/ros_fp_new_env:latest bash -c "cd $DIR && bash"

estimater.py → FoundationPose_ROS/estimater.py

@@ -141,12 +141,13 @@ def guess_translation(self, depth, mask, K):
       return np.zeros((3))
     uc = (us.min()+us.max())/2.0
     vc = (vs.min()+vs.max())/2.0
-    valid = mask.astype(bool) & (depth>=0.001)
+    valid = mask.astype(bool) & (depth>=0.1)
     if not valid.any():
       logging.info(f"valid is empty")
       return np.zeros((3))
 
     zc = np.median(depth[valid])
+
     center = (np.linalg.inv(K)@np.asarray([uc,vc,1]).reshape(3,1))*zc
 
     if self.debug>=2:
@@ -175,13 +176,14 @@ def register(self, K, rgb, depth, ob_mask, ob_id=None, glctx=None, iteration=5):
 
     if self.debug>=2:
       xyz_map = depth2xyzmap(depth, K)
-      valid = xyz_map[...,2]>=0.001
+      valid = xyz_map[...,2]>=0.1
       pcd = toOpen3dCloud(xyz_map[valid], rgb[valid])
       o3d.io.write_point_cloud(f'{self.debug_dir}/scene_raw.ply',pcd)
       cv2.imwrite(f'{self.debug_dir}/ob_mask.png', (ob_mask*255.0).clip(0,255))
 
     normal_map = None
-    valid = (depth>=0.001) & (ob_mask>0)
+
+    valid = (depth>=0.1) & (ob_mask>0)
     if valid.sum()<4:
       logging.info(f'valid too small, return')
       pose = np.eye(4)
@@ -191,7 +193,7 @@ def register(self, K, rgb, depth, ob_mask, ob_id=None, glctx=None, iteration=5):
     if self.debug>=2:
       imageio.imwrite(f'{self.debug_dir}/color.png', rgb)
       cv2.imwrite(f'{self.debug_dir}/depth.png', (depth*1000).astype(np.uint16))
-      valid = xyz_map[...,2]>=0.001
+      valid = xyz_map[...,2]>=0.1
       pcd = toOpen3dCloud(xyz_map[valid], rgb[valid])
       o3d.io.write_point_cloud(f'{self.debug_dir}/scene_complete.ply',pcd)
 
@@ -252,11 +254,15 @@ def track_one(self, rgb, depth, K, iteration, extra={}):
       logging.info("Please init pose by register first")
       raise RuntimeError
     logging.info("Welcome")
-
+    # depth = depth.astype(np.int16)
     depth = torch.as_tensor(depth, device='cuda', dtype=torch.float)
     depth = erode_depth(depth, radius=2, device='cuda')
+
     depth = bilateral_filter_depth(depth, radius=2, device='cuda')
-    logging.info("depth processing done")
+    # cv2.imshow('3', depth_to_vis(depth.detach().cpu().numpy().reshape(480,640), inverse=False))
+    # cv2.waitKey(1)
+
+    # logging.info("depth processing done")
 
     xyz_map = depth2xyzmap_batch(depth[None], torch.as_tensor(K, dtype=torch.float, device='cuda')[None], zfar=np.inf)[0]
 

learning/datasets/h5_dataset.py → FoundationPose_ROS/learning/datasets/h5_dataset.py

@@ -91,7 +91,7 @@ def transform_depth_to_xyzmap(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
       batch.xyz_mapAs = kornia.geometry.transform.warp_perspective(batch.xyz_mapAs, tf_to_crops, dsize=(H,W), mode='nearest', align_corners=False)
     batch.xyz_mapAs = batch.xyz_mapAs.cuda()
     if self.cfg['normalize_xyz']:
-      invalid = batch.xyz_mapAs[:,2:3]<0.001
+      invalid = batch.xyz_mapAs[:,2:3]<0.1
     batch.xyz_mapAs = batch.xyz_mapAs-batch.poseA[:,:3,3].reshape(bs,3,1,1)
     if self.cfg['normalize_xyz']:
       batch.xyz_mapAs *= 1/mesh_radius.reshape(bs,1,1,1)
@@ -104,7 +104,7 @@ def transform_depth_to_xyzmap(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
       batch.xyz_mapBs = kornia.geometry.transform.warp_perspective(batch.xyz_mapBs, tf_to_crops, dsize=(H,W), mode='nearest', align_corners=False)
     batch.xyz_mapBs = batch.xyz_mapBs.cuda()
     if self.cfg['normalize_xyz']:
-      invalid = batch.xyz_mapBs[:,2:3]<0.001
+      invalid = batch.xyz_mapBs[:,2:3]<0.1
     batch.xyz_mapBs = batch.xyz_mapBs-batch.poseA[:,:3,3].reshape(bs,3,1,1)
     if self.cfg['normalize_xyz']:
       batch.xyz_mapBs *= 1/mesh_radius.reshape(bs,1,1,1)

FoundationPose_ROS/mesh/gelatin_box/poisson/nontextured.xml

@@ -0,0 +1,8 @@
+<KinBody name="009_gelatin_box">
+  <Body type="static" name="009_gelatin_box">
+    <Geom type="trimesh">
+      <Render>./nontextured.stl</Render>
+      <Data>./nontextured.stl</Data>
+    </Geom>
+  </Body>
+</KinBody>

FoundationPose_ROS/mesh/gelatin_box/poisson/textured.mtl

@@ -0,0 +1,9 @@
+newmtl texture_textured.png
+Ns 96.078431
+Ka 0.000000 0.000000 0.000000
+Kd 0.640000 0.640000 0.640000
+Ks 0.500000 0.500000 0.500000
+Ni 1.000000
+d 1.000000
+illum 2
+map_Kd textured.png