chensong1995 / HybridPose

HybridPose: 6D Object Pose Estimation under Hybrid Representation (CVPR 2020)
MIT License
412 stars 64 forks source link

creating custom synthetic data to be used for HybridPose following Linemod format of BOP challenge using BlenderProc #91

Open monajalal opened 9 months ago

monajalal commented 9 months ago

Hi Chen,

I want to create synthetic dataset following linemod format for a custom CAD model of my own so I can train hybridpose on it from scratch. I am not sure how to do this and my search so far in the internet brought me to blenderproc given pvnet-rendering is not supported or viable anymore in Ubuntu 22.04. My object is also rather much bigger than BOP challenge objects (it has a diameter of 1.5 meter).

I have this script for randomizing one single object (it is not perfect and doesn't create so many variation in the object pose unfortunately).

Do you have a BlenderProc script or a modification of below scripts that you could share with me so that I could train HybridPose with it? Also, how many images would be needed?

(blenderproc) mona@ada:~/BlenderProc/examples/datasets/bop_challenge$ cat main_lm_uprightcustom.py 
import blenderproc as bproc
import argparse
import os
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('bop_parent_path', default='models_obj', help="Path to the bop datasets parent directory")
parser.add_argument('cc_textures_path', default="/home/mona/BlenderProc/examples/datasets/bop_challenge/texture_dir", help="Path to downloaded cc textures")
parser.add_argument('output_dir', default='/home/mona/BlenderProc/examples/datasets/bop_challenge/output_data', help="Path to where the final files will be saved ")
parser.add_argument('--num_scenes', type=int, default=200, help="How many scenes with 25 images each to generate")
args = parser.parse_args()

bproc.init()

# SIZE=10
SIZE = 1
# 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'), mm2m=True,
#                                             num_of_objs_to_sample=SIZE, obj_ids=[1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
target_bop_objs = bproc.loader.load_bop_objs(bop_dataset_path = os.path.join(args.bop_parent_path, 'lm'), mm2m=True,
                                            num_of_objs_to_sample=SIZE, obj_ids=[1])
# # 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', mm2m = True)
# ycbv_dist_bop_objs = bproc.loader.load_bop_objs(bop_dataset_path = os.path.join(args.bop_parent_path, 'ycbv'), mm2m = True)
# tyol_dist_bop_objs = bproc.loader.load_bop_objs(bop_dataset_path = os.path.join(args.bop_parent_path, 'tyol'), mm2m = True)

# 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])]

# create room
room_planes = [bproc.object.create_primitive('PLANE', scale=[10, 10, 1]),
               bproc.object.create_primitive('PLANE', scale=[10, 10, 1], location=[0, -10, 10], rotation=[-1.570796, 0, 0]),
               bproc.object.create_primitive('PLANE', scale=[10, 10, 1], location=[0, 10, 10], rotation=[1.570796, 0, 0]),
               bproc.object.create_primitive('PLANE', scale=[10, 10, 1], location=[10, 0, 10], rotation=[0, -1.570796, 0]),
               bproc.object.create_primitive('PLANE', scale=[10, 10, 1], location=[-10, 0, 10], rotation=[0, 1.570796, 0])]

# sample light color and strenght from ceiling
light_plane = bproc.object.create_primitive('PLANE', scale=[12, 12, 1], location=[0, 0, 14])
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(500)

# 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])
    # min = np.random.uniform([-1, 1, 0.0], [-2, 2, 0.0])
    # max = np.random.uniform([1, 1, 2], [2, 2, 4])
    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(100)

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=SIZE, 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.0,0.0,0.0],[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.1,
                                                          max_tries=1000,
                                                          )

    # 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 < 30:
        # Sample location
        location = bproc.sampler.shell(center = [0, 0, 0],
                                radius_min = 3,
                                radius_max = 3.1,
                                elevation_min = 80,
                                elevation_max = 88)
                                # 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=SIZE, 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))
        rotation_matrix = bproc.camera.rotation_from_forward_vec(poi - location, inplane_rot=np.random.uniform(-1.396, 1.396))

        # 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.1}, 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()
    nocs_data = bproc.renderer.render_nocs()
    bproc.writer.write_hdf5("output_nocs", nocs_data, append_to_existing_output=True)

    # 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 = 1.0,
                           depths = data["depth"],
                           colors = data["colors"], 
                           color_file_format = "JPEG",
                           ignore_dist_thres = 10)

    for obj in (sampled_target_bop_objs):# + sampled_distractor_bop_objs):      
        obj.hide(True)

I run it using:

(blenderproc) mona@ada:~/BlenderProc/examples/datasets/bop_challenge$ blenderproc run main_lm_uprightcustom.py models_obj/ texture_dir/ output_data --num_scenes=5 --custom-blender-path ~/Downloads/blender-3.5.1-linux-x64

I have Ubuntu 22.04 and I use $ python Python 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] on linux

chensong1995 commented 9 months ago

Hi Mona,

Thanks for your question! Unfortunately, I do not have a BlenderProc to share with you. I simply used the rendering scripts from PVNet. My suggestion is to get an old version of Ubuntu and try the PVNet rendering script, too. I hope this helps! Let me know if you have further concerns.