399 lines
No EOL
16 KiB
Python
Executable file
399 lines
No EOL
16 KiB
Python
Executable file
import blenderproc as bproc
|
||
"""
|
||
renderBOPdataset
|
||
Общая задача: common pipeline
|
||
Реализуемая функция: создание датасета в формате BOP с заданными параметрами рандомизации
|
||
Используется модуль blenderproc
|
||
|
||
02.05.2024 @shalenikol release 0.1
|
||
02.07.2024 @shalenikol release 0.2
|
||
"""
|
||
import numpy as np
|
||
import argparse
|
||
import random
|
||
import os
|
||
import shutil
|
||
import json
|
||
from pathlib import Path
|
||
|
||
import bpy
|
||
|
||
VHACD_PATH = "blenderproc_resources/vhacd"
|
||
DIR_MODELS = "models"
|
||
DIR_MESH = "assets/libs/objects/" #"assets/mesh/"
|
||
FILE_LOG_SCENE = "res.txt"
|
||
FILE_RBS_INFO = "rbs_info.json"
|
||
FILE_GT_COCO = "scene_gt_coco.json"
|
||
EXT_MODELS = ".fbx"
|
||
TEXTURE_TMPL = "*.jpg"
|
||
|
||
Not_Categories_Name = True # наименование категории в COCO-аннотации отсутствует
|
||
|
||
def _get_list_texture(rel_path: str) -> list:
|
||
# local_path/texture/
|
||
loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
|
||
path = os.path.join(loc, rel_path)
|
||
return list(Path(path).absolute().rglob(TEXTURE_TMPL))
|
||
|
||
def _get_path_model(name_model: str) -> str:
|
||
# TODO on name_model find path for mesh (model.fbx)
|
||
# local_path/assets/libs/objects # assets/mesh/
|
||
loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
|
||
return os.path.join(loc, DIR_MESH + name_model + EXT_MODELS)
|
||
|
||
def _get_path_object(name_obj: str) -> str:
|
||
# TODO on name_obj find path for scene object (object.fbx)
|
||
loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
|
||
return os.path.join(loc, DIR_MESH + name_obj + EXT_MODELS)
|
||
|
||
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:
|
||
i = 0
|
||
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)
|
||
i += 1
|
||
# print(f"{i} : {obj.get_name()}")
|
||
|
||
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):
|
||
print(f"********** Series : {r+1}")
|
||
is_texture = True if "texture_path" in rnd_par.models_randomization else False
|
||
if is_texture:
|
||
val = rnd_par.models_randomization["texture_path"]
|
||
l_texture = _get_list_texture(val)
|
||
image = bpy.data.images.load(filepath=str(l_texture[r % len(l_texture)]))
|
||
# один случайный объект в кадре / все заданные объекты
|
||
random_obj = random.choice(range(rnd_par.models.n_item))
|
||
meshs = []
|
||
for i,o in enumerate(all_meshs): # активные модели
|
||
if rnd_par.single_object and i != random_obj:
|
||
continue
|
||
meshs += [o]
|
||
if is_texture:
|
||
mats = o.get_materials()
|
||
for mat in mats:
|
||
# image = bpy.data.images.load(filepath=str(random.choice(l_texture)))
|
||
mat.set_principled_shader_value("Base Color", image)
|
||
|
||
for i,o in enumerate(rnd_par.scene.objs): # объекты сцены
|
||
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["metallic"]
|
||
mat.set_principled_shader_value("Metallic", random.uniform(val[0], val[1]))
|
||
if "texture_path" in rnd_mat: # путь к текстурам (*.jpg)
|
||
val = rnd_mat["texture_path"]
|
||
val = _get_list_texture(val)
|
||
image = bpy.data.images.load(filepath=str(random.choice(val)))
|
||
mat.set_principled_shader_value("Base Color", image)
|
||
else:
|
||
val = rnd_mat["base_color"]
|
||
mat.set_principled_shader_value("Base Color", np.random.uniform(val[0], val[1]))
|
||
# mat.set_principled_shader_value("Base Color", image)
|
||
|
||
# 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=1)
|
||
|
||
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 = False # 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.models_randomization = models_randomization
|
||
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 = []
|
||
if _get_models(rnd_par, rnd_par.dataset_objs) <= 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()) |