add dataset generation script

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)
+

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"
+}

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())