robossembler/framework
7
0
Fork 0
Code Issues 21 Pull requests Projects Releases Packages Wiki Activity Actions

add dataset generation script

Browse source
This commit is contained in:
shalenikol 2024-04-26 08:58:52 +03:00
parent 946e83fd15
commit 9a7d6b9084
3 changed files with 417 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

15
simulation/dataset_generation/README.md Normal file
View file

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

41
simulation/dataset_generation/dataset_cfg.json Normal file
View file

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

361
simulation/dataset_generation/renderBOPdataset.py Executable file
View file

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