6D姿态估计网络【Gen6D复现】CUDA11.3+cuDNN8.2.1

摘要：主要记录一下复现Gen6D的大概过程，以及在其中遇到的一些bug，主要参考了博主牙牙要健康的文章。
参考文章：【6D位姿估计】【深度学习】windows10下Gen6D代码Pytorch实现

一.环境搭建

1.下载CUDA和cuDNN

需要先下载完整的CUDA11.3和对应版本的Cudnn8.2.1
参考链接
【深度学习】windows10环境配置详细教程
解压下载好的cudnn包，将bin、include和lib目录里的文件移动到cuda11.3相对应的bin、include和lib目录中，最后配置环境变量Path。

创建虚拟环境，安装所需的包

# 创建虚拟环境
conda create -y -n gen6d python=3.8
# 查看新环境是否安装成功
conda env list
# 激活环境
activate gen6d
# 下载githup源代码到合适文件夹,并cd到代码文件夹内(科学上网)
git clone https://github.com/liuyuan-pal/Gen6D.git
cd Gen6D
# 安装pytorch包
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
# 安装其他依赖包(需要删除关于pytorch的部分)
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt

requirements.txt删除了pytorch>=1.7.1、torchvision、pytorch3d和cudatoolkit部分，其中pytorch3d安装十分复杂，会单独进行安装讲解。

3. 安装Pytorch3d

pytorch安装参考

1.安装依赖库

# 这里可能需要anaconda的默认源安装
conda install -c conda-forge -c fvcore -c iopath -c bottler fvcore iopath

2. CUB安装配置

这里的cuda是物理机器上实际安装的cuda版本(环境变量里配置的cuda)，不是虚拟环境中的cuda版本。
根据自己的cuda 版本选择对应的cub realase版本下载，对应CUB 1.11.0 (CUDA Toolkit 11.3)：
下载解压后，在环境变量中添加cub的文件的路径(CUB_HOME)：

3.安装pytorch3D：

下载解压pytorch3D源码，注意版本要与pytorch对应，每个版本下有注明其适用的pytorch版本。
修改extra_compile_args = {“cxx”: [“-std=c++14”]}为extra_compile_args = {“cxx”: [“”]}：

管理员身份打开 x64 Native Tools Command Prompt for VS 2019，（需要安装VS2019）依次输入：

activate gen6d
# 进入到pytorch3D目录下
cd external/pytorch3d
set DISTUTILS_USE_SDK=1
set PYTORCH3D_NO_NINJA=1
# 开始安装
python setup.py install

二.基于COLMAP制作个人数据集

参考链接

1.下载COLMAP软件【下载地址】，本文使用的是Windows下的CUDA版本：

解压后双击打开COLMAP.bat，出现如下界面软件安装成功:
拍摄视频采集图像
将拍摄好的视频存放在Gen6D/data\custom\video路径下：
拍摄视频的目标对象一定是要静止的，假如目标对象没有纹理，那背景的纹理必须要足够丰富。

# 激活虚拟环境
conda activate gen6d
# 进入Gen6D工程目录下
cd XXX
# 分割视频
# --transpose 解决图片颠倒
python prepare.py --action video2image --input data/custom/video/XXX.mp4 --output data/custom/XXX/images --frame_inter 10 --image_size 960 --transpose
# eg：python prepare.py --action video2image --input data/custom/video/people.mp4 --output data/custom/people/images --frame_inter 10 --image_size 960

2.恢复相机姿势

# colmap假如没有添加环境变量就需要colmap.bat的完整路径（原作者这里写的是colmap.exe的路径，但是运行会报错，打不开colmap）
python prepare.py --action sfm --database_name custom/XXX--colmap <colmap.bat路径地址>
# eg：python prepare.py --action sfm --database_name custom/people --colmap <colmap.bat路径地址>

3.剪裁点云，确定Z正方向和Z正方向

下载安装CloudCompare
使用CloudCompar打开Gen6D/data/custom/people/colmap/pointcloud.ply
手动指定目标区域：裁剪目标对象点云。

将裁剪后的目标对象点云导出：需要取消remaining部分，只选择segmented，用binary格式保存在Gen6D/data/custom/people/目录下，命名为object_point_cloud.ply。

使用CloudCompar打开刚保存的Gen6D/data/custom/people/object_point_cloud.ply。
同样先裁剪一片区域:

是根据所有选中点云拟合的，因此点云的选取十分关键，下图这种就很糟糕，完全不是目标物体所放置的平面。
X正方向的向量：通过选择俩个点计算出X正方向的向量。

在Gen6D/data/custom/people/目录下创建meta_info.txt将X正方向和Z正方向记录下来(好像不能复制粘贴要手打)。

0.090675 -1.660917 -0.227757
0.560382 -0.0676446 0.825467

三.测试个人数据集

1.安装ffmpeg

参考教程

2.文件组织结构

在data目录下新建custom文件夹，用于放置个人测试数据集，按照以下组织结构放置

Gen6D
|-- data
    |-- custom
       |-- car
           |-- object_point_cloud.ply  # object point cloud
           |-- meta_info.txt           # meta information about z+/x+ directions
           |-- images                  # images
           |-- colmap                  # colmap project

3.运行

# 假如没有添加环境变量就需要ffmpeg.exe的完整路径
# --transpose 解决图片颠倒
python predict.py --cfg configs/gen6d_pretrain.yaml --database custom/XXX --video data/custom/video/XXX.mp4 --resolution 960 --transpose --output data/custom/XXX/test --ffmpeg ffmpeg.exe
# eg: python predict.py --cfg configs/gen6d_pretrain.yaml --database custom/car --video data/custom/vid

4.BUG调试

遇到KeyError 的报错
Traceback (most recent call last):
File "D:\Gen6D-main\predict.py", line 110, in <module>
main(args)
File "D:\Gen6D-main\predict.py", line 33, in main
estimator.build(ref_database, split_type='all')
File "D:\Gen6D-main\estimator.py", line 145, in build
ref_ids = select_reference_img_ids_fps(database, ref_ids_all, self.cfg['ref_view_num'])
File "D:\Gen6D-main\utils\database_utils.py", line 115, in select_reference_img_ids_fps
poses = [database.get_pose(ref_id) for ref_id in ref_ids_all]
File "D:\Gen6D-main\utils\database_utils.py", line 115, in <listcomp>
poses = [database.get_pose(ref_id) for ref_id in ref_ids_all]
File "D:\Gen6D-main\dataset\database.py", line 290, in get_pose
return self.poses[img_id].copy()
KeyError: '16'

解决：
dataset\database.py中CustomDatabase类

if (self.colmap_root / 'sparse' / '0').exists():
    cameras, images, points3d = read_model(str(self.colmap_root / 'sparse' / '0'))
    self.poses, self.Ks, self.img_ids = parse_colmap_project(cameras, images, self.img_fns)
    self.img_ids = list(self.poses.keys())

在源代码出添加**self.img_ids = list(self.poses.keys())**这句代码

将会在 data/custom/mouse/test/video.mp4.看到结果视频。