pkgs/development/python-modules/bpycv/bpycv-test.py

   1 # based on https://github.com/DIYer22/bpycv/blob/c576e01622d87eb3534f73bf1a5686bd2463de97/example/ycb_demo.py
   2 import bpy
   3 import bpycv
   4
   5 import os
   6 import glob
   7 import random
   8 from pathlib import Path
   9
  10 example_data_dir = os.environ['BPY_EXAMPLE_DATA']
  11 out_dir = Path(os.environ['out'])
  12 out_dir.mkdir(parents=True, exist_ok=True)
  13
  14 models = sorted(glob.glob(os.path.join(example_data_dir, "model", "*", "*.obj")))
  15 cat_id_to_model_path = dict(enumerate(sorted(models), 1))
  16
  17 distractors = sorted(glob.glob(os.path.join(example_data_dir, "distractor", "*.obj")))
  18
  19 bpycv.clear_all()
  20 bpy.context.scene.frame_set(1)
  21 bpy.context.scene.render.engine = "CYCLES"
  22 bpy.context.scene.cycles.samples = 32
  23 bpy.context.scene.render.resolution_y = 1024
  24 bpy.context.scene.render.resolution_x = 1024
  25 bpy.context.view_layer.cycles.denoising_store_passes = False
  26
  27 # A transparency stage for holding rigid body
  28 stage = bpycv.add_stage(transparency=True)
  29
  30 bpycv.set_cam_pose(cam_radius=1, cam_deg=45)
  31
  32 hdri_dir = os.path.join(example_data_dir, "background_and_light")
  33 hdri_manager = bpycv.HdriManager(
  34     hdri_dir=hdri_dir, download=False
  35 )  # if download is True, will auto download .hdr file from HDRI Haven
  36 hdri_path = hdri_manager.sample()
  37 bpycv.load_hdri_world(hdri_path, random_rotate_z=True)
  38
  39 # load 5 objects
  40 for 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
  56 for 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
  70 for 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
  74 result = bpycv.render_data()
  75
  76
  77 result.save(dataset_dir=str(out_dir.resolve()), fname="0", save_blend=True)
  78 print(f'Save to "{out_dir}"')
  79 print(f'Open "{out_dir}/vis/" to see visualize result.')
  80