FoundationPose 测试记录

FoundationPose

安装

conda环境下搞了一下午，自己电脑的cuda版本太低了，放弃了，就直接装docker环境了。

docker要注意的是，docker要安装GPU版的，否则会报错。 

docker: Error response from daemon: could not select device driver "" with capabilities: [[gpu]].

参考解决方案 
nvidia-container-toolkit

解析

1. 显示模型

为了显示3D模型，需要重新装下 pyglet的版本，

pip install pyglet==1.5.28

然后运行下面的代码。

import trimesh

mesh = trimesh.load('demo_data/mustard0/mesh/textured_simple.obj')

attributes = dir(mesh)
#attributes 中能包含所有类的所有属性和函数
print(attributes)

print(len(mesh.vertices)) #打印顶点数

scene = trimesh.Scene()
# 添加模型到场景
scene.add_geometry(mesh)

# 显示场景
scene.show()

2.  run_demo 用于有模型的6D姿态估计

mesh = trimesh.load(args.mesh_file) #读取mesh

to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
#to_origin是方向和位置，extents是尺寸大小
bbox = np.stack([-extents/2, extents/2], axis=0).reshape(2,3)

...

#初始化四个模块
scorer = ScorePredictor()
refiner = PoseRefinePredictor()
glctx = dr.RasterizeCudaContext()
est = FoundationPose(model_pts=mesh.vertices, model_normals=mesh.vertex_normals, mesh=mesh, scorer=scorer, refiner=refiner, debug_dir=debug_dir, debug=debug, glctx=glctx)

#初始化数据
reader = YcbineoatReader(video_dir=args.test_scene_dir, shorter_side=None, zfar=np.inf)

 FoundationPose函数是主要进行6D 姿态估计的函数，

model_pts=mesh.vertices #是mesh的顶点坐标

model_normals=mesh.vertex_normals #是mesh的顶点法线方向，是3D向量

 mesh=mesh, #载入 3D 模型

scorer=scorer, #应该是对应论文3.4章节的打分预测

refiner=refiner, #应该对应论文3.3章的精细化预测部分

glctx=glctx，#RasterizeCudaContext函数是使用nvdiffrast库中的RasterizeCudaContext类，用于光栅化CUDA上下文（将3D场景转换为2D图像）

-------------------

RasterizeCudaContext 的使用，解析了一下nvdiffrast_render, 光栅化就是将3D模型向不同的2D图像方向投影。

主要使用dr.rasterize 这一步进行具体执行

rast_out, _ = dr.rasterize(glctx, pos_clip, pos_idx, resolution=np.asarray(output_size))

mustard0 的demo为例

glctx 是dr.RasterizeCudaContext()类

pos_clip 维度 [252, 10983, 4] ， 252个方向，顶点数10983个len(mesh.vertices)，4个通道？

pos_idx  维度 [15728, 3], ，15728是网格面的数量len(mesh.faces) ，每行表示形成面的顶点的索引，意思是三个点成一个面

resolution  = [160,160]  , 期望输出的2D图像维度

得到的 rast_out 维度 (252, 160, 160, 4)， 252个光栅化后的方向，每个图160*160， 4个通道（RGB +透明）

代码中进行了252个方向的光栅化投影，可能就是论文中提到的，粗略姿态方向。

以下是其中3个方向

颜色通道暂时没查出来在mesh中是怎么个储存形式，有个这：

mesh.visual.material.image 

np.array(mesh.visual.material.image.convert('RGB'))

如下所示，RGB和BGR反了

------------------

后面就是正常的输入图片和预测姿态，如果是第一帧，会使用物体的mask信息进行register

    if i==0:
      mask = reader.get_mask(0).astype(bool)
      pose = est.register(K=reader.K, rgb=color, depth=depth, ob_mask=mask, iteration=args.est_refine_iter)

如果不是第一帧，会使用track_one ,根据前面的估计结果进行调整

pose = est.track_one(rgb=color, depth=depth, K=reader.K, iteration=args.track_refine_iter)