pyrox.dev
1# based on https://github.com/DIYer22/bpycv/blob/c576e01622d87eb3534f73bf1a5686bd2463de97/example/ycb_demo.py
2import bpy
3import bpycv
4
5import os
6import glob
7import random
8from pathlib import Path
9
10example_data_dir = os.environ['BPY_EXAMPLE_DATA']
11out_dir = Path(os.environ['out'])
12out_dir.mkdir(parents=True, exist_ok=True)
13
14models = sorted(glob.glob(os.path.join(example_data_dir, "model", "*", "*.obj")))
15cat_id_to_model_path = dict(enumerate(sorted(models), 1))
16
17distractors = sorted(glob.glob(os.path.join(example_data_dir, "distractor", "*.obj")))
18
19bpycv.clear_all()
20bpy.context.scene.frame_set(1)
21bpy.context.scene.render.engine = "CYCLES"
22bpy.context.scene.cycles.samples = 32
23bpy.context.scene.render.resolution_y = 1024
24bpy.context.scene.render.resolution_x = 1024
25bpy.context.view_layer.cycles.denoising_store_passes = False
26
27# A transparency stage for holding rigid body
28stage = bpycv.add_stage(transparency=True)
29
30bpycv.set_cam_pose(cam_radius=1, cam_deg=45)
31
32hdri_dir = os.path.join(example_data_dir, "background_and_light")
33hdri_manager = bpycv.HdriManager(
34 hdri_dir=hdri_dir, download=False
35) # if download is True, will auto download .hdr file from HDRI Haven
36hdri_path = hdri_manager.sample()
37bpycv.load_hdri_world(hdri_path, random_rotate_z=True)
38
39# load 5 objects
40for index in range(5):
41 cat_id = random.choice(list(cat_id_to_model_path))
42 model_path = cat_id_to_model_path[cat_id]
43 obj = bpycv.load_obj(model_path)
44 obj.location = (
45 random.uniform(-0.2, 0.2),
46 random.uniform(-0.2, 0.2),
47 random.uniform(0.1, 0.3),
48 )
49 obj.rotation_euler = [random.uniform(-3.1415, 3.1415) for _ in range(3)]
50 # set each instance a unique inst_id, which is used to generate instance annotation.
51 obj["inst_id"] = cat_id * 1000 + index
52 with bpycv.activate_obj(obj):
53 bpy.ops.rigidbody.object_add()
54
55# load 6 distractors
56for index in range(6):
57 distractor_path = random.choice(distractors)
58 target_size = random.uniform(0.1, 0.3)
59 distractor = bpycv.load_distractor(distractor_path, target_size=target_size)
60 distractor.location = (
61 random.uniform(-0.2, 0.2),
62 random.uniform(-0.2, 0.2),
63 random.uniform(0.1, 0.3),
64 )
65 distractor.rotation_euler = [random.uniform(-3.1415, 3.1415) for _ in range(3)]
66 with bpycv.activate_obj(distractor):
67 bpy.ops.rigidbody.object_add()
68
69# run pyhsic engine for 20 frames
70for i in range(20):
71 bpy.context.scene.frame_set(bpy.context.scene.frame_current + 1)
72
73# render image, instance annoatation and depth in one line code
74result = bpycv.render_data()
75
76
77result.save(dataset_dir=str(out_dir.resolve()), fname="0", save_blend=True)
78print(f'Save to "{out_dir}"')
79print(f'Open "{out_dir}/vis/" to see visualize result.')
80