【6D位姿估计】【深度学习】windows10下Gen6D代码Pytorch实现

提示：最近开始在【6D位姿估计】方面进行研究，记录相关知识点，分享学习中遇到的问题已经解决的方法。

文章目录

【6D位姿估计】【深度学习】windows10下Gen6D代码Pytorch实现
前言
下载源码并安装环境
安装 Pytorch3d
数据集
评估GenMOP/LINEMOD数据集
- 测试GenMOP
- 测试LINEMOD
测试个人数据集
训练网络
总结

前言

Gen6D是由港大&浙大的Liu, Yuan等人在《Gen6D: Generalizable Model-Free 6-DoF Object Pose Estimation from RGB Images【ECCV2022】》【论文地址】一文中提出了不需要CAD模型或者可渲染模型就可以泛化到新物体的位姿估计算法。之前的算法每次都只能针对某个特定的物体（例如PVNet/AAE）或者某个特定的物体类别（例如NOCS）进行位姿估计，而不能泛化到没见过的物体，论文提出的算法只需利用最容易获得的RGB图片，而不需要额外的物体mask或者深度相机得到的深度，就可以泛化到新物体的位姿估计算法。
在详细解析Gen6D网络之前，首要任务是搭建Gen6D【Pytorch-demo地址】所需的运行环境，并完成模型训练和测试工作，展开后续工作才有意义。

论文(源代码)作者的知乎博文

下载源码并安装环境

在win10环境下装anaconda环境，方便搭建专用于Gen6D模型的虚拟环境。

安装GPU版本的pytorch教程，pytorch-gpu版本需要根据个人计算机去安装相应版本。

# 创建虚拟环境
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安装十分复杂，博主会单独进行安装讲解。

# 查看所有安装的包
pip list
conda list

最终的安装的所有包：

检查torch版,已经安装torch-gpu版本

# 查看pytorch版本
import torch
print(torch.__version__)
# 查看cuda版本
print(torch.version.cuda)
# 查看cuda是否可用
print(torch.cuda.is_available())
# 查看可用cuda数量
print(torch.cuda.device_count())

安装 Pytorch3d

安装PyTorch3d参考
1.安装必要依赖库

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

2.cub安装配置：查询对照表：

这里的cuda是物理机器上实际安装的cuda版本(环境变量里配置的cuda)，不是虚拟环境中的cuda版本。

否则即使pytorch3d安装成功，也不能使用gpu。
Windows安装多版本CUDA及其切换参考

根据自己的cuda 版本选择对应的cub realase版本下载，对应CUB 1.11.0 (CUDA Toolkit 11.3)：

下载解压后，在环境变量中添加cub的文件的路径(CUB_HOME)：

3.安装pytorch3D：下载解压pytorch3D源码，注意版本要与pytorch对应，每个版本下有注明其适用的pytorch版本。

# 创建文件夹,将下载好的pytorch3D放到该目录下(重命名)
mkdir -p external

在Windows系统下，不建议使用git clone方式下载源码。

下载0.7.1版本的pytorch3D版本，并解压到Gen6D工程的external/目录下(博主将文件重命名为pytorch3d)：

修改extra_compile_args = {“cxx”: [“-std=c++14”]}为extra_compile_args = {“cxx”: [“”]}：

安装 Visual Studio 2019，管理员身份打开 x64 Native Tools Command Prompt for VS 2019，依次输入：

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

安装完成：

查看pytorch3d版本：

import pytorch3d
print(pytorch3d.__version__)

数据集

官方数据集下载
包括：预训练模型(gen6d_pretrain.tar.gz)、GenMOP数据集(genmop.tar.gz)、处理后的LINEMOD数据集(linemod.tar.gz)、co3d数据集(co3d.tar.gz)、Google扫描物体数据(google_scanned_objects.tar.gz)和ShapeNet渲染图(shapenet.tar.gz)。

数据集和预训练网络权重在Gen6D目录下的组织结构在具体的章节会详细讲解。

评估GenMOP/LINEMOD数据集

在Gen6D目录下新建data文件夹，用于放置数据集和预训练权重。
将gen6d_pretrain.tar.gz、genmop.tar.gz和linemod.tar.gz)解压后按以下组织结构放置：

Gen6D
|-- data
    |-- model
        |-- detector_pretrain
            |-- model_best.pth
        |-- selector_pretrain
            |-- model_best.pth
        |-- refiner_pretrain
            |-- model_best.pth
    |-- GenMOP
        |-- chair 
            ...
    |-- LINEMOD
        |-- cat 
            ...

测试GenMOP

# 测试GenMOP数据集中的tformer
python eval.py --cfg configs/gen6d_pretrain.yaml --object_name genmop/XXX
# eg: python eval.py --cfg configs/gen6d_pretrain.yaml --object_name genmop/cup

在data/vis_final/gen6d_pretrain/genmop/XXX中保存估计姿势的3D边界框的结果：

真实3D边界框用绿色绘制，而预测用蓝色绘制

在data/vis_inter/gen6d_pretrain/genmop/XXX中保存有关检测、视点选择和姿势细化的中间结果：

红色3D边界框代表输入姿势，绿色代表地面实况，蓝色代表当前细化步骤的输出。

测试LINEMOD

# 测试LINEMOD数据集中的benchvise
python eval.py --cfg configs/gen6d_pretrain.yaml --object_name linemod/XXX
# eg：python eval.py --cfg configs/gen6d_pretrain.yaml --object_name linemod/benchvise

在data/vis_final/gen6d_pretrain/linemod/XXX中保存估计姿势的3D边界框的结果：

在data/vis_inter/gen6d_pretrain/genmop/XXX中保存有关检测、视点选择和姿势细化的中间结果：

测试个人数据集

个人数据集的制作参考：博主教程和官方教程。
获得以下文件：

需要安装ffmpeg：参考教程
在data目录下新建custom文件夹，用于放置个人测试数据集，按照以下组织结构放置：

Gen6D
|-- data
    |-- model
    |-- GenMOP
    |-- LINEMOD
    |-- custom
        |-- people
        	...

# 假如没有添加环境变量就需要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/people --video data/custom/video/people.mp4 --resolution 960 --output data/custom/people/test --ffmpeg ffmpeg.exe

在data/custom/XXX/test中保存估计姿势结果(视频转了GIF)：
在这里插入图片描述

训练网络

训练部分时长太久了，等出效果在更新。

总结

尽可能简单、详细的介绍Gen6D的安装流程以及解决了安装过程中可能存在的问题。后续会根据自己学到的知识结合个人理解讲解Gen6D的原理和代码。