diff --git a/simulation/dataset_generation/README.md b/simulation/dataset_generation/README.md new file mode 100644 index 0000000..1d00a6f --- /dev/null +++ b/simulation/dataset_generation/README.md @@ -0,0 +1,15 @@ +## Скрипт генерации датасета + +Скрипт используется в составе web-сервиса для генерации датасетов с использованием заданной пользователем конфигурации. +Должен быть установлен пакет [BlenderProc](https://github.com/DLR-RM/BlenderProc). + +Команда для вызова: +```bash +blenderproc run renderBOPdataset.py --cfg CFG + +options: + --cfg CFG строка json с параметрами конфигурации датасета / путь к json-файлу с конфигурацией +``` + +[Пример файла конфигурации датасета.](dataset_cfg.json) + diff --git a/simulation/dataset_generation/dataset_cfg.json b/simulation/dataset_generation/dataset_cfg.json new file mode 100644 index 0000000..8456392 --- /dev/null +++ b/simulation/dataset_generation/dataset_cfg.json @@ -0,0 +1,41 @@ +{ + "dataSetObjects": ["fork"], + "datasetType": "Object Detection - YOLOv8", + "name": "123123e", + "formBuilder": { + "output": { + "typedataset": "ObjectDetection", + "dataset_path": "eqwfeadszxz", + "models": [{"id": 1, "name": "fork"}], + "models_randomization": { "loc_range_low": [-1, -1, 0.0], "loc_range_high": [1, 1, 2] }, + "scene": { + "objects": [ + {"name": "floor", "collision_shape": "BOX", "loc_xyz":[0,0,0], "rot_euler":[0, 0, 0], + "material_randomization": {"specular":[0,1], "roughness":[0,1], "metallic":[0,1], "base_color":[[0,0,0,1],[1,1,1,1]]} + } + ], + "lights": [ + {"id": 1, "type": "POINT", "loc_xyz":[5,5,5], "rot_euler":[-0.06, 0.61, -0.19], + "color_range_low":[0.5, 0.5, 0.5], "color_range_high":[1, 1, 1], + "energy_range":[400,900] + }, + {"id": 2, "type": "SUN", "loc_xyz":[0,0,0], "rot_euler":[-0.01, 0.01, -0.01], + "color_range_low":[1, 1, 1], "color_range_high":[1, 1, 1], + "energy_range":[2,9] + } + ] + }, + "camera_position": { "center_shell": [0, 0, 0], "radius_range": [0.4, 1.4], "elevation_range": [10, 90] }, + "generation": { + "n_cam_pose": 3, + "n_sample_on_pose": 1, + "n_series": 3, + "image_format": "JPEG", + "image_size_wh": [640, 480] + } + } + }, + "processStatus": "exec", + "local_path": "/home/user/5f4e161b-82d1-41fa-a11c-15d485b01600", + "projectId": "660aaddbf98957a186f9c546" +} diff --git a/simulation/dataset_generation/renderBOPdataset.py b/simulation/dataset_generation/renderBOPdataset.py new file mode 100755 index 0000000..db616fa --- /dev/null +++ b/simulation/dataset_generation/renderBOPdataset.py @@ -0,0 +1,361 @@ +import blenderproc as bproc +""" + renderBOPdataset + Общая задача: common pipeline + Реализуемая функция: создание датасета в формате BOP с заданными параметрами рандомизации + Используется модуль blenderproc + + 26.04.2024 @shalenikol release 0.1 +""" +import numpy as np +import argparse +import random +import os +import shutil +import json + +VHACD_PATH = "blenderproc_resources/vhacd" +DIR_MODELS = "models" +FILE_LOG_SCENE = "res.txt" +FILE_RBS_INFO = "rbs_info.json" +FILE_GT_COCO = "scene_gt_coco.json" + +Not_Categories_Name = True # наименование категории в COCO-аннотации отсутствует + +def _get_path_model(name_model: str) -> str: + # TODO on name_model find path for mesh (model.fbx) + # local_path/assets/mesh/ + return os.path.join(rnd_par.output_dir, "assets/mesh/"+name_model+".fbx") + +def _get_path_object(name_obj: str) -> str: + # TODO on name_obj find path for scene object (object.fbx) + return os.path.join(rnd_par.output_dir, "assets/mesh/"+name_obj+".fbx") + +def convert2relative(height, width, bbox): + """ + YOLO format use relative coordinates for annotation + """ + x, y, w, h = bbox + x += w/2 + y += h/2 + return x/width, y/height, w/width, h/height + +def render() -> int: + for obj in all_meshs: + # Make the object actively participate in the physics simulation + obj.enable_rigidbody(active=True, collision_shape="COMPOUND") + # Also use convex decomposition as collision shapes + obj.build_convex_decomposition_collision_shape(VHACD_PATH) + + objs = all_meshs + rnd_par.scene.objs + + log_txt = os.path.join(rnd_par.output_dir, FILE_LOG_SCENE) + with open(log_txt, "w") as fh: + for i,o in enumerate(objs): + loc = o.get_location() + euler = o.get_rotation_euler() + fh.write(f"{i} : {o.get_name()} {loc} {euler} category_id = {o.get_cp('category_id')}\n") + + # define a light and set its location and energy level + ls = [] + for l in rnd_par.scene.light_data: + light = bproc.types.Light(name=f"l{l['id']}") + light.set_type(l["type"]) + light.set_location(l["loc_xyz"]) #[5, -5, 5]) + light.set_rotation_euler(l["rot_euler"]) #[-0.063, 0.6177, -0.1985]) + ls += [light] + + # define the camera intrinsics + bproc.camera.set_intrinsics_from_blender_params(1, + rnd_par.image_size_wh[0], + rnd_par.image_size_wh[1], + lens_unit="FOV") + + # add segmentation masks (per class and per instance) + bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"]) + + # activate depth rendering + bproc.renderer.enable_depth_output(activate_antialiasing=False) + + res_dir = os.path.join(rnd_par.output_dir, rnd_par.ds_name) + if os.path.isdir(res_dir): + shutil.rmtree(res_dir) + # Цикл рендеринга + # Do multiple times: Position the shapenet objects using the physics simulator and render X images with random camera poses + for r in range(rnd_par.n_series): + # один случайный объект в кадре / все заданные объекты + random_obj = random.choice(range(rnd_par.scene.n_obj)) + meshs = [] + for i,o in enumerate(all_meshs): #objs + if rnd_par.single_object and i != random_obj: + continue + meshs += [o] + rnd_mat = rnd_par.scene.obj_data[i]["material_randomization"] + mats = o.get_materials() #[0] + for mat in mats: + val = rnd_mat["specular"] + mat.set_principled_shader_value("Specular", random.uniform(val[0], val[1])) + val = rnd_mat["roughness"] + mat.set_principled_shader_value("Roughness", random.uniform(val[0], val[1])) + val = rnd_mat["base_color"] + mat.set_principled_shader_value("Base Color", np.random.uniform(val[0], val[1])) + val = rnd_mat["metallic"] + mat.set_principled_shader_value("Metallic", random.uniform(val[0], val[1])) + + # Randomly set the color and energy + for i,l in enumerate(ls): + current = rnd_par.scene.light_data[i] + l.set_color(np.random.uniform(current["color_range_low"], current["color_range_high"])) + energy = current["energy_range"] + l.set_energy(random.uniform(energy[0], energy[1])) + + # Clear all key frames from the previous run + bproc.utility.reset_keyframes() + + # Define a function that samples 6-DoF poses + def sample_pose(obj: bproc.types.MeshObject): + obj.set_location(np.random.uniform(rnd_par.loc_range_low, rnd_par.loc_range_high)) #[-1, -1, 0], [1, 1, 2])) + obj.set_rotation_euler(bproc.sampler.uniformSO3()) + + # Sample the poses of all shapenet objects above the ground without any collisions in-between + bproc.object.sample_poses(meshs, + objects_to_check_collisions = meshs + rnd_par.scene.collision_objects, + sample_pose_func = sample_pose) + + # Run the simulation and fix the poses of the shapenet objects at the end + bproc.object.simulate_physics_and_fix_final_poses(min_simulation_time=4, max_simulation_time=20, check_object_interval=1) + + # Find point of interest, all cam poses should look towards it + poi = bproc.object.compute_poi(meshs) + + coord_max = [0.1, 0.1, 0.1] + coord_min = [0., 0., 0.] + + with open(log_txt, "a") as fh: + fh.write("*****************\n") + fh.write(f"{r}) poi = {poi}\n") + i = 0 + for o in meshs: + i += 1 + loc = o.get_location() + euler = o.get_rotation_euler() + fh.write(f" {i} : {o.get_name()} {loc} {euler}\n") + for j in range(3): + if loc[j] < coord_min[j]: + coord_min[j] = loc[j] + if loc[j] > coord_max[j]: + coord_max[j] = loc[j] + + # Sample up to X camera poses + #an = np.random.uniform(0.78, 1.2) #1. #0.35 + for i in range(rnd_par.n_cam_pose): + # Sample location + location = bproc.sampler.shell(center=rnd_par.center_shell, + radius_min=rnd_par.radius_range[0], + radius_max=rnd_par.radius_range[1], + elevation_min=rnd_par.elevation_range[0], + elevation_max=rnd_par.elevation_range[1]) + # координата, по которой будем сэмплировать положение камеры + j = random.randint(0, 2) + # разовый сдвиг по случайной координате + d = (coord_max[j] - coord_min[j]) / rnd_par.n_sample_on_pose + if location[j] < 0: + d = -d + for _ in range(rnd_par.n_sample_on_pose): + # 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) + bproc.camera.add_camera_pose(cam2world_matrix) + location[j] -= d + # render the whole pipeline + data = bproc.renderer.render() + # Write data to bop format + bproc.writer.write_bop(res_dir, + target_objects = all_meshs, # Optional[List[MeshObject]] = None + depths = data["depth"], + depth_scale = 1.0, + colors = data["colors"], + color_file_format=rnd_par.image_format, + append_to_existing_output = (r>0), + save_world2cam = False) # world coords are arbitrary in most real BOP datasets + # dataset="robo_ds", + + models_dir = os.path.join(res_dir, DIR_MODELS) + os.mkdir(models_dir) + + data = [] + for i,objn in enumerate(rnd_par.models.names): + rec = {} + rec["id"] = i+1 + rec["name"] = objn + rec["model"] = os.path.join(DIR_MODELS, os.path.split(rnd_par.models.filenames[i])[1]) # путь относительный + t = [obj.get_bound_box(local_coords=True).tolist() for obj in all_meshs if obj.get_name() == objn] + rec["cuboid"] = t[0] + data.append(rec) + shutil.copy2(rnd_par.models.filenames[i], models_dir) + f = (os.path.splitext(rnd_par.models.filenames[i]))[0] + ".mtl" # файл материала + if os.path.isfile(f): + shutil.copy2(f, models_dir) + + with open(os.path.join(res_dir, FILE_RBS_INFO), "w") as fh: + json.dump(data, fh, indent=2) + + """ + !!! categories -> name берётся из category_id !!! + см.ниже + blenderproc.python.writer : BopWriterUtility.py + class _BopWriterUtility + def calc_gt_coco + ... + CATEGORIES = [{'id': obj.get_cp('category_id'), 'name': str(obj.get_cp('category_id')), 'supercategory': + dataset_name} for obj in dataset_objects] + + поэтому заменим наименование категории в аннотации + """ + def change_categories_name(dir: str): + coco_file = os.path.join(dir,FILE_GT_COCO) + with open(coco_file, "r") as fh: + data = json.load(fh) + cats = data["categories"] + + for i,cat in enumerate(cats): + cat["name"] = rnd_par.models.names[i] #obj_names[i] + + with open(coco_file, "w") as fh: + json.dump(data, fh, indent=0) + + def explore(path: str): + if not os.path.isdir(path): + return + folders = [ + os.path.join(path, o) + for o in os.listdir(path) + if os.path.isdir(os.path.join(path, o)) + ] + for path_entry in folders: + print(path_entry) + if os.path.isfile(os.path.join(path_entry,FILE_GT_COCO)): + change_categories_name(path_entry) + else: + explore(path_entry) + + if Not_Categories_Name: + explore(res_dir) + return 0 # success + +def _get_models(par, data) -> int: + global all_meshs + + par.models = lambda: None + par.models.n_item = len(data) + if par.models.n_item == 0: + return 0 # no models + + # загрузим объекты + par.models.names = [] # obj_names + par.models.filenames = [] # obj_filenames + i = 1 + for f in data: + nam = f + par.models.names.append(nam) + ff = _get_path_model(nam) + par.models.filenames.append(ff) + if not os.path.isfile(ff): + print(f"Error: no such file '{ff}'") + return -1 + obj = bproc.loader.load_obj(ff) + all_meshs += obj + obj[0].set_cp("category_id", i) # начиная с 1 + i += 1 + return par.models.n_item + +def _get_scene(par, data) -> int: + # load scene + par.scene = lambda: None + objs = data["objects"] + par.scene.n_obj = len(objs) + if par.scene.n_obj == 0: + return 0 # empty scene + lights = data["lights"] + par.scene.n_light = len(lights) + if par.scene.n_light == 0: + return 0 # no lighting + + par.scene.objs = [] + par.scene.collision_objects = [] + for f in objs: + ff = _get_path_object(f["name"]) + if not os.path.isfile(ff): + print(f"Error: no such file '{ff}'") + return -1 + obj = bproc.loader.load_obj(ff) + obj[0].set_cp("category_id", 999) + coll = f["collision_shape"] + if len(coll) > 0: + obj[0].enable_rigidbody(False, collision_shape=coll) + par.scene.collision_objects += obj + par.scene.objs += obj + + if not par.scene.collision_objects: + print("Collision objects not found in the scene") + return 0 + par.scene.obj_data = objs + par.scene.light_data = lights + return par.scene.n_obj + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--cfg", required=True, help="Json-string with dataset parameters") + args = parser.parse_args() + + if args.cfg[-5:] == ".json": + if not os.path.isfile(args.cfg): + print(f"Error: no such file '{args.cfg}'") + exit(-1) + with open(args.cfg, "r") as f: + j_data = f.read() + else: + j_data = args.cfg + try: + cfg = json.loads(j_data) + except json.JSONDecodeError as e: + print(f"JSon error: {e}") + exit(-2) + + ds_cfg = cfg["formBuilder"]["output"] # dataset config + generation = ds_cfg["generation"] + cam_pos = ds_cfg["camera_position"] + models_randomization = ds_cfg["models_randomization"] + + rnd_par = lambda: None + rnd_par.single_object = True + rnd_par.ds_name = cfg["name"] + rnd_par.output_dir = cfg["local_path"] + rnd_par.dataset_objs = cfg["dataSetObjects"] + rnd_par.n_cam_pose = generation["n_cam_pose"] + rnd_par.n_sample_on_pose = generation["n_sample_on_pose"] + rnd_par.n_series = generation["n_series"] + rnd_par.image_format = generation["image_format"] + rnd_par.image_size_wh = generation["image_size_wh"] + rnd_par.center_shell = cam_pos["center_shell"] + rnd_par.radius_range = cam_pos["radius_range"] + rnd_par.elevation_range = cam_pos["elevation_range"] + rnd_par.loc_range_low = models_randomization["loc_range_low"] + rnd_par.loc_range_high = models_randomization["loc_range_high"] + + if not os.path.isdir(rnd_par.output_dir): + print(f"Error: invalid path '{rnd_par.output_dir}'") + exit(-3) + + bproc.init() + + all_meshs = [] + ret = _get_models(rnd_par, rnd_par.dataset_objs) + if ret <= 0: + print("Error: no models in config") + exit(-4) + if _get_scene(rnd_par, ds_cfg["scene"]) == 0: + print("Error: empty scene in config") + exit(-5) + exit(render()) \ No newline at end of file