-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain_lm_upright.py
236 lines (208 loc) · 8.01 KB
/
main_lm_upright.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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
# From https://github.com/DLR-RM/BlenderProc/blob/main/examples/datasets/bop_challenge/main_lm_upright.py
import blenderproc as bproc
import argparse
import os
import numpy as np
parser = argparse.ArgumentParser()
parser.add_argument("bop_parent_path", help="Path to the bop datasets parent directory")
parser.add_argument(
"cc_textures_path",
default="resources/cctextures",
help="Path to downloaded cc textures",
)
parser.add_argument("output_dir", help="Path to where the final files will be saved ")
parser.add_argument(
"--views_per_scene",
type=int,
default=5,
help="How many camera views to generate.",
)
parser.add_argument(
"--num_scenes",
type=int,
default=5,
help="How many scenes to generate",
)
args = parser.parse_args()
bproc.init()
# load bop objects into the scene
target_bop_objs = bproc.loader.load_bop_objs(
bop_dataset_path=os.path.join(args.bop_parent_path, "lm"), object_model_unit="mm"
)
# load distractor bop objects
tless_dist_bop_objs = bproc.loader.load_bop_objs(
bop_dataset_path=os.path.join(args.bop_parent_path, "tless"),
model_type="cad",
object_model_unit="mm",
)
ycbv_dist_bop_objs = bproc.loader.load_bop_objs(
bop_dataset_path=os.path.join(args.bop_parent_path, "ycbv"), object_model_unit="mm"
)
tyol_dist_bop_objs = bproc.loader.load_bop_objs(
bop_dataset_path=os.path.join(args.bop_parent_path, "tyol"), object_model_unit="mm"
)
# load BOP datset intrinsics
bproc.loader.load_bop_intrinsics(
bop_dataset_path=os.path.join(args.bop_parent_path, "lm")
)
# set shading and hide objects
for obj in (
target_bop_objs + tless_dist_bop_objs + ycbv_dist_bop_objs + tyol_dist_bop_objs
):
obj.set_shading_mode("auto")
obj.hide(True)
# create room
room_planes = [
bproc.object.create_primitive("PLANE", scale=[2, 2, 1]),
bproc.object.create_primitive(
"PLANE", scale=[2, 2, 1], location=[0, -2, 2], rotation=[-1.570796, 0, 0]
),
bproc.object.create_primitive(
"PLANE", scale=[2, 2, 1], location=[0, 2, 2], rotation=[1.570796, 0, 0]
),
bproc.object.create_primitive(
"PLANE", scale=[2, 2, 1], location=[2, 0, 2], rotation=[0, -1.570796, 0]
),
bproc.object.create_primitive(
"PLANE", scale=[2, 2, 1], location=[-2, 0, 2], rotation=[0, 1.570796, 0]
),
]
# sample light color and strenght from ceiling
light_plane = bproc.object.create_primitive(
"PLANE", scale=[3, 3, 1], location=[0, 0, 10]
)
light_plane.set_name("light_plane")
light_plane_material = bproc.material.create("light_material")
# sample point light on shell
light_point = bproc.types.Light()
light_point.set_energy(200)
# load cc_textures
cc_textures = bproc.loader.load_ccmaterials(args.cc_textures_path)
# Define a function that samples 6-DoF poses
def sample_pose_func(obj: bproc.types.MeshObject):
min = np.random.uniform([-0.3, -0.3, 0.0], [-0.2, -0.2, 0.0])
max = np.random.uniform([0.2, 0.2, 0.4], [0.3, 0.3, 0.6])
obj.set_location(np.random.uniform(min, max))
obj.set_rotation_euler(bproc.sampler.uniformSO3())
# activate depth rendering without antialiasing and set amount of samples for color rendering
bproc.renderer.enable_depth_output(activate_antialiasing=False)
bproc.renderer.set_max_amount_of_samples(50)
for i in range(args.num_scenes):
# Sample bop objects for a scene
sampled_target_bop_objs = list(
np.random.choice(target_bop_objs, size=15, replace=False)
)
sampled_distractor_bop_objs = list(
np.random.choice(tless_dist_bop_objs, size=3, replace=False)
)
sampled_distractor_bop_objs += list(
np.random.choice(ycbv_dist_bop_objs, size=3, replace=False)
)
sampled_distractor_bop_objs += list(
np.random.choice(tyol_dist_bop_objs, size=3, replace=False)
)
# Randomize materials and set physics
for obj in sampled_target_bop_objs + sampled_distractor_bop_objs:
mat = obj.get_materials()[0]
if obj.get_cp("bop_dataset_name") in ["itodd", "tless"]:
grey_col = np.random.uniform(0.1, 0.9)
mat.set_principled_shader_value(
"Base Color", [grey_col, grey_col, grey_col, 1]
)
mat.set_principled_shader_value("Roughness", np.random.uniform(0, 1.0))
mat.set_principled_shader_value("Specular", np.random.uniform(0, 1.0))
obj.hide(False)
# Sample two light sources
light_plane_material.make_emissive(
emission_strength=np.random.uniform(3, 6),
emission_color=np.random.uniform([0.5, 0.5, 0.5, 1.0], [1.0, 1.0, 1.0, 1.0]),
)
light_plane.replace_materials(light_plane_material)
light_point.set_color(np.random.uniform([0.5, 0.5, 0.5], [1, 1, 1]))
location = bproc.sampler.shell(
center=[0, 0, 0],
radius_min=1,
radius_max=1.5,
elevation_min=5,
elevation_max=89,
)
light_point.set_location(location)
# sample CC Texture and assign to room planes
random_cc_texture = np.random.choice(cc_textures)
for plane in room_planes:
plane.replace_materials(random_cc_texture)
# Sample object poses and check collisions
bproc.object.sample_poses(
objects_to_sample=sampled_target_bop_objs + sampled_distractor_bop_objs,
sample_pose_func=sample_pose_func,
max_tries=1000,
)
# Define a function that samples the initial pose of a given object above the ground
def sample_initial_pose(obj: bproc.types.MeshObject):
obj.set_location(
bproc.sampler.upper_region(
objects_to_sample_on=room_planes[0:1],
min_height=1,
max_height=4,
face_sample_range=[0.4, 0.6],
)
)
obj.set_rotation_euler(np.random.uniform([0, 0, 0], [0, 0, np.pi * 2]))
# Sample objects on the given surface
placed_objects = bproc.object.sample_poses_on_surface(
objects_to_sample=sampled_target_bop_objs + sampled_distractor_bop_objs,
surface=room_planes[0],
sample_pose_func=sample_initial_pose,
min_distance=0.01,
max_distance=0.2,
)
# BVH tree used for camera obstacle checks
bop_bvh_tree = bproc.object.create_bvh_tree_multi_objects(
sampled_target_bop_objs + sampled_distractor_bop_objs
)
cam_poses = 0
while cam_poses < args.views_per_scene:
# Sample location
location = bproc.sampler.shell(
center=[0, 0, 0],
radius_min=0.35,
radius_max=1.5,
elevation_min=5,
elevation_max=89,
)
# Determine point of interest in scene as the object closest to the mean of a subset of objects
poi = bproc.object.compute_poi(
np.random.choice(sampled_target_bop_objs, size=10, replace=False)
)
# Compute rotation based on vector going from location towards poi
rotation_matrix = bproc.camera.rotation_from_forward_vec(
poi - location, inplane_rot=np.random.uniform(-0.7854, 0.7854)
)
# Add homog cam pose based on location an rotation
cam2world_matrix = bproc.math.build_transformation_mat(
location, rotation_matrix
)
# Check that obstacles are at least 0.3 meter away from the camera and make sure the view interesting enough
if bproc.camera.perform_obstacle_in_view_check(
cam2world_matrix, {"min": 0.3}, bop_bvh_tree
):
# Persist camera pose
bproc.camera.add_camera_pose(cam2world_matrix, frame=cam_poses)
cam_poses += 1
# render the whole pipeline
data = bproc.renderer.render()
# Write data in bop format
bproc.writer.write_bop(
os.path.join(args.output_dir, "bop_data"),
target_objects=sampled_target_bop_objs,
dataset="lm",
depth_scale=0.1,
depths=data["depth"],
colors=data["colors"],
color_file_format="JPEG",
ignore_dist_thres=10,
frames_per_chunk=args.views_per_scene,
num_worker=6,
)
for obj in sampled_target_bop_objs + sampled_distractor_bop_objs:
obj.hide(True)