/
run_dataset_recording.py
167 lines (140 loc) · 4.93 KB
/
run_dataset_recording.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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
import argparse
from colorama import Fore, Style
from focalpose.config import LOCAL_DATA_DIR
from focalpose.recording.record_dataset import record_dataset
def make_cfg(cfg_name,
resume_ds_name='',
debug=False,
distributed=False,
overwrite=False,
datasets_dir=LOCAL_DATA_DIR):
datasets_dir = datasets_dir / 'synt_datasets'
datasets_dir.mkdir(exist_ok=True)
cfg = argparse.ArgumentParser('').parse_args([])
cfg.overwrite = overwrite
cfg.ds_name = 'default_dataset'
n_frames = 1e6
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.train_ratio = 0.95
cfg.distributed = distributed
cfg.n_workers = 12
cfg.n_processes_per_gpu = 10
cfg.scene_cls = 'focalpose.recording.bop_recording_scene.BopRecordingScene'
cfg.scene_kwargs = dict(
gpu_renderer=True,
texture_ds='shapenet',
domain_randomization=True,
n_objects_interval=(1, 1),
proba_falling=0.0,
border_check=True,
n_textures_cache=100,
objects_xyz_interval=((-0.15, -0.15, 0.), (0.15, 0.15, 0.)),
camera_distance_interval=(0.8, 2.4),
focal_interval=(200, 1000),
resolution=(640, 640),
textures_on_objects=True,
)
cfg.ds_name = f'{cfg_name}-1M'
if cfg_name == 'pix3d-sofa':
n_frames = 1e6
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds='pix3d-sofa',
)
elif cfg_name == 'pix3d-bed':
n_frames = 1e6
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds='pix3d-bed',
)
elif cfg_name == 'pix3d-table':
n_frames = 1e6
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds='pix3d-table',
)
elif 'pix3d-chair' in cfg_name:
cfg.scene_kwargs['camera_distance_interval'] = (0.8, 3.4)
n_frames = 50000
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds=cfg_name,
)
elif cfg_name == 'pix3d':
n_frames = 1e6
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds='pix3d',
)
elif 'stanfordcars' in cfg_name:
n_frames = 75000
cfg.scene_kwargs['camera_distance_interval'] = (0.8, 3.0)
cfg.scene_kwargs['resolution'] = (300, 200)
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds=cfg_name,
)
elif 'compcars' in cfg_name:
n_frames = 100000
cfg.scene_kwargs['camera_distance_interval'] = (0.8, 3.0)
cfg.scene_kwargs['resolution'] = (300, 200)
cfg.n_frames_per_chunk = 100
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
cfg.ds_name = f'{cfg_name}-1M'
cfg.scene_kwargs.update(
urdf_ds=cfg_name,
)
elif resume_ds_name:
pass
else:
raise ValueError('Unknown config')
if debug:
n_frames = 10
cfg.overwrite = True
cfg.ds_name = 'debug'
cfg.n_frames_per_chunk = 1
cfg.n_chunks = n_frames // cfg.n_frames_per_chunk
if resume_ds_name:
cfg.resume = datasets_dir / resume_ds_name
cfg.ds_name = resume_ds_name
assert cfg.resume.exists()
else:
cfg.resume = ''
cfg.ds_dir = datasets_dir / cfg.ds_name
return cfg
def main():
parser = argparse.ArgumentParser('Dataset recording')
parser.add_argument('--config', default='', type=str)
parser.add_argument('--resume', default='', type=str)
parser.add_argument('--debug', action='store_true')
parser.add_argument('--local', action='store_true')
parser.add_argument('--overwrite', action='store_true')
args = parser.parse_args()
print(f"{Fore.RED}using config {args.config} {Style.RESET_ALL}")
cfg = make_cfg(args.config,
resume_ds_name=args.resume,
debug=args.debug,
distributed=not args.local,
overwrite=args.overwrite)
for k, v in vars(cfg).items():
print(k, v)
if cfg.resume:
print(f"RESUMING {Fore.RED} {cfg.ds_name} {Style.RESET_ALL} \n ")
else:
print(f"STARTING DATASET RECORDING {Fore.GREEN} {cfg.ds_name} {Style.RESET_ALL} \n ")
record_dataset(cfg)
if __name__ == '__main__':
main()