【6D位姿估计】FoundationPose 跑通demo 训练记录

前言

本文记录在FoundationPose中，跑通基于CAD模型为输入的demo，输出位姿信息，可视化结果。

然后分享NeRF物体重建部分的训练，以及RGBD图为输入的demo。

1、搭建环境

方案1：基于docker镜像（推荐）

首先下载开源代码：https://github.com/NVlabs/FoundationPose

然后执行下面命令，拉取镜像，并构建镜像环境

  1. cd docker/
  2. docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose
  3. bash docker/run_container.sh
  4. bash build_all.sh

构建完成后，可以用docker exec 进入镜像容器中。

方案2：基于Conda（比较麻烦）

首先安装 Eigen3

  1. cd $HOME && wget -q https://gitlab.com/libeigen/eigen/-/archive/3.4.0/eigen-3.4.0.tar.gz && \
  2. tar -xzf eigen-3.4.0.tar.gz && \
  3. cd eigen-3.4.0 && mkdir build && cd build
  4. cmake .. -Wno-dev -DCMAKE_BUILD_TYPE=Release -DCMAKE_CXX_FLAGS=-std=c++14 ..
  5. sudo make install
  6. cd $HOME && rm -rf eigen-3.4.0 eigen-3.4.0.tar.gz

然后参考下面命令，创建conda环境

  1. # create conda environment
  2. create -n foundationpose python=3.9
  3.  
  4. # activate conda environment
  5. conda activate foundationpose
  6.  
  7. # install dependencies
  8. python -m pip install -r requirements.txt
  9.  
  10. # Install NVDiffRast
  11. python -m pip install --quiet --no-cache-dir git+https://github.com/NVlabs/nvdiffrast.git
  12.  
  13. # Kaolin (Optional, needed if running model-free setup)
  14. python -m pip install --quiet --no-cache-dir kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.0.0_cu118.html
  15.  
  16. # PyTorch3D
  17. python -m pip install --quiet --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu118_pyt200/download.html
  18.  
  19. # Build extensions
  20. CMAKE_PREFIX_PATH=$CONDA_PREFIX/lib/python3.9/site-packages/pybind11/share/cmake/pybind11 bash build_all_conda.sh

2、基于CAD模型为输入的demo

首先，下载模型权重，里面包括两个文件夹；点击下载（模型权重）

在工程目录中创建weights/，将下面两个文件夹放到里面。

然后，下载测试数据，里面包括两个压缩文件；点击下载（demo数据）

在工程目录中创建demo_data/，加压文件，将下面两个文件放到里面。

 运行 run_demo.py，实现CAD模型为输入的demo

python run_demo.py --debug 2

如果是服务器运行，没有可视化的，需要注释两行代码：

    if debug>=1:
      center_pose = [email protected](to_origin)
      vis = draw_posed_3d_box(reader.K, img=color, ob_in_cam=center_pose, bbox=bbox)
      vis = draw_xyz_axis(color, ob_in_cam=center_pose, scale=0.1, K=reader.K, thickness=3, transparency=0, is_input_rgb=True)
      # cv2.imshow('1', vis[...,::-1])
      # cv2.waitKey(1)

然后看到demo_data/mustard0/，里面生成了ob_in_cam、track_vis文件夹

ob_in_cam 是位姿估计的结果，用txt文件存储，示例文件：

  1. 6.073544621467590332e-01 -2.560715079307556152e-01 7.520291209220886230e-01 -4.481770694255828857e-01
  2. -7.755840420722961426e-01 -3.960975110530853271e-01 4.915038347244262695e-01 1.187708452343940735e-01
  3. 1.720167100429534912e-01 -8.817789554595947266e-01 -4.391765296459197998e-01 8.016449213027954102e-01
  4. 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00 1.000000000000000000e+00

track_vis 是可视化结果，能看到多张图片：

3、NeRF物体重建训练

下载训练数据，Linemod和YCB-V两个公开数据集的示例：

点击下载（RGBD参考数据）

示例1：训练Linemod数据集

修改代码bundlesdf/run_nerf.py，修改为use_refined_mask=False，即98行：

mesh = run_one_ob(base_dir=base_dir, cfg=cfg, use_refined_mask=False)

 然后执行命令：

python bundlesdf/run_nerf.py --ref_view_dir /DATASET/lm_ref_views --dataset linemod

 如果是服务器运行，没有可视化的，需要安装xvfb

  1. sudo apt-get update
  2. sudo apt-get install -y xvfb

 然后执行命令：

xvfb-run -s "-screen 0 1024x768x24" python bundlesdf/run_nerf.py --ref_view_dir model_free_ref_views/lm_ref_views --dataset linemod

因为训练NeRF需要渲染的，使用xvfb进行模拟。

能看到打印信息：

  1. bundlesdf/run_nerf.py:61: DeprecationWarning: Starting with ImageIO v3 the behavior of this function will switch to that of iio.v3.imread. To keep the current behavior (and make this warning disappear) use `import imageio.v2 as imageio` or call `imageio.v2.imread` directly.
  2.   rgb = imageio.imread(color_file)
  3. [compute_scene_bounds()] compute_scene_bounds_worker start
  4. [compute_scene_bounds()] compute_scene_bounds_worker done
  5. [compute_scene_bounds()] merge pcd
  6. [compute_scene_bounds()] compute_translation_scales done
  7. translation_cvcam=[0.00024226 0.00356217 0.00056694], sc_factor=19.274929219577043
  8. [build_octree()] Octree voxel dilate_radius:1
  9. [__init__()] level:0, vox_pts:torch.Size([1, 3]), corner_pts:torch.Size([8, 3])
  10. [__init__()] level:1, vox_pts:torch.Size([8, 3]), corner_pts:torch.Size([27, 3])
  11. [__init__()] level:2, vox_pts:torch.Size([64, 3]), corner_pts:torch.Size([125, 3])
  12. [draw()] level:2
  13. [draw()] level:2
  14. level 0, resolution: 32
  15. level 1, resolution: 37
  16. level 2, resolution: 43
  17. level 3, resolution: 49
  18. level 4, resolution: 56
  19. level 5, resolution: 64
  20. level 6, resolution: 74
  21. level 7, resolution: 85
  22. level 8, resolution: 98
  23. level 9, resolution: 112
  24. level 10, resolution: 128
  25. level 11, resolution: 148
  26. level 12, resolution: 169
  27. level 13, resolution: 195
  28. level 14, resolution: 223
  29. level 15, resolution: 256
  30. GridEncoder: input_dim=3 n_levels=16 level_dim=2 resolution=32 -> 256 per_level_scale=1.1487 params=(26463840, 2) gridtype=hash align_corners=False
  31. sc_factor 19.274929219577043
  32. translation [0.00024226 0.00356217 0.00056694]
  33. [__init__()] denoise cloud
  34. [__init__()] Denoising rays based on octree cloud
  35. [__init__()] bad_mask#=3
  36. rays torch.Size([128387, 12])
  37. [train()] train progress 0/1001
  38. [train_loop()] Iter: 0, valid_samples: 524161/524288, valid_rays: 2048/2048, loss: 309.0942383, rgb_loss: 0.0216732, rgb0_loss: 0.0000000, fs_rgb_loss: 0.0000000, depth_loss: 0.0000000, depth_loss0: 0.0000000, fs_loss: 301.6735840, point_cloud_loss: 0.0000000, point_cloud_normal_loss: 0.0000000, sdf_loss: 7.2143111, eikonal_loss: 0.0000000, variation_loss: 0.0000000, truncation(meter): 0.0100000, pose_reg: 0.0000000, reg_features: 0.1152707,
  39.  
  40. [train()] train progress 100/1001
  41. [train()] train progress 200/1001
  42. [train()] train progress 300/1001
  43. [train()] train progress 400/1001
  44. [train()] train progress 500/1001
  45. Saved checkpoints at model_free_ref_views/lm_ref_views/ob_0000001/nerf/model_latest.pth
  46. [train_loop()] Iter: 500, valid_samples: 518554/524288, valid_rays: 2026/2048, loss: 1.0530750, rgb_loss: 0.0009063, rgb0_loss: 0.0000000, fs_rgb_loss: 0.0000000, depth_loss: 0.0000000, depth_loss0: 0.0000000, fs_loss: 0.2142579, point_cloud_loss: 0.0000000, point_cloud_normal_loss: 0.0000000, sdf_loss: 0.8360301, eikonal_loss: 0.0000000, variation_loss: 0.0000000, truncation(meter): 0.0100000, pose_reg: 0.0000000, reg_features: 0.0008409,
  47.  
  48. [extract_mesh()] query_pts:torch.Size([42875, 3]), valid:42875
  49. [extract_mesh()] Running Marching Cubes
  50. [extract_mesh()] done V:(4536, 3), F:(8986, 3)
  51. [train()] train progress 600/1001
  52. [train()] train progress 700/1001
  53. [train()] train progress 800/1001
  54. [train()] train progress 900/1001
  55. [train()] train progress 1000/1001
  56. Saved checkpoints at model_free_ref_views/lm_ref_views/ob_0000001/nerf/model_latest.pth
  57. [train_loop()] Iter: 1000, valid_samples: 519351/524288, valid_rays: 2029/2048, loss: 0.4827633, rgb_loss: 0.0006563, rgb0_loss: 0.0000000, fs_rgb_loss: 0.0000000, depth_loss: 0.0000000, depth_loss0: 0.0000000, fs_loss: 0.0935674, point_cloud_loss: 0.0000000, point_cloud_normal_loss: 0.0000000, sdf_loss: 0.3876466, eikonal_loss: 0.0000000, variation_loss: 0.0000000, truncation(meter): 0.0100000, pose_reg: 0.0000000, reg_features: 0.0001022,
  58.  
  59. [extract_mesh()] query_pts:torch.Size([42875, 3]), valid:42875
  60. [extract_mesh()] Running Marching Cubes
  61. [extract_mesh()] done V:(5265, 3), F:(10328, 3)
  62. [extract_mesh()] query_pts:torch.Size([42875, 3]), valid:42875
  63. [extract_mesh()] Running Marching Cubes
  64. [extract_mesh()] done V:(5265, 3), F:(10328, 3)
  65. [()] OpenGL_accelerate module loaded
  66. [()] Using accelerated ArrayDatatype
  67. [mesh_texture_from_train_images()] Texture: Texture map computation
  68. project train_images 0/16
  69. project train_images 1/16
  70. project train_images 2/16
  71. project train_images 3/16
  72. project train_images 4/16
  73. project train_images 5/16
  74. project train_images 6/16
  75. project train_images 7/16
  76. project train_images 8/16
  77. project train_images 9/16
  78. project train_images 10/16
  79. project train_images 11/16
  80. project train_images 12/16
  81. project train_images 13/16
  82. project train_images 14/16
  83. project train_images 15/16

重点留意，损失的变化：

[train()] train progress 0/1001
[train_loop()] Iter: 0, valid_samples: 524161/524288, valid_rays: 2048/2048, loss: 309.0942383, rgb_loss: 0.0216732, rgb0_loss: 0.0000000, fs_rgb_loss: 0.0000000, depth_loss: 0.0000000, depth_loss0: 0.0000000, fs_loss: 301.6735840, point_cloud_loss: 0.0000000, point_cloud_normal_loss: 0.0000000, sdf_loss: 7.2143111, eikonal_loss: 0.0000000, variation_loss: 0.0000000, truncation(meter): 0.0100000, pose_reg: 0.0000000, reg_features: 0.1152707,

[train()] train progress 100/1001
[train()] train progress 200/1001
[train()] train progress 300/1001
[train()] train progress 400/1001
[train()] train progress 500/1001
Saved checkpoints at model_free_ref_views/lm_ref_views/ob_0000001/nerf/model_latest.pth
[train_loop()] Iter: 500, valid_samples: 518554/524288, valid_rays: 2026/2048, loss: 1.0530750, rgb_loss: 0.0009063, rgb0_loss: 0.0000000, fs_rgb_loss: 0.0000000, depth_loss: 0.0000000, depth_loss0: 0.0000000, fs_loss: 0.2142579, point_cloud_loss: 0.0000000, point_cloud_normal_loss: 0.0000000, sdf_loss: 0.8360301, eikonal_loss: 0.0000000, variation_loss: 0.0000000, truncation(meter): 0.0100000, pose_reg: 0.0000000, reg_features: 0.0008409,

默认训练1000轮，训练也挺快的。

在lm_ref_views/ob_0000001/中生成了nerf文件夹，存放下面文件：

在lm_ref_views/ob_0000001/model中，生成了的model.obj，后续模型推理或demo直接使用它。

示例2：训练YCB-V数据集

python bundlesdf/run_nerf.py --ref_view_dir /DATASET/ycbv/ref_views_16 --dataset ycbv

如果是服务器运行，没有可视化的，需要安装xvfb

  1. sudo apt-get update
  2. sudo apt-get install -y xvfb

 然后执行命令：

xvfb-run -s "-screen 0 1024x768x24" python bundlesdf/run_nerf.py --ref_view_dir /DATASET/ycbv/ref_views_16 --dataset ycbv

因为训练NeRF需要渲染的，使用xvfb进行模拟。

4、RGBD图输入demo

这里以Linemod数据集为示例，首先下面测试数据集

点击下载（测试数据）

然后加压文件，存放路径：FoundationPose-main/model_free_ref_views/lm_test_all

官方代码有问题，需要替换两个代码：datareader.py、run_linemod.py

run_linemod.py

  1. # Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
  2. #
  3. # NVIDIA CORPORATION and its licensors retain all intellectual property
  4. # and proprietary rights in and to this software, related documentation
  5. # and any modifications thereto.  Any use, reproduction, disclosure or
  6. # distribution of this software and related documentation without an express
  7. # license agreement from NVIDIA CORPORATION is strictly prohibited.
  8.  
  9.  
  10. from Utils import *
  11. import json,uuid,joblib,os,sys
  12. import scipy.spatial as spatial
  13. from multiprocessing import Pool
  14. import multiprocessing
  15. from functools import partial
  16. from itertools import repeat
  17. import itertools
  18. from datareader import *
  19. from estimater import *
  20. code_dir = os.path.dirname(os.path.realpath(__file__))
  21. sys.path.append(f'{code_dir}/mycpp/build')
  22. import yaml
  23. import re
  24.  
  25.  
  26. def get_mask(reader, i_frame, ob_id, detect_type):
  27.   if detect_type=='box':
  28.     mask = reader.get_mask(i_frame, ob_id)
  29.     H,W = mask.shape[:2]
  30.     vs,us = np.where(mask>0)
  31.     umin = us.min()
  32.     umax = us.max()
  33.     vmin = vs.min()
  34.     vmax = vs.max()
  35.     valid = np.zeros((H,W), dtype=bool)
  36.     valid[vmin:vmax,umin:umax] = 1
  37.   elif detect_type=='mask':
  38.     mask = reader.get_mask(i_frame, ob_id)
  39.     if mask is None:
  40.       return None
  41.     valid = mask>0
  42.   elif detect_type=='detected':
  43.     mask = cv2.imread(reader.color_files[i_frame].replace('rgb','mask_cosypose'), -1)
  44.     valid = mask==ob_id
  45.   else:
  46.     raise RuntimeError
  47.   return valid
  48.  
  49.  
  50.  
  51. def run_pose_estimation_worker(reader, i_frames, est:FoundationPose=None, debug=0, ob_id=None, device='cuda:0'):
  52.   torch.cuda.set_device(device)
  53.   est.to_device(device)
  54.   est.glctx = dr.RasterizeCudaContext(device=device)
  55.  
  56.   result = NestDict()
  57.  
  58.   for i, i_frame in enumerate(i_frames):
  59.     logging.info(f"{i}/{len(i_frames)}, i_frame:{i_frame}, ob_id:{ob_id}")
  60.     print("\n### ", f"{i}/{len(i_frames)}, i_frame:{i_frame}, ob_id:{ob_id}")
  61.     video_id = reader.get_video_id()
  62.     color = reader.get_color(i_frame)
  63.     depth = reader.get_depth(i_frame)
  64.     id_str = reader.id_strs[i_frame]
  65.     H,W = color.shape[:2]
  66.  
  67.     debug_dir =est.debug_dir
  68.  
  69.     ob_mask = get_mask(reader, i_frame, ob_id, detect_type=detect_type)
  70.     if ob_mask is None:
  71.       logging.info("ob_mask not found, skip")
  72.       result[video_id][id_str][ob_id] = np.eye(4)
  73.       return result
  74.  
  75.     est.gt_pose = reader.get_gt_pose(i_frame, ob_id)
  76.  
  77.     pose = est.register(K=reader.K, rgb=color, depth=depth, ob_mask=ob_mask, ob_id=ob_id)
  78.     logging.info(f"pose:\n{pose}")
  79.  
  80.     if debug>=3:
  81.       m = est.mesh_ori.copy()
  82.       tmp = m.copy()
  83.       tmp.apply_transform(pose)
  84.       tmp.export(f'{debug_dir}/model_tf.obj')
  85.  
  86.     result[video_id][id_str][ob_id] = pose
  87.  
  88.   return result, pose
  89.  
  90.  
  91. def run_pose_estimation():
  92.   wp.force_load(device='cuda')
  93.   reader_tmp = LinemodReader(opt.linemod_dir, split=None)
  94.   print("## opt.linemod_dir:", opt.linemod_dir)
  95.  
  96.   debug = opt.debug
  97.   use_reconstructed_mesh = opt.use_reconstructed_mesh
  98.   debug_dir = opt.debug_dir
  99.  
  100.   res = NestDict()
  101.   glctx = dr.RasterizeCudaContext()
  102.   mesh_tmp = trimesh.primitives.Box(extents=np.ones((3)), transform=np.eye(4)).to_mesh()
  103.   est = FoundationPose(model_pts=mesh_tmp.vertices.copy(), model_normals=mesh_tmp.vertex_normals.copy(), symmetry_tfs=None, mesh=mesh_tmp, scorer=None, refiner=None, glctx=glctx, debug_dir=debug_dir, debug=debug)
  104.  
  105.   # ob_id
  106.   match = re.search(r'\d+$', opt.linemod_dir)
  107.   if match:
  108.       last_number = match.group()
  109.       ob_id = int(last_number)
  110.   else:
  111.       print("No digits found at the end of the string")
  112.       
  113.   # for ob_id in reader_tmp.ob_ids:
  114.   if ob_id:
  115.     if use_reconstructed_mesh:
  116.       print("## ob_id:", ob_id)
  117.       print("## opt.linemod_dir:", opt.linemod_dir)
  118.       print("## opt.ref_view_dir:", opt.ref_view_dir)
  119.       mesh = reader_tmp.get_reconstructed_mesh(ref_view_dir=opt.ref_view_dir)
  120.     else:
  121.       mesh = reader_tmp.get_gt_mesh(ob_id)
  122.     # symmetry_tfs = reader_tmp.symmetry_tfs[ob_id]  # !!!!!!!!!!!!!!!!
  123.  
  124.     args = []
  125.  
  126.     reader = LinemodReader(opt.linemod_dir, split=None)
  127.     video_id = reader.get_video_id()
  128.     # est.reset_object(model_pts=mesh.vertices.copy(), model_normals=mesh.vertex_normals.copy(), symmetry_tfs=symmetry_tfs, mesh=mesh)  # raw
  129.     est.reset_object(model_pts=mesh.vertices.copy(), model_normals=mesh.vertex_normals.copy(), mesh=mesh) # !!!!!!!!!!!!!!!!
  130.  
  131.     print("### len(reader.color_files):", len(reader.color_files))
  132.     for i in range(len(reader.color_files)):
  133.       args.append((reader, [i], est, debug, ob_id, "cuda:0"))
  134.  
  135.     # vis Data
  136.     to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
  137.     bbox = np.stack([-extents/2, extents/2], axis=0).reshape(2,3)
  138.     os.makedirs(f'{opt.linemod_dir}/track_vis', exist_ok=True)
  139.  
  140.     outs = []
  141.     i = 0
  142.     for arg in args[:200]:
  143.       print("### num:", i)
  144.       out, pose = run_pose_estimation_worker(*arg)
  145.       outs.append(out)
  146.       center_pose = [email protected](to_origin)
  147.       img_color = reader.get_color(i)
  148.       vis = draw_posed_3d_box(reader.K, img=img_color, ob_in_cam=center_pose, bbox=bbox)
  149.       vis = draw_xyz_axis(img_color, ob_in_cam=center_pose, scale=0.1, K=reader.K, thickness=3, transparency=0, is_input_rgb=True)
  150.       imageio.imwrite(f'{opt.linemod_dir}/track_vis/{reader.id_strs[i]}.png', vis)
  151.       i = i + 1
  152.  
  153.     for out in outs:
  154.       for video_id in out:
  155.         for id_str in out[video_id]:
  156.           for ob_id in out[video_id][id_str]:
  157.             res[video_id][id_str][ob_id] = out[video_id][id_str][ob_id]
  158.  
  159.   with open(f'{opt.debug_dir}/linemod_res.yml','w') as ff:
  160.     yaml.safe_dump(make_yaml_dumpable(res), ff)
  161.     print("Save linemod_res.yml OK !!!")
  162.  
  163.  
  164. if __name__=='__main__':
  165.   parser = argparse.ArgumentParser()
  166.   code_dir = os.path.dirname(os.path.realpath(__file__))
  167.   parser.add_argument('--linemod_dir', type=str, default="/guopu/FoundationPose-main/model_free_ref_views/lm_test_all/000015", help="linemod root dir") # lm_test_all  lm_test
  168.   parser.add_argument('--use_reconstructed_mesh', type=int, default=1)
  169.   parser.add_argument('--ref_view_dir', type=str, default="/guopu/FoundationPose-main/model_free_ref_views/lm_ref_views/ob_0000015")
  170.   parser.add_argument('--debug', type=int, default=0)
  171.   parser.add_argument('--debug_dir', type=str, default=f'/guopu/FoundationPose-main/model_free_ref_views/lm_test_all/debug') # lm_test_all  lm_test
  172.   opt = parser.parse_args()
  173.   set_seed(0)
  174.  
  175.   detect_type = 'mask'   # mask / box / detected
  176.   run_pose_estimation()
  177.

datareader.py

  1. # Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
  2. #
  3. # NVIDIA CORPORATION and its licensors retain all intellectual property
  4. # and proprietary rights in and to this software, related documentation
  5. # and any modifications thereto.  Any use, reproduction, disclosure or
  6. # distribution of this software and related documentation without an express
  7. # license agreement from NVIDIA CORPORATION is strictly prohibited.
  8.  
  9.  
  10. from Utils import *
  11. import json,os,sys
  12.  
  13.  
  14. BOP_LIST = ['lmo','tless','ycbv','hb','tudl','icbin','itodd']
  15. BOP_DIR = os.getenv('BOP_DIR')
  16.  
  17. def get_bop_reader(video_dir, zfar=np.inf):
  18.   if 'ycbv' in video_dir or 'YCB' in video_dir:
  19.     return YcbVideoReader(video_dir, zfar=zfar)
  20.   if 'lmo' in video_dir or 'LINEMOD-O' in video_dir:
  21.     return LinemodOcclusionReader(video_dir, zfar=zfar)
  22.   if 'tless' in video_dir or 'TLESS' in video_dir:
  23.     return TlessReader(video_dir, zfar=zfar)
  24.   if 'hb' in video_dir:
  25.     return HomebrewedReader(video_dir, zfar=zfar)
  26.   if 'tudl' in video_dir:
  27.     return TudlReader(video_dir, zfar=zfar)
  28.   if 'icbin' in video_dir:
  29.     return IcbinReader(video_dir, zfar=zfar)
  30.   if 'itodd' in video_dir:
  31.     return ItoddReader(video_dir, zfar=zfar)
  32.   else:
  33.     raise RuntimeError
  34.  
  35.  
  36. def get_bop_video_dirs(dataset):
  37.   if dataset=='ycbv':
  38.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/ycbv/test/*'))
  39.   elif dataset=='lmo':
  40.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/lmo/lmo_test_bop19/test/*'))
  41.   elif dataset=='tless':
  42.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/tless/tless_test_primesense_bop19/test_primesense/*'))
  43.   elif dataset=='hb':
  44.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/hb/hb_test_primesense_bop19/test_primesense/*'))
  45.   elif dataset=='tudl':
  46.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/tudl/tudl_test_bop19/test/*'))
  47.   elif dataset=='icbin':
  48.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/icbin/icbin_test_bop19/test/*'))
  49.   elif dataset=='itodd':
  50.     video_dirs = sorted(glob.glob(f'{BOP_DIR}/itodd/itodd_test_bop19/test/*'))
  51.   else:
  52.     raise RuntimeError
  53.   return video_dirs
  54.  
  55.  
  56.  
  57. class YcbineoatReader:
  58.   def __init__(self,video_dir, downscale=1, shorter_side=None, zfar=np.inf):
  59.     self.video_dir = video_dir
  60.     self.downscale = downscale
  61.     self.zfar = zfar
  62.     self.color_files = sorted(glob.glob(f"{self.video_dir}/rgb/*.png"))
  63.     self.K = np.loadtxt(f'{video_dir}/cam_K.txt').reshape(3,3)
  64.     self.id_strs = []
  65.     for color_file in self.color_files:
  66.       id_str = os.path.basename(color_file).replace('.png','')
  67.       self.id_strs.append(id_str)
  68.     self.H,self.W = cv2.imread(self.color_files[0]).shape[:2]
  69.  
  70.     if shorter_side is not None:
  71.       self.downscale = shorter_side/min(self.H, self.W)
  72.  
  73.     self.H = int(self.H*self.downscale)
  74.     self.W = int(self.W*self.downscale)
  75.     self.K[:2] *= self.downscale
  76.  
  77.     self.gt_pose_files = sorted(glob.glob(f'{self.video_dir}/annotated_poses/*'))
  78.  
  79.     self.videoname_to_object = {
  80.       'bleach0': "021_bleach_cleanser",
  81.       'bleach_hard_00_03_chaitanya': "021_bleach_cleanser",
  82.       'cracker_box_reorient': '003_cracker_box',
  83.       'cracker_box_yalehand0': '003_cracker_box',
  84.       'mustard0': '006_mustard_bottle',
  85.       'mustard_easy_00_02': '006_mustard_bottle',
  86.       'sugar_box1': '004_sugar_box',
  87.       'sugar_box_yalehand0': '004_sugar_box',
  88.       'tomato_soup_can_yalehand0': '005_tomato_soup_can',
  89.     }
  90.  
  91.  
  92.   def get_video_name(self):
  93.     return self.video_dir.split('/')[-1]
  94.  
  95.   def __len__(self):
  96.     return len(self.color_files)
  97.  
  98.   def get_gt_pose(self,i):
  99.     try:
  100.       pose = np.loadtxt(self.gt_pose_files[i]).reshape(4,4)
  101.       return pose
  102.     except:
  103.       logging.info("GT pose not found, return None")
  104.       return None
  105.  
  106.  
  107.   def get_color(self,i):
  108.     color = imageio.imread(self.color_files[i])[...,:3]
  109.     color = cv2.resize(color, (self.W,self.H), interpolation=cv2.INTER_NEAREST)
  110.     return color
  111.  
  112.   def get_mask(self,i):
  113.     mask = cv2.imread(self.color_files[i].replace('rgb','masks'),-1)
  114.     if len(mask.shape)==3:
  115.       for c in range(3):
  116.         if mask[...,c].sum()>0:
  117.           mask = mask[...,c]
  118.           break
  119.     mask = cv2.resize(mask, (self.W,self.H), interpolation=cv2.INTER_NEAREST).astype(bool).astype(np.uint8)
  120.     return mask
  121.  
  122.   def get_depth(self,i):
  123.     depth = cv2.imread(self.color_files[i].replace('rgb','depth'),-1)/1e3
  124.     depth = cv2.resize(depth, (self.W,self.H), interpolation=cv2.INTER_NEAREST)
  125.     depth[(depth<0.1) | (depth>=self.zfar)] = 0
  126.     return depth
  127.  
  128.  
  129.   def get_xyz_map(self,i):
  130.     depth = self.get_depth(i)
  131.     xyz_map = depth2xyzmap(depth, self.K)
  132.     return xyz_map
  133.  
  134.   def get_occ_mask(self,i):
  135.     hand_mask_file = self.color_files[i].replace('rgb','masks_hand')
  136.     occ_mask = np.zeros((self.H,self.W), dtype=bool)
  137.     if os.path.exists(hand_mask_file):
  138.       occ_mask = occ_mask | (cv2.imread(hand_mask_file,-1)>0)
  139.  
  140.     right_hand_mask_file = self.color_files[i].replace('rgb','masks_hand_right')
  141.     if os.path.exists(right_hand_mask_file):
  142.       occ_mask = occ_mask | (cv2.imread(right_hand_mask_file,-1)>0)
  143.  
  144.     occ_mask = cv2.resize(occ_mask, (self.W,self.H), interpolation=cv2.INTER_NEAREST)
  145.  
  146.     return occ_mask.astype(np.uint8)
  147.  
  148.   def get_gt_mesh(self):
  149.     ob_name = self.videoname_to_object[self.get_video_name()]
  150.     YCB_VIDEO_DIR = os.getenv('YCB_VIDEO_DIR')
  151.     mesh = trimesh.load(f'{YCB_VIDEO_DIR}/models/{ob_name}/textured_simple.obj')
  152.     return mesh
  153.  
  154.  
  155. class BopBaseReader:
  156.   def __init__(self, base_dir, zfar=np.inf, resize=1):
  157.     self.base_dir = base_dir
  158.     self.resize = resize
  159.     self.dataset_name = None
  160.     self.color_files = sorted(glob.glob(f"{self.base_dir}/rgb/*"))
  161.     if len(self.color_files)==0:
  162.       self.color_files = sorted(glob.glob(f"{self.base_dir}/gray/*"))
  163.     self.zfar = zfar
  164.  
  165.     self.K_table = {}
  166.     with open(f'{self.base_dir}/scene_camera.json','r') as ff:
  167.       info = json.load(ff)
  168.     for k in info:
  169.       self.K_table[f'{int(k):06d}'] = np.array(info[k]['cam_K']).reshape(3,3)
  170.       self.bop_depth_scale = info[k]['depth_scale']
  171.  
  172.     if os.path.exists(f'{self.base_dir}/scene_gt.json'):
  173.       with open(f'{self.base_dir}/scene_gt.json','r') as ff:
  174.         self.scene_gt = json.load(ff)
  175.       self.scene_gt = copy.deepcopy(self.scene_gt)   # Release file handle to be pickle-able by joblib
  176.       assert len(self.scene_gt)==len(self.color_files)
  177.     else:
  178.       self.scene_gt = None
  179.  
  180.     self.make_id_strs()
  181.  
  182.  
  183.   def make_scene_ob_ids_dict(self):
  184.     with open(f'{BOP_DIR}/{self.dataset_name}/test_targets_bop19.json','r') as ff:
  185.       self.scene_ob_ids_dict = {}
  186.       data = json.load(ff)
  187.       for d in data:
  188.         if d['scene_id']==self.get_video_id():
  189.           id_str = f"{d['im_id']:06d}"
  190.           if id_str not in self.scene_ob_ids_dict:
  191.             self.scene_ob_ids_dict[id_str] = []
  192.           self.scene_ob_ids_dict[id_str] += [d['obj_id']]*d['inst_count']
  193.  
  194.  
  195.   def get_K(self, i_frame):
  196.     K = self.K_table[self.id_strs[i_frame]]
  197.     if self.resize!=1:
  198.       K[:2,:2] *= self.resize
  199.     return K
  200.  
  201.  
  202.   def get_video_dir(self):
  203.     video_id = int(self.base_dir.rstrip('/').split('/')[-1])
  204.     return video_id
  205.  
  206.   def make_id_strs(self):
  207.     self.id_strs = []
  208.     for i in range(len(self.color_files)):
  209.       name = os.path.basename(self.color_files[i]).split('.')[0]
  210.       self.id_strs.append(name)
  211.  
  212.  
  213.   def get_instance_ids_in_image(self, i_frame:int):
  214.     ob_ids = []
  215.     if self.scene_gt is not None:
  216.       name = int(os.path.basename(self.color_files[i_frame]).split('.')[0])
  217.       for k in self.scene_gt[str(name)]:
  218.         ob_ids.append(k['obj_id'])
  219.     elif self.scene_ob_ids_dict is not None:
  220.       return np.array(self.scene_ob_ids_dict[self.id_strs[i_frame]])
  221.     else:
  222.       mask_dir = os.path.dirname(self.color_files[0]).replace('rgb','mask_visib')
  223.       id_str = self.id_strs[i_frame]
  224.       mask_files = sorted(glob.glob(f'{mask_dir}/{id_str}_*.png'))
  225.       ob_ids = []
  226.       for mask_file in mask_files:
  227.         ob_id = int(os.path.basename(mask_file).split('.')[0].split('_')[1])
  228.         ob_ids.append(ob_id)
  229.     ob_ids = np.asarray(ob_ids)
  230.     return ob_ids
  231.  
  232.  
  233.   def get_gt_mesh_file(self, ob_id):
  234.     raise RuntimeError("You should override this")
  235.  
  236.  
  237.   def get_color(self,i):
  238.     color = imageio.imread(self.color_files[i])
  239.     if len(color.shape)==2:
  240.       color = np.tile(color[...,None], (1,1,3))  # Gray to RGB
  241.     if self.resize!=1:
  242.       color = cv2.resize(color, fx=self.resize, fy=self.resize, dsize=None)
  243.     return color
  244.  
  245.  
  246.   def get_depth(self,i, filled=False):
  247.     if filled:
  248.       depth_file = self.color_files[i].replace('rgb','depth_filled')
  249.       depth_file = f'{os.path.dirname(depth_file)}/0{os.path.basename(depth_file)}'
  250.       depth = cv2.imread(depth_file,-1)/1e3
  251.     else:
  252.       depth_file = self.color_files[i].replace('rgb','depth').replace('gray','depth')
  253.       depth = cv2.imread(depth_file,-1)*1e-3*self.bop_depth_scale
  254.     if self.resize!=1:
  255.       depth = cv2.resize(depth, fx=self.resize, fy=self.resize, dsize=None, interpolation=cv2.INTER_NEAREST)
  256.     depth[depth<0.1] = 0
  257.     depth[depth>self.zfar] = 0
  258.     return depth
  259.  
  260.   def get_xyz_map(self,i):
  261.     depth = self.get_depth(i)
  262.     xyz_map = depth2xyzmap(depth, self.get_K(i))
  263.     return xyz_map
  264.  
  265.  
  266.   def get_mask(self, i_frame:int, ob_id:int, type='mask_visib'):
  267.     '''
  268.     @type: mask_visib (only visible part) / mask (projected mask from whole model)
  269.     '''
  270.     pos = 0
  271.     name = int(os.path.basename(self.color_files[i_frame]).split('.')[0])
  272.     if self.scene_gt is not None:
  273.       for k in self.scene_gt[str(name)]:
  274.         if k['obj_id']==ob_id:
  275.           break
  276.         pos += 1
  277.       mask_file = f'{self.base_dir}/{type}/{name:06d}_{pos:06d}.png'
  278.       if not os.path.exists(mask_file):
  279.         logging.info(f'{mask_file} not found')
  280.         return None
  281.     else:
  282.       # mask_dir = os.path.dirname(self.color_files[0]).replace('rgb',type)
  283.       # mask_file = f'{mask_dir}/{self.id_strs[i_frame]}_{ob_id:06d}.png'
  284.       raise RuntimeError
  285.     mask = cv2.imread(mask_file, -1)
  286.     if self.resize!=1:
  287.       mask = cv2.resize(mask, fx=self.resize, fy=self.resize, dsize=None, interpolation=cv2.INTER_NEAREST)
  288.     return mask>0
  289.  
  290.  
  291.   def get_gt_mesh(self, ob_id:int):
  292.     mesh_file = self.get_gt_mesh_file(ob_id)
  293.     mesh = trimesh.load(mesh_file)
  294.     mesh.vertices *= 1e-3
  295.     return mesh
  296.  
  297.  
  298.   def get_model_diameter(self, ob_id):
  299.     dir = os.path.dirname(self.get_gt_mesh_file(self.ob_ids[0]))
  300.     info_file = f'{dir}/models_info.json'
  301.     with open(info_file,'r') as ff:
  302.       info = json.load(ff)
  303.     return info[str(ob_id)]['diameter']/1e3
  304.  
  305.  
  306.  
  307.   def get_gt_poses(self, i_frame, ob_id):
  308.     gt_poses = []
  309.     name = int(self.id_strs[i_frame])
  310.     for i_k, k in enumerate(self.scene_gt[str(name)]):
  311.       if k['obj_id']==ob_id:
  312.         cur = np.eye(4)
  313.         cur[:3,:3] = np.array(k['cam_R_m2c']).reshape(3,3)
  314.         cur[:3,3] = np.array(k['cam_t_m2c'])/1e3
  315.         gt_poses.append(cur)
  316.     return np.asarray(gt_poses).reshape(-1,4,4)
  317.  
  318.  
  319.   def get_gt_pose(self, i_frame:int, ob_id, mask=None, use_my_correction=False):
  320.     ob_in_cam = np.eye(4)
  321.     best_iou = -np.inf
  322.     best_gt_mask = None
  323.     name = int(self.id_strs[i_frame])
  324.     for i_k, k in enumerate(self.scene_gt[str(name)]):
  325.       if k['obj_id']==ob_id:
  326.         cur = np.eye(4)
  327.         cur[:3,:3] = np.array(k['cam_R_m2c']).reshape(3,3)
  328.         cur[:3,3] = np.array(k['cam_t_m2c'])/1e3
  329.         if mask is not None:  # When multi-instance exists, use mask to determine which one
  330.           gt_mask = cv2.imread(f'{self.base_dir}/mask_visib/{self.id_strs[i_frame]}_{i_k:06d}.png', -1).astype(bool)
  331.           intersect = (gt_mask*mask).astype(bool)
  332.           union = (gt_mask+mask).astype(bool)
  333.           iou = float(intersect.sum())/union.sum()
  334.           if iou>best_iou:
  335.             best_iou = iou
  336.             best_gt_mask = gt_mask
  337.             ob_in_cam = cur
  338.         else:
  339.           ob_in_cam = cur
  340.           break
  341.  
  342.  
  343.     if use_my_correction:
  344.       if 'ycb' in self.base_dir.lower() and 'train_real' in self.color_files[i_frame]:
  345.         video_id = self.get_video_id()
  346.         if ob_id==1:
  347.           if video_id in [12,13,14,17,24]:
  348.             ob_in_cam = [email protected]_tfs[ob_id][1]
  349.     return ob_in_cam
  350.  
  351.  
  352.   def load_symmetry_tfs(self):
  353.     dir = os.path.dirname(self.get_gt_mesh_file(self.ob_ids[0]))
  354.     info_file = f'{dir}/models_info.json'
  355.     with open(info_file,'r') as ff:
  356.       info = json.load(ff)
  357.     self.symmetry_tfs = {}
  358.     self.symmetry_info_table = {}
  359.     for ob_id in self.ob_ids:
  360.       self.symmetry_info_table[ob_id] = info[str(ob_id)]
  361.       self.symmetry_tfs[ob_id] = symmetry_tfs_from_info(info[str(ob_id)], rot_angle_discrete=5)
  362.     self.geometry_symmetry_info_table = copy.deepcopy(self.symmetry_info_table)
  363.  
  364.  
  365.   def get_video_id(self):
  366.     return int(self.base_dir.split('/')[-1])
  367.  
  368.  
  369. class LinemodOcclusionReader(BopBaseReader):
  370.   def __init__(self,base_dir='/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD-O/lmo_test_all/test/000002', zfar=np.inf):
  371.     super().__init__(base_dir, zfar=zfar)
  372.     self.dataset_name = 'lmo'
  373.     self.K = list(self.K_table.values())[0]
  374.     self.ob_ids = [1,5,6,8,9,10,11,12]
  375.     self.ob_id_to_names = {
  376.       1: 'ape',
  377.       2: 'benchvise',
  378.       3: 'bowl',
  379.       4: 'camera',
  380.       5: 'water_pour',
  381.       6: 'cat',
  382.       7: 'cup',
  383.       8: 'driller',
  384.       9: 'duck',
  385.       10: 'eggbox',
  386.       11: 'glue',
  387.       12: 'holepuncher',
  388.       13: 'iron',
  389.       14: 'lamp',
  390.       15: 'phone',
  391.     }
  392.     # self.load_symmetry_tfs()
  393.  
  394.   def get_gt_mesh_file(self, ob_id):
  395.     mesh_dir = f'{BOP_DIR}/{self.dataset_name}/models/obj_{ob_id:06d}.ply'
  396.     return mesh_dir
  397.  
  398.  
  399.  
  400. class LinemodReader(LinemodOcclusionReader):
  401.   def __init__(self, base_dir='/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD/lm_test_all/test/000001', zfar=np.inf, split=None):
  402.     super().__init__(base_dir, zfar=zfar)
  403.     self.dataset_name = 'lm'
  404.     if split is not None:  # train/test
  405.       print("## split is not None")
  406.       with open(f'/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD/Linemod_preprocessed/data/{self.get_video_id():02d}/{split}.txt','r') as ff:
  407.         lines = ff.read().splitlines()
  408.       self.color_files = []
  409.       for line in lines:
  410.         id = int(line)
  411.         self.color_files.append(f'{self.base_dir}/rgb/{id:06d}.png')
  412.       self.make_id_strs()
  413.  
  414.     self.ob_ids = np.setdiff1d(np.arange(1,16), np.array([7,3])).tolist()  # Exclude bowl and mug
  415.     # self.load_symmetry_tfs()
  416.  
  417.  
  418.   def get_gt_mesh_file(self, ob_id):
  419.     root = self.base_dir
  420.     print(f'{root}/../')
  421.     print(f'{root}/lm_models')
  422.     print(f'{root}/lm_models/models/obj_{ob_id:06d}.ply')
  423.     while 1:
  424.       if os.path.exists(f'{root}/lm_models'):
  425.         mesh_dir = f'{root}/lm_models/models/obj_{ob_id:06d}.ply'
  426.         break
  427.       else:
  428.         root = os.path.abspath(f'{root}/../')
  429.         mesh_dir = f'{root}/lm_models/models/obj_{ob_id:06d}.ply'
  430.         break
  431.     return mesh_dir
  432.  
  433.  
  434.   def get_reconstructed_mesh(self, ref_view_dir):
  435.     mesh = trimesh.load(os.path.abspath(f'{ref_view_dir}/model/model.obj'))
  436.     return mesh
  437.  
  438.  
  439. class YcbVideoReader(BopBaseReader):
  440.   def __init__(self, base_dir, zfar=np.inf):
  441.     super().__init__(base_dir, zfar=zfar)
  442.     self.dataset_name = 'ycbv'
  443.     self.K = list(self.K_table.values())[0]
  444.  
  445.     self.make_id_strs()
  446.  
  447.     self.ob_ids = np.arange(1,22).astype(int).tolist()
  448.     YCB_VIDEO_DIR = os.getenv('YCB_VIDEO_DIR')
  449.     self.ob_id_to_names = {}
  450.     self.name_to_ob_id = {}
  451.     # names = sorted(os.listdir(f'{YCB_VIDEO_DIR}/models/'))
  452.     if os.path.exists(f'{YCB_VIDEO_DIR}/models/'):
  453.         names = sorted(os.listdir(f'{YCB_VIDEO_DIR}/models/'))
  454.         for i,ob_id in enumerate(self.ob_ids):
  455.           self.ob_id_to_names[ob_id] = names[i]
  456.           self.name_to_ob_id[names[i]] = ob_id
  457.     else:
  458.         names = []
  459.          
  460.     if 0:
  461.     # if 'BOP' not in self.base_dir:
  462.       with open(f'{self.base_dir}/../../keyframe.txt','r') as ff:
  463.         self.keyframe_lines = ff.read().splitlines()
  464.  
  465.     # self.load_symmetry_tfs()
  466.     '''for ob_id in self.ob_ids:
  467.       if ob_id in [1,4,6,18]:   # Cylinder
  468.         self.geometry_symmetry_info_table[ob_id] = {
  469.           'symmetries_continuous': [
  470.               {'axis':[0,0,1], 'offset':[0,0,0]},
  471.             ],
  472.           'symmetries_discrete': euler_matrix(0, np.pi, 0).reshape(1,4,4).tolist(),
  473.           }
  474.       elif ob_id in [13]:
  475.         self.geometry_symmetry_info_table[ob_id] = {
  476.           'symmetries_continuous': [
  477.               {'axis':[0,0,1], 'offset':[0,0,0]},
  478.             ],
  479.           }
  480.       elif ob_id in [2,3,9,21]:   # Rectangle box
  481.         tfs = []
  482.         for rz in [0, np.pi]:
  483.           for rx in [0,np.pi]:
  484.             for ry in [0,np.pi]:
  485.               tfs.append(euler_matrix(rx, ry, rz))
  486.         self.geometry_symmetry_info_table[ob_id] = {
  487.           'symmetries_discrete': np.asarray(tfs).reshape(-1,4,4).tolist(),
  488.           }
  489.       else:
  490.         pass'''
  491.  
  492.   def get_gt_mesh_file(self, ob_id):
  493.     if 'BOP' in self.base_dir:
  494.       mesh_file = os.path.abspath(f'{self.base_dir}/../../ycbv_models/models/obj_{ob_id:06d}.ply')
  495.     else:
  496.       mesh_file = f'{self.base_dir}/../../ycbv_models/models/obj_{ob_id:06d}.ply'
  497.     return mesh_file
  498.  
  499.  
  500.   def get_gt_mesh(self, ob_id:int, get_posecnn_version=False):
  501.     if get_posecnn_version:
  502.       YCB_VIDEO_DIR = os.getenv('YCB_VIDEO_DIR')
  503.       mesh = trimesh.load(f'{YCB_VIDEO_DIR}/models/{self.ob_id_to_names[ob_id]}/textured_simple.obj')
  504.       return mesh
  505.     mesh_file = self.get_gt_mesh_file(ob_id)
  506.     mesh = trimesh.load(mesh_file, process=False)
  507.     mesh.vertices *= 1e-3
  508.     tex_file = mesh_file.replace('.ply','.png')
  509.     if os.path.exists(tex_file):
  510.       from PIL import Image
  511.       im = Image.open(tex_file)
  512.       uv = mesh.visual.uv
  513.       material = trimesh.visual.texture.SimpleMaterial(image=im)
  514.       color_visuals = trimesh.visual.TextureVisuals(uv=uv, image=im, material=material)
  515.       mesh.visual = color_visuals
  516.     return mesh
  517.  
  518.  
  519.   def get_reconstructed_mesh(self, ob_id, ref_view_dir):
  520.     mesh = trimesh.load(os.path.abspath(f'{ref_view_dir}/ob_{ob_id:07d}/model/model.obj'))
  521.     return mesh
  522.  
  523.  
  524.   def get_transform_reconstructed_to_gt_model(self, ob_id):
  525.     out = np.eye(4)
  526.     return out
  527.  
  528.  
  529.   def get_visible_cloud(self, ob_id):
  530.     file = os.path.abspath(f'{self.base_dir}/../../models/{self.ob_id_to_names[ob_id]}/visible_cloud.ply')
  531.     pcd = o3d.io.read_point_cloud(file)
  532.     return pcd
  533.  
  534.  
  535.   def is_keyframe(self, i):
  536.     color_file = self.color_files[i]
  537.     video_id = self.get_video_id()
  538.     frame_id = int(os.path.basename(color_file).split('.')[0])
  539.     key = f'{video_id:04d}/{frame_id:06d}'
  540.     return (key in self.keyframe_lines)
  541.  
  542.  
  543.  
  544. class TlessReader(BopBaseReader):
  545.   def __init__(self, base_dir, zfar=np.inf):
  546.     super().__init__(base_dir, zfar=zfar)
  547.     self.dataset_name = 'tless'
  548.  
  549.     self.ob_ids = np.arange(1,31).astype(int).tolist()
  550.     self.load_symmetry_tfs()
  551.  
  552.  
  553.   def get_gt_mesh_file(self, ob_id):
  554.     mesh_file = f'{self.base_dir}/../../../models_cad/obj_{ob_id:06d}.ply'
  555.     return mesh_file
  556.  
  557.  
  558.   def get_gt_mesh(self, ob_id):
  559.     mesh = trimesh.load(self.get_gt_mesh_file(ob_id))
  560.     mesh.vertices *= 1e-3
  561.     mesh = trimesh_add_pure_colored_texture(mesh, color=np.ones((3))*200)
  562.     return mesh
  563.  
  564.  
  565. class HomebrewedReader(BopBaseReader):
  566.   def __init__(self, base_dir, zfar=np.inf):
  567.     super().__init__(base_dir, zfar=zfar)
  568.     self.dataset_name = 'hb'
  569.     self.ob_ids = np.arange(1,34).astype(int).tolist()
  570.     self.load_symmetry_tfs()
  571.     self.make_scene_ob_ids_dict()
  572.  
  573.  
  574.   def get_gt_mesh_file(self, ob_id):
  575.     mesh_file = f'{self.base_dir}/../../../hb_models/models/obj_{ob_id:06d}.ply'
  576.     return mesh_file
  577.  
  578.  
  579.   def get_gt_pose(self, i_frame:int, ob_id, use_my_correction=False):
  580.     logging.info("WARN HomeBrewed doesn't have GT pose")
  581.     return np.eye(4)
  582.  
  583.  
  584.  
  585. class ItoddReader(BopBaseReader):
  586.   def __init__(self, base_dir, zfar=np.inf):
  587.     super().__init__(base_dir, zfar=zfar)
  588.     self.dataset_name = 'itodd'
  589.     self.make_id_strs()
  590.  
  591.     self.ob_ids = np.arange(1,29).astype(int).tolist()
  592.     self.load_symmetry_tfs()
  593.     self.make_scene_ob_ids_dict()
  594.  
  595.  
  596.   def get_gt_mesh_file(self, ob_id):
  597.     mesh_file = f'{self.base_dir}/../../../itodd_models/models/obj_{ob_id:06d}.ply'
  598.     return mesh_file
  599.  
  600.  
  601. class IcbinReader(BopBaseReader):
  602.   def __init__(self, base_dir, zfar=np.inf):
  603.     super().__init__(base_dir, zfar=zfar)
  604.     self.dataset_name = 'icbin'
  605.     self.ob_ids = np.arange(1,3).astype(int).tolist()
  606.     self.load_symmetry_tfs()
  607.  
  608.   def get_gt_mesh_file(self, ob_id):
  609.     mesh_file = f'{self.base_dir}/../../../icbin_models/models/obj_{ob_id:06d}.ply'
  610.     return mesh_file
  611.  
  612.  
  613. class TudlReader(BopBaseReader):
  614.   def __init__(self, base_dir, zfar=np.inf):
  615.     super().__init__(base_dir, zfar=zfar)
  616.     self.dataset_name = 'tudl'
  617.     self.ob_ids = np.arange(1,4).astype(int).tolist()
  618.     self.load_symmetry_tfs()
  619.  
  620.   def get_gt_mesh_file(self, ob_id):
  621.     mesh_file = f'{self.base_dir}/../../../tudl_models/models/obj_{ob_id:06d}.ply'
  622.     return mesh_file
  623.  
  624.

运行run_linemod.py：

python run_linemod.py

能看到文件夹model_free_ref_views/lm_test_all/000015/track_vis/

里面存放可视化结果：

分享完成~

本文先介绍到这里，后面会分享“6D位姿估计”的其它数据集、算法、代码、具体应用示例。

注：本文转载自blog.csdn.net的一颗小树x的文章"https://blog.csdn.net/qq_41204464/article/details/138619210"。版权归原作者所有，此博客不拥有其著作权，亦不承担相应法律责任。如有侵权，请联系我们删除。
复制链接
相关推荐
发表评论

评论记录：

未查询到任何数据！