【代码复现】Instance-Adaptive and Geometric-Aware Keypoint Learning for Category-Level 6D Object Pose Estimation

安装AG-Pose环境

  1. 创建环境:

    1
    conda create --name agpose python=3.9

  2. 激活环境:

    1
    conda activate agpose

  3. 安装PyTorch:

    1
    pip3 install torch==1.12 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu113

  4. 安装其他依赖:

    1
    2
    3
    4
    5
    6
    pip3 install gorilla-core==0.2.5.3
    pip3 install opencv-python

    git clone https://github.com/Leeiieeo/AG-Pose.git
    cd AG-Pose/model/pointnet2
    python setup.py install

  5. Clone mentian/object-deformnet并安装nn_distance,需要使用该仓库中的代码进行数据处理:

    1
    2
    3
    git clone https://github.com/mentian/object-deformnet.git
    cd object-deformnet/lib/nn_distance
    python setup.py install

数据集下载

  1. 下载NOCS数据集到数据集文件夹中:

    • http://download.cs.stanford.edu/orion/nocs/camera_composed_depth.zip
    • http://download.cs.stanford.edu/orion/nocs/camera_train.zip
    • http://download.cs.stanford.edu/orion/nocs/camera_val25K.zip
    • http://download.cs.stanford.edu/orion/nocs/gts.zip
    • http://download.cs.stanford.edu/orion/nocs/obj_models.zip
    • http://download.cs.stanford.edu/orion/nocs/real_test.zip
    • http://download.cs.stanford.edu/orion/nocs/real_train.zip
  2. 下载JiehongLin/Self-DPDN的分割结果到仓库文件夹中:

    • https://drive.google.com/file/d/1hNmNRr7YRCgg-c_qdvaIzKEd2g4Kac3w/
  3. 下载使用mentian/object-deformnet进行数据预处理时用到的数据到仓库文件夹中:

    • https://drive.google.com/file/d/1p72NdY4Bie_sra9U8zoUNI4fTrQZdbnc/
  4. 下载运行AG-Pose时使用到的分割结果:

    • http://home.ustc.edu.cn/~llinxiao/segmentation_results.zip
  5. 下载NOCS数据集中缺失的物体模型(From: https://github.com/mentian/object-deformnet/issues/3#issuecomment-698858332)到数据集文件夹中:

    • https://drive.google.com/file/d/1rWkxEVJJh_kWIqxudn_i6sJ-hhd0E7TV/

数据集处理

JiehongLin/Self-DPDN

遵循JiehongLin/Self-DPDN的处理方式,首先解压上述文件,并组织如下,注意,data文件夹位于Self-DPDN文件夹下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
data
├── camera -> /data1/dataset/nocs/tanmx/dpdn/camera
│   ├── train
│   └── val
├── camera_full_depths -> /data1/dataset/nocs/tanmx/dpdn/camera_full_depths/
│   ├── train
│   └── val
├── gts -> /data1/dataset/nocs/tanmx/dpdn/gts/
│   ├── real_test
│   └── val
├── mean_shapes.npy
├── obj_models -> /data1/dataset/nocs/tanmx/dpdn/obj_models/
│   ├── real_test
│   ├── real_train
│   ├── train
│   └── val
├── real -> /data1/dataset/nocs/tanmx/dpdn/real/
│   ├── test
│   └── train
└── segmentation_results -> /data1/dataset/nocs/tanmx/dpdn/segmentation_results/
├── test_trainedwithMask
├── test_trainedwoMask
└── train_trainedwoMask

21 directories, 1 file

进入obj_models/val/02876657/d3b53f56b4a7b3b3c9f016d57db96408文件夹,其中内容为

1
2
$ ls
bbox.txt model.mtl model.obj

查看model.obj文件:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
$ cat model.obj
newmtl Material.57774a50243e5d429f3a4ea105ec3930
d 0.5
Tr 0.5
Kd 0.392157 0.584314 0.929412
newmtl Material.90b00bf69165eda831fe702671104217
Tr 0.0
Kd 0.6 0.117647 0.117647
newmtl Material.dc7a3ca9ca091ddbf48733c3a604f557
Tr 0.0
Kd 1.0 1.0 1.0
newmtl Material.b2cdaf4394d1a4f850458878b81c7fa1
Tr 0.0
Kd 0.8 0.6 0.0

解压https://drive.google.com/file/d/1rWkxEVJJh_kWIqxudn_i6sJ-hhd0E7TV/中的文件:

1
2
3
4
5
$ unzip obj_models.zip
$ ls
02876657 obj_models.zip
$ ls 02876657/d3b53f56b4a7b3b3c9f016d57db96408/
bbox.txt info.txt model_bad.mtl model_bad.obj model.mtl model.obj

删除原文件夹obj_models/val/02876657/d3b53f56b4a7b3b3c9f016d57db96408中的内容,将02876657/d3b53f56b4a7b3b3c9f016d57db96408文件夹中的内容复制到该文件夹中:

1
2
3
4
$ rm obj_models/val/02876657/d3b53f56b4a7b3b3c9f016d57db96408/*
$ cp 02876657/d3b53f56b4a7b3b3c9f016d57db96408/* obj_models/val/02876657/d3b53f56b4a7b3b3c9f016d57db96408/
$ ls obj_models/val/02876657/d3b53f56b4a7b3b3c9f016d57db96408
bbox.txt info.txt model_bad.mtl model_bad.obj model.mtl model.obj

执行:

1
python data_processing.py

执行后目录变为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
data
├── camera -> /data1/dataset/nocs/tanmx/dpdn/camera
│   ├── train
│   ├── train_list_all.txt
│   ├── train_list.txt
│   ├── val
│   └── val_list_all.txt
├── camera_full_depths -> /data1/dataset/nocs/tanmx/dpdn/camera_full_depths/
│   ├── train
│   └── val
├── gts -> /data1/dataset/nocs/tanmx/dpdn/gts/
│   ├── real_test
│   └── val
├── mean_shapes.npy
├── obj_models -> /data1/dataset/nocs/tanmx/dpdn/obj_models/
│   ├── real_test
│   ├── real_train
│   ├── train
│   └── val
├── real -> /data1/dataset/nocs/tanmx/dpdn/real/
│   ├── test
│   ├── test_list_all.txt
│   ├── train
│   ├── train_list_all.txt
│   └── train_list.txt
└── segmentation_results -> /data1/dataset/nocs/tanmx/dpdn/segmentation_results
├── test_trainedwithMask
├── test_trainedwoMask
└── train_trainedwoMask

21 directories, 7 files

执行:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
$ ls -l data/camera/
total 9712
drwxr-xr-x 27502 tanmx tanmx 577536 Jun 14 2019 train
-rw-rw-r-- 1 tanmx tanmx 4675000 Apr 29 11:00 train_list_all.txt
-rw-rw-r-- 1 tanmx tanmx 4235159 Apr 29 15:57 train_list.txt
drwxr-xr-x 2502 tanmx tanmx 69632 Jun 14 2019 val
-rw-rw-r-- 1 tanmx tanmx 375000 Apr 29 11:00 val_list_all.txt

$ ls -l -h data/real/
total 228
drwxrwxr-x 8 tanmx tanmx 4096 Nov 13 2018 test
-rw-rw-r-- 1 tanmx tanmx 49572 Apr 29 11:00 test_list_all.txt
drwxrwxr-x 9 tanmx tanmx 4096 Jun 14 2019 train
-rw-rw-r-- 1 tanmx tanmx 82042 Apr 29 11:00 train_list_all.txt
-rw-rw-r-- 1 tanmx tanmx 82042 Apr 29 16:09 train_list.txt

mentian/object-deformnet

遵循mentian/object-deformnet的处理方式,首先解压上述文件,并组织如下,注意,data文件夹位于object-deformnet文件夹下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
data
├── camera -> /data1/dataset/nocs/tanmx/spd/camera/
│   ├── train
│   └── val
├── deformnet_eval -> /data1/dataset/nocs/tanmx/spd/deformnet_eval/
│   ├── camera
│   ├── mrcnn_results
│   ├── nocs_results
│   └── real
├── gts -> /data1/dataset/nocs/tanmx/spd/gts/
│   ├── real_test
│   └── val
├── obj_models -> /data1/dataset/nocs/tanmx/spd/obj_models/
│   ├── real_test
│   ├── real_train
│   ├── train
│   └── val
├── pose_dataset.py
├── real -> /data1/dataset/nocs/tanmx/spd/real/
│   ├── test
│   └── train
└── shape_dataset.py

19 directories, 2 files

按顺序执行:

1
2
3
cd preprocess
python shape_data.py
python pose_data.py

执行python shape_data.py后目录变为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
data
├── camera -> /data1/dataset/nocs/tanmx/spd/camera/
│   ├── train
│   └── val
├── deformnet_eval -> /data1/dataset/nocs/tanmx/spd/deformnet_eval/
│   ├── camera
│   ├── mrcnn_results
│   ├── nocs_results
│   └── real
├── gts -> /data1/dataset/nocs/tanmx/spd/gts/
│   ├── real_test
│   └── val
├── obj_models -> /data1/dataset/nocs/tanmx/spd/obj_models/
│   ├── camera_train.pkl
│   ├── camera_val.pkl
│   ├── mug_meta.pkl
│   ├── real_test
│   ├── real_test.pkl
│   ├── real_train
│   ├── real_train.pkl
│   ├── ShapeNetCore_2048.h5
│   ├── ShapeNetCore_4096.h5
│   ├── train
│   └── val
├── pose_dataset.py
├── real -> /data1/dataset/nocs/tanmx/spd/real/
│   ├── test
│   └── train
└── shape_dataset.py

19 directories, 9 files

执行python pose_data.py后目录变为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
data
├── camera -> /data1/dataset/nocs/tanmx/spd/camera
│   ├── train
│   ├── train_list_all.txt
│   ├── train_list.txt
│   ├── val
│   ├── val_list_all.txt
│   └── val_list.txt
├── deformnet_eval -> /data1/dataset/nocs/tanmx/spd/deformnet_eval
│   ├── camera
│   ├── mrcnn_results
│   ├── nocs_results
│   └── real
├── gts -> /data1/dataset/nocs/tanmx/spd/gts/
│   ├── real_test
│   └── val
├── obj_models -> /data1/dataset/nocs/tanmx/spd/obj_models/
│   ├── camera_train.pkl
│   ├── camera_val.pkl
│   ├── mug_meta.pkl
│   ├── real_test
│   ├── real_test.pkl
│   ├── real_train
│   ├── real_train.pkl
│   ├── ShapeNetCore_2048.h5
│   ├── ShapeNetCore_4096.h5
│   ├── train
│   └── val
├── pose_dataset.py
├── real -> /data1/dataset/nocs/tanmx/spd/real/
│   ├── test
│   ├── test_list_all.txt
│   ├── test_list.txt
│   ├── train
│   ├── train_list_all.txt
│   └── train_list.txt
└── shape_dataset.py

19 directories, 17 files

执行:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
$ ls -l data/camera/
total 9996
drwxr-xr-x 27502 tanmx tanmx 577536 Jun 14 2019 train
-rw-rw-r-- 1 tanmx tanmx 4675000 Apr 29 14:00 train_list_all.txt
-rw-rw-r-- 1 tanmx tanmx 4235159 Apr 29 18:42 train_list.txt
drwxr-xr-x 2502 tanmx tanmx 69632 Jun 14 2019 val
-rw-rw-r-- 1 tanmx tanmx 375000 Apr 29 14:00 val_list_all.txt
-rw-rw-r-- 1 tanmx tanmx 288300 Apr 29 19:15 val_list.txt

$ ls -l data/real/
total 280
drwxrwxr-x 8 tanmx tanmx 4096 Nov 13 2018 test
-rw-rw-r-- 1 tanmx tanmx 49572 Apr 29 14:00 test_list_all.txt
-rw-rw-r-- 1 tanmx tanmx 49572 Apr 29 19:21 test_list.txt
drwxrwxr-x 9 tanmx tanmx 4096 Jun 14 2019 train
-rw-rw-r-- 1 tanmx tanmx 82042 Apr 29 14:00 train_list_all.txt
-rw-rw-r-- 1 tanmx tanmx 82042 Apr 29 18:53 train_list.txt

$ ls -l data/obj_models/
total 116060
-rw-rw-r-- 1 tanmx tanmx 26695895 Apr 29 11:32 camera_train.pkl
-rw-rw-r-- 1 tanmx tanmx 4634293 Apr 29 11:36 camera_val.pkl
-rw-rw-r-- 1 tanmx tanmx 23197 Apr 29 11:37 mug_meta.pkl
drwxrwxr-x 2 tanmx tanmx 4096 Sep 27 2019 real_test
-rw-rw-r-- 1 tanmx tanmx 443600 Apr 29 11:37 real_test.pkl
drwxrwxr-x 2 tanmx tanmx 4096 Sep 27 2019 real_train
-rw-rw-r-- 1 tanmx tanmx 443584 Apr 29 11:37 real_train.pkl
-rw-rw-r-- 1 tanmx tanmx 28950893 Apr 29 12:21 ShapeNetCore_2048.h5
-rw-rw-r-- 1 tanmx tanmx 57618690 Apr 29 11:39 ShapeNetCore_4096.h5
drwxr-xr-x 12 tanmx tanmx 4096 Sep 27 2019 train
drwxr-xr-x 12 tanmx tanmx 4096 Sep 27 2019 val

运行代码

需要JiehongLin/Self-DPDN处理后的数据集及其list和mentian/object-deformnet处理后的物体模型pkl,组织如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
data
├── camera -> /data1/dataset/nocs/tanmx/agpose/camera
│   ├── train
│   ├── train_list_all.txt
│   ├── train_list.txt
│   ├── val
│   └── val_list_all.txt
├── camera_full_depths -> /data1/dataset/nocs/tanmx/agpose/camera_full_depths/
│   ├── train
│   └── val
├── gts -> /data1/dataset/nocs/tanmx/agpose/gts/
│   ├── real_test
│   └── val
├── obj_models -> /data1/dataset/nocs/tanmx/agpose/obj_models/
│   ├── camera_train.pkl
│   ├── real_test
│   ├── real_train
│   ├── real_train.pkl
│   ├── train
│   └── val
├── real -> /data1/dataset/nocs/tanmx/agpose/real/
│   ├── test
│   ├── test_list_all.txt
│   ├── train
│   ├── train_list_all.txt
│   └── train_list.txt
└── segmentation_results -> /data1/dataset/nocs/tanmx/agpose/segmentation_results
├── CAMERA25
└── REAL275

20 directories, 8 files

之后按照仓库README中的命令运行代码即可。