Skip to content

Commit 89aca6f

Browse files
mxh1999Tai-Wang
mxh1999
and
Tai-Wang
authored
[Feature] Preliminary Support for ArkitScenes Dataset (#75)
* Implement portable eval script for server * lint format * implement visualizer for debugging * Update README * docstring * Refine some docstrings. * refine docstring * refine docstring * update README for single gpu training/testing * fix: rotate yaw instead of pitch * fix: scale for euler 3d box * lint format * Minor fixes: 1) some info may have inaccurate tokens_positive 2) create_positive_maps in batch > 1 training * Visual test complete. Need Full pipline test * Pass Pipeline test * Lint * Arkitscenes data preliminary support --------- Co-authored-by: Tai-Wang <[email protected]>
1 parent e7b7f60 commit 89aca6f

File tree

6 files changed

+75
-30
lines changed

6 files changed

+75
-30
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ demo/data/*
125125
data/scannet
126126
data/3rscan
127127
data/matterport3d
128+
data/arkitscenes
128129
data/*.pkl
129130
data/*.json
130131
exps/

‎data/README.md‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ Detailed steps are shown as follows.
99

1010
3. Download Matterport3D data [HERE](https://github.com/niessner/Matterport). Link or move the folder to this level of directory.
1111

12-
4. Download EmbodiedScan data and extract it here. Currently, please fill in the [form](https://docs.google.com/forms/d/e/1FAIpQLScUXEDTksGiqHZp31j7Zp7zlCNV7p_08uViwP_Nbzfn3g6hhw/viewform?usp=sf_link), and we will reply with the data download link.
12+
4. Download ARKitScenes data [HERE](https://github.com/apple/ARKitScenes). Link or move the folder to this level of directory.
13+
14+
5. Download EmbodiedScan data and extract it here. Currently, please fill in the [form](https://docs.google.com/forms/d/e/1FAIpQLScUXEDTksGiqHZp31j7Zp7zlCNV7p_08uViwP_Nbzfn3g6hhw/viewform?usp=sf_link), and we will reply with the data download link.
1315

1416
The directory structure should be as below.
1517

@@ -25,17 +27,24 @@ data
2527
├── matterport3d
2628
│ ├── <scene_id>
2729
│ ├── ...
30+
├── arkitscenes
31+
│ ├── Training
32+
│ | ├── <scene_id>
33+
│ | ├── ...
34+
│ ├── Validation
35+
│ | ├── <scene_id>
36+
│ | ├── ...
2837
├── embodiedscan_occupancy
2938
├── embodiedscan_infos_train.pkl
3039
├── embodiedscan_infos_val.pkl
3140
├── embodiedscan_infos_test.pkl
32-
├── embodiedscan_infos_train_vg.json
33-
├── embodiedscan_infos_val_vg.json
34-
├── embodiedscan_infos_test_vg.json
35-
├── embodiedscan_infos_train_mini_vg.json (mini set)
36-
├── embodiedscan_infos_val_mini_vg.json (mini set)
37-
├── embodiedscan_infos_train_vg_all.json (w/ complex prompts)
38-
├── embodiedscan_infos_val_vg_all.json (w/ complex prompts)
41+
├── embodiedscan_train_vg.json
42+
├── embodiedscan_val_vg.json
43+
├── embodiedscan_test_vg.json
44+
├── embodiedscan_train_mini_vg.json (mini set)
45+
├── embodiedscan_val_mini_vg.json (mini set)
46+
├── embodiedscan_train_vg_all.json (w/ complex prompts)
47+
├── embodiedscan_val_vg_all.json (w/ complex prompts)
3948
```
4049

4150
5. Enter the project root directory, extract images by running
@@ -68,13 +77,21 @@ data
6877
├── matterport3d
6978
│ ├── <scene_id>
7079
│ ├── ...
80+
├── arkitscenes
81+
│ ├── Training
82+
│ | ├── <scene_id>
83+
│ | ├── ...
84+
│ ├── Validation
85+
│ | ├── <scene_id>
86+
│ | ├── ...
7187
├── embodiedscan_occupancy
72-
├── embodiedscan_infos_train_full.pkl
73-
├── embodiedscan_infos_val_full.pkl
74-
├── embodiedscan_infos_train_full_vg.json
75-
├── embodiedscan_infos_val_full_vg.json
76-
├── embodiedscan_infos_train_mini_vg.json
77-
├── embodiedscan_infos_val_mini_vg.json
88+
├── embodiedscan_infos_train.pkl
89+
├── embodiedscan_infos_val.pkl
90+
├── embodiedscan_infos_test.pkl
91+
├── embodiedscan_train_vg.json
92+
├── embodiedscan_val_vg.json
93+
├── embodiedscan_train_mini_vg.json
94+
├── embodiedscan_val_mini_vg.json
7895
```
7996

8097
6. Also extract EmbodiedScan occupancy annotations here by running

‎embodiedscan/datasets/embodiedscan_dataset.py‎

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -223,22 +223,33 @@ def parse_ann_info(self, info: dict) -> dict:
223223
mask_filename = os.path.join(self.data_prefix.get('img_path', ''),
224224
ann_dataset, building, 'occupancy',
225225
f'visible_occupancy_{region}.pkl')
226+
elif ann_dataset == 'arkitscenes':
227+
occ_filename = None
228+
mask_filename = None
226229
else:
227230
raise NotImplementedError
228231

229-
gt_occ = np.load(occ_filename)
230-
for i in range(gt_occ.shape[0]):
231-
cls_id = self.occ_label_mapping[gt_occ[i][3]]
232-
if cls_id < 0:
233-
cls_id = 255
234-
gt_occ[i][3] = cls_id
232+
if occ_filename is None:
233+
gt_occ = np.zeros((0, 4), dtype=np.int64)
234+
else:
235+
gt_occ = np.load(occ_filename)
236+
for i in range(gt_occ.shape[0]):
237+
cls_id = self.occ_label_mapping[gt_occ[i][3]]
238+
if cls_id < 0:
239+
cls_id = 255
240+
gt_occ[i][3] = cls_id
235241
ann_info['gt_occupancy'] = gt_occ
236242

237-
ann_info['visible_occupancy_masks'] = []
238-
occ_masks = mmengine.load(mask_filename)
239-
for i in range(len(info['images'])):
240-
ann_info['visible_occupancy_masks'].append(
241-
occ_masks[i]['visible_occupancy'])
243+
if mask_filename is None:
244+
ann_info['visible_occupancy_masks'] = [
245+
[] for i in range(len(info['images']))
246+
]
247+
else:
248+
ann_info['visible_occupancy_masks'] = []
249+
occ_masks = mmengine.load(mask_filename)
250+
for i in range(len(info['images'])):
251+
ann_info['visible_occupancy_masks'].append(
252+
occ_masks[i]['visible_occupancy'])
242253

243254
ann_info['gt_bboxes_3d'] = self.box_type_3d(
244255
ann_info['gt_bboxes_3d'],

‎embodiedscan/explorer.py‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from embodiedscan.visualization.img_drawer import ImageDrawer
1313
from embodiedscan.visualization.utils import _9dof_to_box, _box_add_thickness
1414

15-
DATASETS = ['scannet', '3rscan', 'matterport3d']
15+
DATASETS = ['scannet', '3rscan', 'matterport3d', 'arkitscenes']
1616

1717

1818
class EmbodiedScanExplorer:
@@ -65,7 +65,7 @@ def __init__(self,
6565
if self.verbose:
6666
print('Dataset root')
6767
for dataset in DATASETS:
68-
print(dataset, ':', self.data_root[dataset])
68+
print(dataset, ':', self.data_root.get(dataset, None))
6969

7070
if self.verbose:
7171
print('Loading')
@@ -118,6 +118,10 @@ def __init__(self,
118118
building, region = splits[1], splits[2]
119119
dirpath = os.path.join(self.data_root['matterport3d'],
120120
building)
121+
elif dataset == 'arkitscenes':
122+
split, region = splits[1], splits[2]
123+
dirpath = os.path.join(self.data_root['arkitscenes'],
124+
split, region)
121125
else:
122126
region = splits[1]
123127
dirpath = os.path.join(self.data_root[dataset], region)
@@ -168,6 +172,8 @@ def list_cameras(self, scene):
168172
elif dataset == 'matterport3d':
169173
cam_name = img_path.split(
170174
'/')[-1][:-8] + img_path.split('/')[-1][-7:-4]
175+
elif dataset == 'arkitscenes':
176+
cam_name = img_path.split('/')[-1][:-4]
171177
else:
172178
cam_name = img_path.split('/')[-1][:-4]
173179
res.append(cam_name)
@@ -252,6 +258,9 @@ def render_scene(self, scene_name, render_box=False):
252258
elif dataset == 'matterport3d':
253259
filepath = os.path.join(self.data_root['matterport3d'], building,
254260
'region_segmentations', f'{region}.ply')
261+
elif dataset == 'arkitscenes':
262+
filepath = os.path.join(self.data_root['arkitscenes'], building,
263+
region, f'{region}_3dod_mesh.ply')
255264
else:
256265
raise NotImplementedError
257266

@@ -311,6 +320,8 @@ def render_continuous_scene(self,
311320
elif dataset == 'matterport3d':
312321
cam_name = img_path.split(
313322
'/')[-1][:-8] + img_path.split('/')[-1][-7:-4]
323+
elif dataset == 'arkitscenes':
324+
cam_name = img_path.split('/')[-1][:-4]
314325
else:
315326
cam_name = img_path.split('/')[-1][:-4]
316327
if cam_name == start_cam:

‎embodiedscan/tutorial.ipynb‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,14 +60,14 @@
6060
"# # If there are multiple composed dataset and multiple annotation files\n",
6161
"# # this requires 'path/to/dataset' contain the name of composed dataset.\n",
6262
"# explorer = EmbodiedScanExplorer(\n",
63-
"# \tdataroot=['path/to/scannet', 'path/to/3rscan', 'path/to/matterport3d'],\n",
63+
"# \tdataroot=['path/to/scannet', 'path/to/3rscan', 'path/to/matterport3d', 'path/to/arkitscenes'],\n",
6464
"# \tann_file=['path/to/train.pkl', 'path/to/val.pkl'],\n",
6565
"# \tverbose=True,\t# print log or not\n",
6666
"# )\n",
6767
"\n",
6868
"# or\n",
6969
"explorer = EmbodiedScanExplorer(\n",
70-
"\tdata_root={'scannet' : 'demo/data/scannet', '3rscan' : 'demo/data/3rscan', 'matterport3d': 'demo/data/matterport3d'},\n",
70+
"\tdata_root={'scannet' : 'demo/data/scannet', '3rscan' : 'demo/data/3rscan', 'matterport3d': 'demo/data/matterport3d', 'arkitscenes': 'demo/data/arkitscenes'},\n",
7171
"\tann_file=['demo/data/train.pkl', 'demo/data/val.pkl'],\n",
7272
"\tverbose=True,\t# print log or not\n",
7373
")"

‎embodiedscan/visualization/continuous_drawer.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,9 @@ def begin(self):
6969
elif dataset == 'matterport3d':
7070
pcdpath = os.path.join(self.dir, building, 'region_segmentations',
7171
f'{region}.ply')
72+
elif dataset == 'arkitscenes':
73+
pcdpath = os.path.join(self.dir, building, region,
74+
f'{region}_3dod_mesh.ply')
7275
else:
7376
self.demo = True
7477
self.drawed_boxes = []
@@ -121,7 +124,9 @@ def draw_next(self, vis):
121124
if 'depth_cam2img' in img:
122125
depth_intrinsic = img['depth_cam2img']
123126
else:
124-
depth_intrinsic = self.scene['depth_cam2img']
127+
depth_intrinsic = self.scene.get('depth_cam2img', None)
128+
if depth_intrinsic is None:
129+
depth_intrinsic = intrinsic
125130
depth_shift = 1000.0
126131
if self.dataset == 'matterport3d':
127132
depth_shift = 4000.0

0 commit comments

Comments
 (0)