Skip to content

Commit ab7e882

Browse files
qq456cvbqq456cvb
committed
evaluate on nocs
1 parent 9460bb3 commit ab7e882

File tree

11 files changed

+851
-32
lines changed

11 files changed

+851
-32
lines changed

‎.gitignore‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
1-
outputs/
1+
checkpoints/
22
**/__pycache__

‎README.md‎

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,33 +42,64 @@ CVPR 2022
4242
This is the official code implementation of CPPF, including both training and testing. Inference on custom datasets is also supported.
4343

4444
# Installation
45-
You can run the following command to setup an environment, tested on Ubuntu 16.04:
45+
You can run the following command to setup an environment, tested on Ubuntu 18.04:
4646

4747
```
4848
conda create -n cppf python=3.8
49-
conda install pytorch cudatoolkit=10.2 -c pytorch-lts
50-
pip install tqdm opencv-python scipy matplotlib open3d==0.12.0 hydra-core pyrender cupy-cuda102 PyOpenGL-accelerate
49+
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch-lts
50+
pip install tqdm opencv-python scipy matplotlib open3d==0.12.0 hydra-core pyrender cupy-cuda102 PyOpenGL-accelerate OpenEXR
5151
CXX=g++-7 CC=gcc-7 pip install MinkowskiEngine==0.5.4 -v
5252
```
5353
We use [Hydra](https://hydra.cc/) configuration system to run scripts.
5454
Notice that we use pyrender with OSMesa support, you may need to install OSMesa after running ```pip install pyrender```, more details can be found [here](https://pyrender.readthedocs.io/en/latest/install/index.html).
5555

5656
# Train on ShapeNet Objects
57-
First, download [ShapeNet v2](https://shapenet.org/) dataset and modify the ``shapenet_root`` key in ``config/config.yaml`` to point to the location of the dataset.
57+
<details>
58+
<summary>Data Preparation</summary>
59+
60+
Download [ShapeNet v2](https://shapenet.org/) dataset and modify the ``shapenet_root`` key in ``config/config.yaml`` to point to the location of the dataset.
61+
62+
</details>
63+
64+
<details>
65+
<summary>Train on NOCS REAL275 objects</summary>
5866

5967
To train on synthetic ShapeNet objects that appear in NOCS REAL275, run:
6068
```
6169
python train.py category=bottle,bowl,camera,can,laptop,mug -m
6270
```
71+
72+
For laptops, an auxiliary segmentation is needed to ensure a unique pose. Please refer to <a href='#laptop-aux'>Auxiliary Segmentation for Laptops</a>/
73+
</details>
74+
75+
<details>
76+
<summary>Train on SUN RGB-D objects</summary>
77+
6378
To train on synthetic ShapeNet objects that appear in SUN RGB-D, run:
6479
```
6580
python train.py category=bathtub,bed,bookshelf,chair,sofa,table -m
6681
```
82+
</details>
83+
84+
<details>
85+
<summary id='laptop-aux'>Auxiliary Segmentation for Laptops</summary>
86+
87+
For Laptops, geometry alone cannot determine the pose unambiguously, we rely on an auxiliary segmentation network that segments out the lid and the keyboard base.
88+
89+
To train the segmenter network, first download our Blender physically rendered laptop images from [Google Drive]() and place it under ``data/laptop``. Then run the following command:
90+
```
91+
python train_laptop_aux.py
92+
```
93+
</details>
94+
6795

6896
# Pretrained Models
6997
# Test on NOCS REAL275
7098

71-
## With Instance Segmentation Mask
99+
<details>
100+
<summary>With Instance Segmentation Mask</summary>
101+
102+
</details>
72103

73104
## With Bounding Box Mask
74105

‎config/config.yaml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ regress_right: False
77
tr_num_bins: 32
88
rot_num_bins: 36
99
opt:
10-
weight_decay: 0
10+
weight_decay: 1e-4
1111
lr: 1e-3
1212
up_sym: False
1313
right_sym: False
@@ -19,11 +19,11 @@ defaults:
1919
- category: bowl
2020

2121
knn: 60
22-
shapenet_root: ../disk/ShapeNetCore.v2
22+
shapenet_root: ./data/ShapeNetCore.v2
2323

2424
hydra:
2525
run:
26-
dir: ./outputs/${category}/
26+
dir: ./checkpoints/${category}/
2727
sweep:
28-
dir: outputs
28+
dir: checkpoints
2929
subdir: ${category}

‎config/laptop_aux.yaml‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
category: laptop
2+
max_epoch: 20
3+
batch_size: 16
4+
data_root: ./data/laptop
5+
opt:
6+
weight_decay: 0
7+
lr: 1e-3
8+
9+
hydra:
10+
run:
11+
dir: ./checkpoints/${category}_aux/

‎data/laptop_nonreal.txt‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
03642806/1b5ec8c5ca53caa5ad6be7ca8dcfc55b
2+
03642806/1ce688d90a2010a69718283200011d2a
3+
03642806/1d66c26055573243a0efaa0482f35281
4+
03642806/1d1947a3104bdbf4752cca7e96f0bf86
5+
03642806/2c61f0ba3236fe356dae27c417fa89b
6+
03642806/2c697db8a982242d54a8d8e3b6972768
7+
03642806/2cafa6107fff001f7a470f4d014a755f
8+
03642806/2d38fa685e71721986a7e0484f5238e4
9+
03642806/2fbc09c414c80a2d50577cf04f3bf74a
10+
03642806/3b2bddb0e8ff57c85831a0624cf5a945
11+
03642806/3b2db36aaa2546b99c7c402f274622c
12+
03642806/3b9a99663babc301e2ce1e6d2467cb03
13+
03642806/3dd1b4cdc6d3b33f676394b063daab40
14+
03642806/3e9af28eb2d6e216a4e3429ccb8eaf16
15+
03642806/4f3575df3821e08c466909b3e9553909
16+
03642806/5ea74fdff0f69a60546d7b6431984ed4
17+
03642806/6a9cffdd120720a6512143cd347a4650
18+
03642806/6ae14173f5bface824a85bcf396c330
19+
03642806/6b61ef17b4f45050b598e8984f11eb0c
20+
03642806/6c1cba2e5d2ebb552209327b12ad3fd
21+
03642806/6d88f11ac37f07a86a7e0484f5238e4
22+
03642806/6da67be7e83afd3be3b927dfeefbac4e
23+
03642806/6f1d1ad1d8238e84957b6a4f3924d982
24+
03642806/6f20fd2e8d8d518da5f7b1c78f76e23e
25+
03642806/6fd01f7db45bff66cf01bc59d215f0
26+
03642806/7a6b074fec1b8b9c810b27ad0ef73ed
27+
03642806/7b4260884a1dfd76b080af510dd640b
28+
03642806/7bdac195f81588ff743bd44138e7d2e5
29+
03642806/7d733a513d905a1a56fd4d28526341bf
30+
03642806/7d2718656efb3920d1954a7ac233ac80
31+
03642806/7e5b970c83dd97b7823eead1c8e7b3b4
32+
03642806/7ebe5cc71bdeebd9dd70b866e1cfbf35
33+
03642806/8b10449b4c25fb0939773ee5b5578e3
34+
03642806/8e38d4d1f849f024296462783ef76d22
35+
03642806/8fefaa8179404c4ae165aa9cd54a0a29
36+
03642806/9b4ab67eb448c49c11ced4a54f2e6229
37+
03642806/9cb17f33ba3c057255d5baa5be4ca05
38+
03642806/9cb54ee28aec14013cb02989e2da5a2a
39+
03642806/9d48ab8c41174e60888cad7f6c0e6001
40+
03642806/9ddc0f876f54a784f3655a2ce983ab1f
41+
03642806/9f4a0de5fb60ca3e1ec0c47eb137bcb
42+
03642806/9fa387d7f442b96e75e60c00fabe2744
43+
03642806/9fc5b76d363ca64ed03066fc8168e9c6
44+
03642806/16c49793f432cd4b33e4e0fe8cce118e
45+
03642806/21c0187bf2174c1684c4330348773d0a
46+
03642806/21e03f21a8724ab5bfa1d058f8ea2f0a
47+
03642806/25f391ae823c118963439329667c7723
48+
03642806/26ad81059039dfcca3e30e62a8e6f77f
49+
03642806/33b9567a9f534747cdd9e16bfb7eee3d
50+
03642806/36bf7cd2652ccf48dd42bb2503ef78ef
51+
03642806/37b5e63e9f80e2402052a3b22ea3f616
52+
03642806/39b54fd684a4ed3e6f0a9a10dc885e17
53+
03642806/40e0263822860cc5d69e26904fa68e7f
54+
03642806/41e263d2af5c667f988abef1fd117e7
55+
03642806/42f71a2df30eabb1b8183a4a81361b94
56+
03642806/51dad6ff8c00b7be288d8f1144f7979f
57+
03642806/66e3b7c7f2e8e9297fd8853234f5e918
58+
03642806/68adcb59ad19c431676394b063daab40
59+
03642806/69c1b09e3cd5056d11f9b5b6530d6e16
60+
03642806/73b436084ec1697627d962ee5b0a40c1
61+
03642806/74bddbc79d758bcbda39f7cbbf0da6b
62+
03642806/76d74a604c1ddacd6093b210438ee0f7
63+
03642806/79b420b8fdf6a37db747d76ca5db985e
64+
03642806/79ed190f0e35ccdf98ce68842a487b19
65+
03642806/80c013931b7f7b49636ee1998fc5378
66+
03642806/97be95bc027c84da5a7ade0e7a3afdc8
67+
03642806/97df7577c2690ab82290f0f0820b73a5
68+
03642806/151fccb37c2d51b0b82fa571a1cdbf24
69+
03642806/195e2a65e5396271bc7eb7ec22211551
70+
03642806/229f440bbd8f5482355323b019d9d87e
71+
03642806/241ec8a746dd1cfc78f71a335ebabfa5
72+
03642806/283a115028fffe6e4e67e2e8459ebde
73+
03642806/454a88f4b2fe4472f9cacf0b2853f40
74+
03642806/464edfe14e9fa45c3394926146371698
75+
03642806/476cd9a7a0a7c491c96e2e556a25ab0c
76+
03642806/517de75577ac6e8a42b9615216f9a30d
77+
03642806/520d98e360cf44ec8139dd63d55edc44
78+
03642806/546b40d15cd483873339d8259ddfa7bb
79+
03642806/606f4b61d81085d1adbb6cf5d72fbe37
80+
03642806/679bb4a724b1e2bc3339d8259ddfa7bb
81+
03642806/760b1f82f7133ca1de7d2830be38218d
82+
03642806/894e47a6adb76680d4eb7e68e898dc44
83+
03642806/928c86eabc0be624c2bf2dcc31ba1713
84+
03642806/1312ea502b4e9b51701c1f58e22b85e8
85+
03642806/1622e466794cf72998536cff5312b747
86+
03642806/2134ad3fc25a6284193a4c984002ed32
87+
03642806/2211a40cc77a085362c091e763f81d3
88+
03642806/3909e96930cb1e32a4f2e5eab55eddea
89+
03642806/5173aa7f75ff3cf1b55fde51a411949f
90+
03642806/5430ba1b4995af02717dccfa4b0eec5a
91+
03642806/6103fd8c11413be975e42c1faa1148bb
92+
03642806/6146d9fd9460370c578b3fc82443ab96
93+
03642806/7312c9dc3c2802e926db2cd5ecb2df2b
94+
03642806/9297e108b9780fe7da0ccb8671dd235b
95+
03642806/9980e8f99eb0098bc6b48dee64de8434
96+
03642806/17069b6604fc28bfa2f5beb253216d5b
97+
03642806/19117e74b575cdca82555f4c45537277
98+
03642806/25192ef0befdd75d291d2d0d560233f4
99+
03642806/35462b6ea0597bf555e8ddeaf1791474
100+
03642806/69850d70fcaf5a464161e123a102b12b
101+
03642806/75936cfb6c3ee142205444155ae5c84d
102+
03642806/92702a4de72aac3e20685b5a7e34b501
103+
03642806/129237ff9e4f5d5c95b5eddcdeed8f43
104+
03642806/471005df47f0b4a6625cbbe81bda8c5c
105+
03642806/490245e63baf78fb1bf5877a865437f7
106+
03642806/689667fca044e15b941a99145cf33e72
107+
03642806/942156d763bd7e867cff7d16e20f126a
108+
03642806/4348768ff9802fc5ce1decf313ece97f
109+
03642806/5515081bfe708008edf29b343d09db45
110+
03642806/11494281ca09a63d3b54554be306bc70
111+
03642806/93958423b98be8b538ff1b6d120c56aa
112+
03642806/5085268208bb23ba7b037ded105df073
113+
03642806/5458437965b3d5da2462e129dafd29fb
114+
03642806/a4b410734514306ac401e233323032d6
115+
03642806/a4eb5dea7eda6ab58e5480b52f7861ca
116+
03642806/a7f983f1d0642745135a402b573354e4
117+
03642806/a59d3d87068d313c2656684d670220c2
118+
03642806/a702d347f32e0d3a229fe85eb79d32b
119+
03642806/a76724964bc58c73ec00b83eb025362c
120+
03642806/a96983022f36a0a83ad38b993d9e745f
121+
03642806/aaebdb95de0459f66a51abe62a0f8cce
122+
03642806/ab9321d89df4d58e78c854de246ffde1
123+
03642806/ac3ca8ae5e3203117f2411014f221968
124+
03642806/b3a2081fc19c7bbac241f229aa037cc2
125+
03642806/b3c5bb4eb4c797ae6984b767317d33c4
126+
03642806/b5f6fd84a3f44ddb1aa47689117a61e1
127+
03642806/b9b521a7e25caf837b18e772ada5a339
128+
03642806/b36c5ee8b04b4b0243cecf206f8bdb34
129+
03642806/b210db4bca21536ed6daca3908cc94ab
130+
03642806/b211cfb105e9f97e6436916a86a90ed7
131+
03642806/b272eb1541844f55ad7b15cd0061e5ee
132+
03642806/b806daf849a5dba289c212008d2a390e
133+
03642806/b913c80eca0afa4a1b7e0829f0e4fe61
134+
03642806/b436271050d647052f8d6d501b18a4b5
135+
03642806/b233163860361eda8cfacef5204026d6
136+
03642806/bb33d26f324d19d94845e0946708405d
137+
03642806/be6a1fba2165bdfe2f81cb887dc35578
138+
03642806/be6af6c0bdaee6e06ec6fb380960ec65
139+
03642806/c0c055ba43c3528a3cb4d5b227cc213b
140+
03642806/c3eac3a391b81bc8387188c556ebf654
141+
03642806/c012da029408fb8d498c10f665c444af
142+
03642806/c912c50b812268e8fc9e3f0d993c8121
143+
03642806/ca98fa6fe9c91b8b8b0a7ba4456e7659
144+
03642806/caa4afd404f24d21275c1147a304ed86
145+
03642806/cb1e3a990782678b4b6682da890df381
146+
03642806/cb6ae19022733ef1956ad4a5a9b8400c
147+
03642806/cb090bd99ed76f9689661310be87a70d
148+
03642806/cc0535a34cdc7d676bf98d15712168f
149+
03642806/cc596eaa22c3e6cf37e1b9ea967d5071
150+
03642806/ccf62b71e7a9aedfb162fa62b36ac9a4
151+
03642806/ce13a85fb71694fcb611830890d7aa97
152+
03642806/d2eabcc071cec6f420bec9be02f008ca
153+
03642806/d7e7e6651a23afc68ba4e518219eb66a
154+
03642806/d8ef2a8a3ccff3fa5cad2ceae48f2da5
155+
03642806/d88e4093a73064b6b44d9812f259e403
156+
03642806/d91a3191d427bf0467710679e532af1
157+
03642806/d122ecd3c34e87eaf76120a09d6a0db
158+
03642806/d380e53dafa66c1a86a7e0484f5238e4
159+
03642806/d838fc238fc640b83bba289a9fdd2a7d
160+
03642806/d1388dc83fa45da33a8222bbf59f8f16
161+
03642806/dafca3237829fef9f6fef4a0ba2e7186
162+
03642806/dbc61cbed5f7f2b33c1abb78f1519c49
163+
03642806/dbcd5a88a9d4f1d7579cfe4420588034
164+
03642806/dc264eab83ca12b3da4c0d8596dff972
165+
03642806/dd205f0f8bc78bae359d7b39cfebb287
166+
03642806/dd282ed5d4956590b713df0d260c2628
167+
03642806/dda1832a36858f06ea791b47ef8b531a
168+
03642806/de2e95eac460c361e862e3cac45aa769
169+
03642806/de9a39e8ef9c313bbba0ffaa383992dc
170+
03642806/df624231ef05b83d54c7ab6aa5674ccc
171+
03642806/e0cc767bf10134509fd6a63b3b6f6cd4
172+
03642806/e4c34c87ed1bc2191ef7a71d6e01357e
173+
03642806/e8f89f185dd8165b9b07343517e0f88e
174+
03642806/e9e28a11f71337fc201115f39f20d1ff
175+
03642806/e34a4bd815f38858fb15d1c4afffaed3
176+
03642806/e466c2c86a439c1faebff3b001eb4a27
177+
03642806/e1773be8392c5fc4f7d4d3ced798ad90
178+
03642806/e083105e9c2a28bb0c3a03d0a1f182f
179+
03642806/e555880523dc11425b440d1b0fe8ea5c
180+
03642806/ebc59a9b4291ce274c3c121820623e4e
181+
03642806/ec3223edbdef2db93334fb6668ccd834
182+
03642806/ee31f8d77002d8869b07343517e0f88e
183+
03642806/eeca77a4ae8aa0ce3bef19ecc8621d89
184+
03642806/ef5b312fc20f1b20aab089a6db538ba7
185+
03642806/ef6d92c90aeabf5becae27d182a3e41c
186+
03642806/f00ec45425ee33fe18363ee01bdc990
187+
03642806/f3d162798a285521da0ccb8671dd235b
188+
03642806/f7c26b8c94ba8214397c35f585745a82
189+
03642806/f9f484e4364663d61bb5a64b9a0f552b
190+
03642806/f28fe2f836cb0dfaa8f4918d18e28074
191+
03642806/f53a54a64dc475195f2e56f31ec62a24
192+
03642806/f72dc1ffeae0168aadcfd37206a0d18b
193+
03642806/f14056ee8bbebeecc1b05209f08e5ec6
194+
03642806/fb641fe28d0f5e73d51fce576fbd4600
195+
03642806/fd54faba3bd95b6d6d59d46ccb8deb56

‎models/voting.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import cupy as cp
2-
3-
2+
import numpy as np
3+
44
ppf_kernel = cp.RawKernel(r'''
55
#include "/home/neil/ppf_matching/helper_math.cuh"
66
#define M_PI 3.14159265358979323846264338327950288
@@ -166,4 +166,4 @@
166166
outputs[idx] = diff_x + diff_y + diff_z;
167167
}
168168
}
169-
''', 'findpeak')
169+
''', 'findpeak')

‎nocs/eval.py‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,43 @@
1-
1+
import argparse
2+
import glob
3+
import os
4+
from tqdm import tqdm
5+
import pickle
6+
import numpy as np
7+
from utils.util import compute_degree_cm_mAP
8+
from inference import synset_names
29

310

411
if __name__ == '__main__':
5-
pass
12+
parser = argparse.ArgumentParser()
13+
parser.add_argument('--pred_dir', default='data/nocs_prediction', help='Directory for pose predictions')
14+
args = parser.parse_args()
15+
16+
result_pkl_list = glob.glob(os.path.join(args.pred_dir, 'results_*.pkl'))
17+
result_pkl_list = sorted(result_pkl_list)[:]
18+
assert len(result_pkl_list)
19+
20+
final_results = []
21+
for pkl_path in tqdm(result_pkl_list):
22+
with open(pkl_path, 'rb') as f:
23+
result = pickle.load(f)
24+
if not 'gt_handle_visibility' in result:
25+
result['gt_handle_visibility'] = np.ones_like(result['gt_class_ids'])
26+
print('can\'t find gt_handle_visibility in the pkl.')
27+
else:
28+
assert len(result['gt_handle_visibility']) == len(result['gt_class_ids']), "{} {}".format(result['gt_handle_visibility'], result['gt_class_ids'])
29+
30+
31+
if type(result) is list:
32+
final_results += result
33+
elif type(result) is dict:
34+
final_results.append(result)
35+
else:
36+
assert False
37+
38+
iou_3d_aps, pose_aps, pose_pred_matches, pose_gt_matches = compute_degree_cm_mAP(final_results, synset_names, args.pred_dir + '_map',
39+
degree_thresholds = [5, 10, 15],
40+
shift_thresholds= [5, 10, 15],
41+
iou_3d_thresholds=np.linspace(0, 1, 101),
42+
iou_pose_thres=0.1,
43+
use_matches_for_pose=True)

0 commit comments

Comments
 (0)