diff --git a/PoseEstimation/.gitkeep b/PoseEstimation/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/PoseEstimation/BOPdataset.md b/PoseEstimation/BOPdataset.md new file mode 100644 index 0000000..2043d6a --- /dev/null +++ b/PoseEstimation/BOPdataset.md @@ -0,0 +1,44 @@ +--- +id: BOP_dataset +title: script for create BOP dataset +--- + +## Структура входных данных: +``` +/ + input_obj/asm_element_edge.mtl # файл материала + input_obj/asm_element_edge.obj # меш-объект + input_obj/fork.mtl + input_obj/fork.obj + input_obj/... + resources/sklad.blend # файл сцены + objs2BOPdataset.py # этот скрипт +``` + +## Пример команды запуска скрипта: +``` +cd / +blenderproc run objs2BOPdataset.py resources/sklad.blend input_obj output --imgs 333 +``` +- resources/sklad.blend : файл сцены +- input_obj : каталог с меш-файлами +- output : выходной каталог +- imgs : количество пакетов по 9 кадров в каждом (в примере 333 * 9 = 2997) + +## Структура BOP датасета на выходе: +``` +output/ + bop_data/ + train_pbr/ + 000000/ + depth/... # файлы глубины + mask/... # файлы маски + mask_visib/... # файлы маски видимости + rgb/... # файлы изображений RGB + scene_camera.json + scene_gt.json + scene_gt_coco.json + scene_gt_info.json + camera.json # внутренние параметры камеры (для всего датасета) + res.txt # протокол создания пакетов датасета +``` diff --git a/PoseEstimation/objs2BOPdataset.py b/PoseEstimation/objs2BOPdataset.py new file mode 100644 index 0000000..64e55e5 --- /dev/null +++ b/PoseEstimation/objs2BOPdataset.py @@ -0,0 +1,261 @@ +import blenderproc as bproc +""" + objs2BOPdataset + Общая задача: распознавание 6D позы объекта (6D pose estimation) + Реализуемая функция: создание датасета в формате BOP для серии заданных объектов (*.obj) в заданной сцене (*.blend) + Используется модуль blenderproc + + 29.08.2023 @shalenikol release 0.1 + 12.10.2023 @shalenikol release 0.2 +""" +import sys +import numpy as np +import argparse +import random +import os +import shutil +import json + +Not_Categories_Name = True # наименование категории в аннотации отсутствует + +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 + +parser = argparse.ArgumentParser() +parser.add_argument('scene', nargs='?', default="resources/sklad.blend", help="Path to the scene object.") +parser.add_argument('obj_path', nargs='?', default="resources/in_obj", help="Path to the object files.") +parser.add_argument('output_dir', nargs='?', default="output", help="Path to where the final files, will be saved") +parser.add_argument('vhacd_path', nargs='?', default="blenderproc_resources/vhacd", help="The directory in which vhacd should be installed or is already installed.") +parser.add_argument('-single_object', nargs='?', type= bool, default=True, help="One object per frame.") +parser.add_argument('--imgs', default=2, type=int, help="The number of times the objects should be rendered.") +args = parser.parse_args() + +if not os.path.isdir(args.obj_path): + print(f"{args.obj_path} : no object directory") + sys.exit() + +if not os.path.isdir(args.output_dir): + os.mkdir(args.output_dir) + +single_object = args.single_object + +bproc.init() + +# ? загрузим свет из сцены +#cam = bproc.loader.load_blend(args.scene, data_blocks=["cameras"]) +#lights = bproc.loader.load_blend(args.scene, data_blocks=["lights"]) + +# загрузим объекты +list_files = os.listdir(args.obj_path) +obj_names = [] +obj_filenames = [] +all_meshs = [] +nObj = 0 +for f in list_files: + if (os.path.splitext(f))[1] == ".obj": + f = os.path.join(args.obj_path, f) # путь к файлу объекта + if os.path.isfile(f): + obj = bproc.loader.load_obj(f) + all_meshs += obj + obj_names += [obj[0].get_name()] + obj_filenames += [f] + nObj += 1 + +if nObj == 0: + print("Objects not found") + sys.exit() + +for i,obj in enumerate(all_meshs): + #print(f"{i} *** {obj}") + obj.set_cp("category_id", i+1) + +# загрузим сцену +scene = bproc.loader.load_blend(args.scene, data_blocks=["objects"]) + +# найдём объекты коллизии (пол и т.д.) +obj_type = ["floor", "obj"] +collision_objects = [] +#floor = None +for o in scene: + o.set_cp("category_id", 999) + s = o.get_name() + for type in obj_type: + if s.find(type) >= 0: + collision_objects += [o] + o.enable_rigidbody(False, collision_shape='BOX') +if not collision_objects: + print("Collision objects not found in the scene") + sys.exit() + +#floor.enable_rigidbody(False, collision_shape='BOX') + +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(args.vhacd_path) + +objs = all_meshs + scene + +with open(os.path.join(args.output_dir,"res.txt"), "w") as fh: +# fh.write(str(type(scene[0]))+"\n") + i = 0 + for o in objs: + i += 1 + 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 +light = bproc.types.Light() +light.set_type("POINT") +light.set_location([5, -5, 5]) +#light.set_energy(900) +#light.set_color([0.7, 0.7, 0.7]) + +light1 = bproc.types.Light(name="light1") +light1.set_type("SUN") +light1.set_location([0, 0, 0]) +light1.set_rotation_euler([-0.063, 0.6177, -0.1985]) +#light1.set_energy(7) +light1.set_color([1, 1, 1]) + +# define the camera intrinsics +bproc.camera.set_intrinsics_from_blender_params(1, 640, 480, lens_unit="FOV") + +# add segmentation masks (per class and per instance) +bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"]) +#bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name", "bop_dataset_name"], +# default_values={"category_id": 0, "bop_dataset_name": None}) + +# activate depth rendering +bproc.renderer.enable_depth_output(activate_antialiasing=False) + +res_dir = os.path.join(args.output_dir, "bop_data") +if os.path.isdir(res_dir): + shutil.rmtree(res_dir) +# Цикл рендеринга +n_cam_location = 3 #5 # количество случайных локаций камеры +n_cam_poses = 3 #3 # количество сэмплов для каждой локации камеры +# Do multiple times: Position the shapenet objects using the physics simulator and render X images with random camera poses +for r in range(args.imgs): + # один случайный объект в кадре / все заданные объекты + meshs = [random.choice(all_meshs)] if single_object else all_meshs[:] + + # Randomly set the color and energy + light.set_color(np.random.uniform([0.5, 0.5, 0.5], [1, 1, 1])) + light.set_energy(random.uniform(500, 1000)) + light1.set_energy(random.uniform(3, 11)) + + for i,o in enumerate(meshs): #objs + mat = o.get_materials()[0] + mat.set_principled_shader_value("Specular", random.uniform(0, 1)) + mat.set_principled_shader_value("Roughness", random.uniform(0, 1)) + mat.set_principled_shader_value("Base Color", np.random.uniform([0, 0, 0, 1], [1, 1, 1, 1])) + mat.set_principled_shader_value("Metallic", random.uniform(0, 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([-1, -1.5, 0.2], [1, 2, 1.2])) #[-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 + [floor], sample_pose_func = sample_pose) + bproc.object.sample_poses(meshs, objects_to_check_collisions = meshs + 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(os.path.join(args.output_dir,"res.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(n_cam_location): + # Sample location + location = bproc.sampler.shell(center=[0, 0, 0], + radius_min=1.1, + radius_max=2.2, + elevation_min=5, + elevation_max=89) + # координата, по которой будем сэмплировать положение камеры + j = random.randint(0, 2) + # разовый сдвиг по случайной координате + d = (coord_max[j] - coord_min[j]) / n_cam_poses + if location[j] < 0: + d = -d + for k in range(n_cam_poses): + # 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 + #world_matrix = bproc.math.build_transformation_mat([2.3, -0.4, 0.66], [1.396, 0., an]) + #bproc.camera.add_camera_pose(world_matrix) + #an += 0.2 + + # 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='JPEG', + append_to_existing_output = (r>0), + save_world2cam = False) # world coords are arbitrary in most real BOP datasets + # dataset="robo_ds", +""" + !!! 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] + +поэтому заменим наименование категории в аннотации +""" +if Not_Categories_Name: + coco_file = os.path.join(res_dir,"train_pbr/000000/scene_gt_coco.json") + with open(coco_file, "r") as fh: + data = json.load(fh) + cats = data["categories"] + #print(f"type(cat) = {type(cat)} cat : {cat}") + i = 0 + for cat in cats: + cat["name"] = obj_names[i] + i += 1 + #print(cat) + with open(coco_file, "w") as fh: + json.dump(data, fh, indent=0)