import omni.replicator.core as rep

# 環境の設定
rep.settings.set_stage_up_axis("Z") # 上軸をZ軸に設定
rep.settings.set_stage_meters_per_unit(1.0) # 1単位をメートルに設定

# カメラの作成
camera = rep.create.camera(clipping_range=(0.001, 100000))
render_product = rep.create.render_product(camera, resolution=(512, 512))

# RGB画像とバウンディングボックスアノテーションをfigureフォルダに保存
writer = rep.WriterRegistry.get("BasicWriter")
writer.initialize(output_dir="figure", rgb=True,bounding_box_2d_tight=True)
writer.attach([render_product])

# ディストラクタアセットの作成
can = rep.create.from_usd("omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Props/YCB/Axis_Aligned/002_master_chef_can.usd")
drill = rep.create.from_usd("omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Props/YCB/Axis_Aligned/035_power_drill.usd")
scissors = rep.create.from_usd("omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Props/YCB/Axis_Aligned/037_scissors.usd")
banana = rep.create.from_usd("omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Props/YCB/Axis_Aligned/011_banana.usd")
cracker_box = rep.create.from_usd("omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Props/YCB/Axis_Aligned/003_cracker_box.usd")

# 背景用のplaneの作成
plane = rep.create.plane(scale=100)

# 人形アセットの取得
figure = rep.get.prims(semantics = [("class", "figure")])

# 500回分の撮影                          
with rep.trigger.on_frame(num_frames=500, rt_subframes=30):

# カメラのランダム配置
    with camera:
        rep.modify.pose(
            position=rep.distribution.uniform((-1, -1, 0.5), (1, 1, 2.0)), # カメラのランダム配置
            look_at=(0, 0, 0)   # カメラの注視点
        )

# 背景画像の切り替え
    with plane:
        rep.randomizer.texture(
            textures=rep.distribution.sequence([
                "omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Materials/Textures/Patterns/nv_bamboo_desktop.jpg",
                "omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Materials/Textures/Patterns/nv_brick_tile.jpg",
                "omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Materials/Textures/Patterns/nv_plastic_foam.jpg",
                "omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Materials/Textures/Patterns/nv_tile_hexagonal_green_white.jpg",
                "omniverse://localhost/NVIDIA/Assets/Isaac/4.2/Isaac/Materials/Textures/Patterns/nv_wood_fence_nonuniform.jpg"
                ]),
            texture_scale=[(100, 100)]
            )
        
#　人形アセットのランダム配置
    with figure:
        rep.modify.pose(
            rotation_z=rep.distribution.uniform(0, 360), # 人形のランダムな回転
        )
        rep.randomizer.scatter_2d(
            rep.create.plane(scale=1.0,visible=False), # 不可視のplaneを作成し、その表面にディストラクタをランダム配置
            check_for_collisions = True, # ディストラクタ同士が重ならないようにする
            )

# ディストラクタのランダム配置
    with rep.create.group([can, drill, scissors, banana, cracker_box]):
        rep.modify.pose(
            rotation=rep.distribution.uniform((0, 0, 0),(360,360,360))) # ディストラクタのランダムな回転
        rep.randomizer.scatter_2d(
            rep.create.plane(scale=1.3,visible=False), # 不可視のplaneを作成し、その表面にディストラクタをランダム配置
            check_for_collisions = True, # ディストラクタ同士が重ならないようにする
            no_coll_prims=[figure], # ディストラクタと人形アセットが重ならないようにする
            offset = 0.2, # ディストラクタが床にめり込まないようにする為、planeの表面から0.2m離れた距離でランダム配置するようにする。
            )