dataset_generation and train_models as modules
This commit is contained in:
parent
69b8512d6b
commit
649a426c25
20 changed files with 412 additions and 4342 deletions
412
dataset_generation/renderBOPdataset2.py
Executable file
412
dataset_generation/renderBOPdataset2.py
Executable file
|
@ -0,0 +1,412 @@
|
|||
import blenderproc as bproc
|
||||
"""
|
||||
renderBOPdataset2
|
||||
Общая задача: common pipeline
|
||||
Реализуемая функция: создание датасета в формате BOP с заданными параметрами рандомизации
|
||||
Используется модуль blenderproc
|
||||
|
||||
02.05.2024 @shalenikol release 0.1
|
||||
02.07.2024 @shalenikol release 0.2
|
||||
28.10.2024 @shalenikol release 0.3
|
||||
"""
|
||||
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/"
|
||||
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)
|
||||
return os.path.join(rnd_par.details_dir, 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(os.path.dirname(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)
|
||||
res_dir = rnd_par.output_dir
|
||||
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["name"]
|
||||
par.models.names.append(nam)
|
||||
ff = f["fbx"] # _get_path_model(nam)
|
||||
par.models.filenames.append(ff)
|
||||
if not os.path.isfile(ff):
|
||||
print(f"Error: no such file '{ff}'")
|
||||
return -1
|
||||
# !!! dir with meshs
|
||||
par.details_dir = os.path.split(ff)[0]
|
||||
|
||||
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
|
||||
if len(rnd_par.details_dir) == 0:
|
||||
return 0 # no path to details
|
||||
|
||||
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("--form", required=True, help="Json-string with dataset parameters")
|
||||
parser.add_argument("--path", required=True, help="Output path")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.form[-5:] == ".json":
|
||||
if not os.path.isfile(args.form):
|
||||
print(f"Error: no such file '{args.form}'")
|
||||
exit(-1)
|
||||
with open(args.form, "r") as f:
|
||||
j_data = f.read()
|
||||
else:
|
||||
j_data = args.form
|
||||
try:
|
||||
cfg = json.loads(j_data)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"JSon error: {e}")
|
||||
exit(-2)
|
||||
|
||||
# output_dir = args.path
|
||||
|
||||
ds_cfg = cfg["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.output_dir = args.path # cfg["local_path"]
|
||||
|
||||
if not os.path.isdir(rnd_par.output_dir):
|
||||
print(f"Error: invalid path '{rnd_par.output_dir}'")
|
||||
exit(-3)
|
||||
|
||||
rnd_par.single_object = False # True
|
||||
rnd_par.details_dir = ""
|
||||
# rnd_par.ds_name = os.path.split(rnd_par.output_dir)[1] # cfg["name"]
|
||||
rnd_par.dataset_objs = ds_cfg["datasetObjects"]["details"] # ["knight"]
|
||||
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"]
|
||||
|
||||
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())
|
|
@ -1,105 +0,0 @@
|
|||
# Инструкция для запуска
|
||||
|
||||
Должен быть установлен пакет [BlenderProc](https://github.com/DLR-RM/BlenderProc)
|
||||
|
||||
## Создание датасета в формате YoloV4 для заданного объекта
|
||||
|
||||
Команда для запуска:
|
||||
|
||||
```
|
||||
blenderproc run obj2Yolov4dataset.py [obj] [output_dir] [--imgs 1]
|
||||
```
|
||||
- obj: файл описания объекта *.obj
|
||||
- output_dir: выходной каталог
|
||||
- --imgs 1: количество изображений на выходе
|
||||
|
||||
## Создание датасета в формате YoloV4 для серии заданных объектов в заданной сцене
|
||||
|
||||
Команда для запуска:
|
||||
```
|
||||
blenderproc run objs2Yolov4dataset.py [scene] [obj_path] [output_dir] [vhacd_path] [--imgs 1]
|
||||
```
|
||||
- scene: путь к файлу описания сцены (*.blend)
|
||||
- obj_path: путь к каталогу с файлами описания детектируемых объектов *.obj
|
||||
- output_dir: выходной каталог
|
||||
- vhacd_path: каталог, в котором должен быть установлен или уже установлен vhacd (по умолчанию blenderproc_resources/vhacd)
|
||||
- --imgs 1: количество серий рендеринга (по 15 изображений каждая) на выходе (например, если imgs=100, то будет получено 1500 изображений)
|
||||
|
||||
Файл описания сцены обязательно должен содержать плоскость (с именем 'floor'), на которую будут сэмплированы объекты для обнаружения.
|
||||
|
||||
Должен быть собран пакет [darknet](https://github.com/AlexeyAB/darknet) для работы на заданном ПО и оборудовании (CPU, GPU ...)
|
||||
|
||||
---
|
||||
|
||||
## Обучение нейросети и получение файла с её весами
|
||||
|
||||
Команда для запуска:
|
||||
```
|
||||
darknet detector train [data] [cfg] [weight]
|
||||
```
|
||||
- data: файл с описанием датасета (*.data)
|
||||
- cfg: файл с описанием нейросети
|
||||
- weight: файл весов нейросети
|
||||
|
||||
Для обучения нужно загрузить файл с предобученными весами (162 MB): [yolov4.conv.137](https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.conv.137)
|
||||
Для разного количества детектируемых объектов в выборке нужны свои файлы [data](https://gitlab.com/robossembler/framework/-/blob/master/ObjectDetection/yolov4_objs2.data) и [cfg](https://gitlab.com/robossembler/framework/-/blob/master/ObjectDetection/yolov4_objs2.cfg).
|
||||
|
||||
---
|
||||
|
||||
## Команда для обнаружения объектов нейросетью с обученными весами
|
||||
* вариант 1 (в файле t.txt - список изображений):
|
||||
```
|
||||
darknet detector test yolov4_objs2.data yolov4_test.cfg yolov4_objs2_final.weights -dont_show -ext_output < t.txt > res.txt
|
||||
```
|
||||
|
||||
* вариант 2 (файл 000015.jpg - тестовое изображение):
|
||||
```
|
||||
darknet detector test yolov4_objs2.data yolov4_test.cfg yolov4_objs2_final.weights -dont_show -ext_output 000015.jpg > res.txt
|
||||
```
|
||||
* вариант 3 (в файле t.txt - список изображений):
|
||||
```
|
||||
darknet detector test yolov4_objs2.data yolov4_test.cfg yolov4_objs2_final.weights -dont_show -ext_output -out res.json < t.txt
|
||||
```
|
||||
|
||||
Файл res.txt после запуска варианта 2:
|
||||
|
||||
> net.optimized_memory = 0
|
||||
> mini_batch = 1, batch = 1, time_steps = 1, train = 0
|
||||
> Create CUDA-stream - 0
|
||||
> Create cudnn-handle 0
|
||||
> nms_kind: greedynms (1), beta = 0.600000
|
||||
> nms_kind: greedynms (1), beta = 0.600000
|
||||
> nms_kind: greedynms (1), beta = 0.600000
|
||||
>
|
||||
> seen 64, trained: 768 K-images (12 Kilo-batches_64)
|
||||
> Detection layer: 139 - type = 28
|
||||
> Detection layer: 150 - type = 28
|
||||
> Detection layer: 161 - type = 28
|
||||
>000015.jpg: Predicted in 620.357000 milli-seconds.
|
||||
>fork.001: 94% (left_x: 145 top_y: -0 width: 38 height: 18)
|
||||
>asm_element_edge.001: 28% (left_x: 195 top_y: 320 width: 40 height: 61)
|
||||
>start_link.001: 87% (left_x: 197 top_y: 313 width: 39 height: 68)
|
||||
>doking_link.001: 99% (left_x: 290 top_y: 220 width: 32 height: 21)
|
||||
>start_link.001: 90% (left_x: 342 top_y: 198 width: 33 height: 34)
|
||||
>doking_link.001: 80% (left_x: 342 top_y: 198 width: 32 height: 34)
|
||||
>assemb_link.001: 100% (left_x: 426 top_y: 410 width: 45 height: 61)
|
||||
|
||||
|
||||
Файл res.json после запуска варианта 3:
|
||||
>[
|
||||
{
|
||||
"frame_id":1,
|
||||
"filename":"img_test/000001.jpg",
|
||||
"objects": [
|
||||
{"class_id":5, "name":"asm_element_edge.001", "relative_coordinates":{"center_x":0.498933, "center_y":0.502946, "width":0.083075, "height":0.073736}, "confidence":0.999638},
|
||||
{"class_id":4, "name":"grip-tool.001", "relative_coordinates":{"center_x":0.858856, "center_y":0.031339, "width":0.043919, "height":0.064563}, "confidence":0.996551}
|
||||
]
|
||||
},
|
||||
{
|
||||
"frame_id":2,
|
||||
"filename":"img_test/000002.jpg",
|
||||
"objects": [
|
||||
{"class_id":1, "name":"start_link.001", "relative_coordinates":{"center_x":0.926026, "center_y":0.728457, "width":0.104029, "height":0.132757}, "confidence":0.995811},
|
||||
{"class_id":0, "name":"assemb_link.001", "relative_coordinates":{"center_x":0.280403, "center_y":0.129059, "width":0.029980, "height":0.025067}, "confidence":0.916782}
|
||||
]
|
||||
}
|
|
@ -1,144 +0,0 @@
|
|||
import blenderproc as bproc
|
||||
"""
|
||||
obj2Yolov4dataset
|
||||
Общая задача: обнаружение объекта (Object detection)
|
||||
Реализуемая функция: создание датасета в формате YoloV4 для заданного объекта (*.obj)
|
||||
Используется модуль blenderproc
|
||||
|
||||
24.01.2023 @shalenikol release 0.1
|
||||
22.02.2023 @shalenikol release 0.2 исправлен расчёт x,y в convert2relative
|
||||
"""
|
||||
import numpy as np
|
||||
import argparse
|
||||
import random
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
|
||||
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/robossembler-asset.obj", help="Path to the object file.")
|
||||
parser.add_argument('output_dir', nargs='?', default="output", help="Path to where the final files, will be saved")
|
||||
parser.add_argument('--imgs', default=1, type=int, help="The number of times the objects should be rendered.")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.isdir(args.output_dir):
|
||||
os.mkdir(args.output_dir)
|
||||
|
||||
bproc.init()
|
||||
|
||||
# load the objects into the scene
|
||||
obj = bproc.loader.load_obj(args.scene)[0]
|
||||
obj.set_cp("category_id", 1)
|
||||
|
||||
# Randomly perturbate the material of the object
|
||||
mat = obj.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))
|
||||
|
||||
# Create a new light
|
||||
light = bproc.types.Light()
|
||||
light.set_type("POINT")
|
||||
# Sample its location around the object
|
||||
light.set_location(bproc.sampler.shell(
|
||||
center=obj.get_location(),
|
||||
radius_min=1,
|
||||
radius_max=5,
|
||||
elevation_min=1,
|
||||
elevation_max=89
|
||||
))
|
||||
# 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(100, 1000))
|
||||
|
||||
bproc.camera.set_resolution(640, 480)
|
||||
|
||||
# Sample five camera poses
|
||||
poses = 0
|
||||
tries = 0
|
||||
while tries < 10000 and poses < args.imgs:
|
||||
# Sample random camera location around the object
|
||||
location = bproc.sampler.shell(
|
||||
center=obj.get_location(),
|
||||
radius_min=1,
|
||||
radius_max=4,
|
||||
elevation_min=1,
|
||||
elevation_max=89
|
||||
)
|
||||
# Compute rotation based lookat point which is placed randomly around the object
|
||||
lookat_point = obj.get_location() + np.random.uniform([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5])
|
||||
rotation_matrix = bproc.camera.rotation_from_forward_vec(lookat_point - 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)
|
||||
|
||||
# Only add camera pose if object is still visible
|
||||
if obj in bproc.camera.visible_objects(cam2world_matrix):
|
||||
bproc.camera.add_camera_pose(cam2world_matrix)
|
||||
poses += 1
|
||||
tries += 1
|
||||
|
||||
# Enable transparency so the background becomes transparent
|
||||
bproc.renderer.set_output_format(enable_transparency=True)
|
||||
# add segmentation masks (per class and per instance)
|
||||
bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"])
|
||||
|
||||
# Render RGB images
|
||||
data = bproc.renderer.render()
|
||||
|
||||
# Write data to coco file
|
||||
res_dir = os.path.join(args.output_dir, 'coco_data')
|
||||
bproc.writer.write_coco_annotations(res_dir,
|
||||
instance_segmaps=data["instance_segmaps"],
|
||||
instance_attribute_maps=data["instance_attribute_maps"],
|
||||
color_file_format='JPEG',
|
||||
colors=data["colors"],
|
||||
append_to_existing_output=True)
|
||||
|
||||
#загрузим аннотацию
|
||||
with open(os.path.join(res_dir,"coco_annotations.json"), "r") as fh:
|
||||
y = json.load(fh)
|
||||
|
||||
# список имен объектов
|
||||
with open(os.path.join(res_dir,"obj.names"), "w") as fh:
|
||||
for cat in y["categories"]:
|
||||
fh.write(cat["name"]+"\n")
|
||||
|
||||
# содадим или очистим папку data для датасета
|
||||
res_data = os.path.join(res_dir, 'data')
|
||||
if os.path.isdir(res_data):
|
||||
for f in os.listdir(res_data):
|
||||
os.remove(os.path.join(res_data, f))
|
||||
else:
|
||||
os.mkdir(res_data)
|
||||
|
||||
# список имен файлов с изображениями
|
||||
s = []
|
||||
with open(os.path.join(res_dir,"images.txt"), "w") as fh:
|
||||
for i in y["images"]:
|
||||
filename = i["file_name"]
|
||||
shutil.copy(os.path.join(res_dir,filename),res_data)
|
||||
fh.write(filename.replace('images','data')+"\n")
|
||||
s.append((os.path.split(filename))[1])
|
||||
|
||||
# предполагается, что "images" и "annotations" следуют в одном и том же порядке
|
||||
c = 0
|
||||
for i in y["annotations"]:
|
||||
bbox = i["bbox"]
|
||||
im_h = i["height"]
|
||||
im_w = i["width"]
|
||||
rel = convert2relative(im_h,im_w,bbox)
|
||||
fn = (os.path.splitext(s[c]))[0] # только имя файла
|
||||
with open(os.path.join(res_data,fn+".txt"), "w") as fh:
|
||||
# формат: <target> <x-center> <y-center> <width> <height>
|
||||
fh.write("0 "+'{:-f} {:-f} {:-f} {:-f}'.format(rel[0],rel[1],rel[2],rel[3])+"\n")
|
||||
c += 1
|
|
@ -1,296 +0,0 @@
|
|||
import blenderproc as bproc
|
||||
"""
|
||||
objs2Yolov4dataset
|
||||
Общая задача: обнаружение объекта (Object detection)
|
||||
Реализуемая функция: создание датасета в формате YoloV4 для серии заданных объектов (*.obj) в заданной сцене (*.blend)
|
||||
Используется модуль blenderproc
|
||||
|
||||
17.02.2023 @shalenikol release 0.1
|
||||
22.02.2023 @shalenikol release 0.2 исправлен расчёт x,y в convert2relative
|
||||
"""
|
||||
import sys
|
||||
import numpy as np
|
||||
import argparse
|
||||
import random
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
|
||||
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('--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)
|
||||
|
||||
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)
|
||||
meshs = []
|
||||
i = 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):
|
||||
meshs += bproc.loader.load_obj(f)
|
||||
i += 1
|
||||
|
||||
if i == 0:
|
||||
print("Objects not found")
|
||||
sys.exit()
|
||||
|
||||
for i,o in enumerate(meshs):
|
||||
o.set_cp("category_id", i+1)
|
||||
|
||||
# загрузим сцену
|
||||
scene = bproc.loader.load_blend(args.scene, data_blocks=["objects"])
|
||||
#scene = bproc.loader.load_obj(args.scene)
|
||||
|
||||
# найдём пол
|
||||
floor = None
|
||||
for o in scene:
|
||||
o.set_cp("category_id", 999)
|
||||
s = o.get_name()
|
||||
if s.find("floor") >= 0:
|
||||
floor = o
|
||||
if floor == None:
|
||||
print("Floor not found in the scene")
|
||||
sys.exit()
|
||||
|
||||
floor.enable_rigidbody(False, collision_shape='BOX')
|
||||
|
||||
objs = meshs + scene
|
||||
|
||||
for obj in 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)
|
||||
|
||||
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}\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])
|
||||
"""
|
||||
# Sample its location around the object
|
||||
light.set_location(bproc.sampler.shell(
|
||||
center=obj.get_location(),
|
||||
radius_min=2.5,
|
||||
radius_max=5,
|
||||
elevation_min=1,
|
||||
elevation_max=89
|
||||
))
|
||||
"""
|
||||
|
||||
# define the camera intrinsics
|
||||
bproc.camera.set_intrinsics_from_blender_params(1, 640, 480, lens_unit="FOV")
|
||||
bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"])
|
||||
|
||||
res_dir = os.path.join(args.output_dir, 'coco_data')
|
||||
# Цикл рендеринга
|
||||
n_cam_location = 5 # количество случайных локаций камеры
|
||||
n_cam_poses = 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):
|
||||
# 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(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)
|
||||
|
||||
# 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=3.3,
|
||||
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 coco file
|
||||
bproc.writer.write_coco_annotations(res_dir,
|
||||
instance_segmaps=data["instance_segmaps"],
|
||||
instance_attribute_maps=data["instance_attribute_maps"],
|
||||
color_file_format='JPEG',
|
||||
colors=data["colors"],
|
||||
append_to_existing_output=True)
|
||||
|
||||
#загрузим аннотацию
|
||||
with open(os.path.join(res_dir,"coco_annotations.json"), "r") as fh:
|
||||
y = json.load(fh)
|
||||
|
||||
# список имен объектов
|
||||
n_obj = 0
|
||||
obj_list = []
|
||||
with open(os.path.join(res_dir,"obj.names"), "w") as fh:
|
||||
for cat in y["categories"]:
|
||||
if cat["id"] < 999:
|
||||
n = cat["name"]
|
||||
i = cat["id"]
|
||||
obj_list.append([n,i,n_obj])
|
||||
fh.write(n+"\n")
|
||||
n_obj += 1
|
||||
|
||||
# содадим или очистим папку data для датасета
|
||||
res_data = os.path.join(res_dir, 'data')
|
||||
if os.path.isdir(res_data):
|
||||
for f in os.listdir(res_data):
|
||||
os.remove(os.path.join(res_data, f))
|
||||
else:
|
||||
os.mkdir(res_data)
|
||||
|
||||
# список имен файлов с изображениями
|
||||
fn_image = os.path.join(res_dir,"images.txt")
|
||||
img_list = []
|
||||
with open(fn_image, "w") as fh:
|
||||
for i in y["images"]:
|
||||
filename = i["file_name"]
|
||||
shutil.copy(os.path.join(res_dir,filename),res_data)
|
||||
fh.write(filename.replace('images','data')+"\n")
|
||||
img_list.append([i["id"], (os.path.split(filename))[1]])
|
||||
|
||||
# создадим 2 списка имен файлов для train и valid
|
||||
n_image_in_series = n_cam_location * n_cam_poses # количество изображений в серии
|
||||
i = 0
|
||||
fh = open(fn_image, "r")
|
||||
f1 = open(os.path.join(res_dir,"i_train.txt"), "w")
|
||||
f2 = open(os.path.join(res_dir,"i_val.txt"), "w")
|
||||
for line in fh:
|
||||
i += 1
|
||||
if i % n_image_in_series == 0:
|
||||
f2.write(line)
|
||||
else:
|
||||
f1.write(line)
|
||||
fh.close()
|
||||
f1.close()
|
||||
f2.close()
|
||||
|
||||
# заполним файлы с метками bbox
|
||||
for i in y["annotations"]:
|
||||
cat_id = i["category_id"]
|
||||
if cat_id < 999:
|
||||
im_id = i["image_id"]
|
||||
bbox = i["bbox"]
|
||||
im_h = i["height"]
|
||||
im_w = i["width"]
|
||||
rel = convert2relative(im_h,im_w,bbox)
|
||||
|
||||
# находим индекс списка с нужным изображением
|
||||
j = next(k for k, (x, _) in enumerate(img_list) if x == im_id)
|
||||
filename = img_list[j][1]
|
||||
fn = (os.path.splitext(filename))[0] # только имя файла
|
||||
with open(os.path.join(res_data,fn+".txt"), "a") as fh:
|
||||
# находим индекс списка с нужным объектом
|
||||
j = next(k for k, (_, x, _) in enumerate(obj_list) if x == cat_id)
|
||||
# формат: <target> <x-center> <y-center> <width> <height>
|
||||
fh.write(f"{obj_list[j][2]} {rel[0]} {rel[1]} {rel[2]} {rel[3]}\n")
|
||||
|
||||
# создадим файл описания датасета для darknet
|
||||
with open(os.path.join(res_dir,"yolov4_objs2.data"), "w") as fh:
|
||||
fh.write(f"classes = {n_obj}\n")
|
||||
fh.write("train = i_train.txt\n")
|
||||
fh.write("valid = i_val.txt\n")
|
||||
fh.write("names = obj.names\n")
|
||||
fh.write("backup = backup\n")
|
||||
fh.write("eval = coco\n")
|
File diff suppressed because it is too large
Load diff
|
@ -1,7 +0,0 @@
|
|||
classes= 1
|
||||
train = i_train.txt
|
||||
valid = i_val.txt
|
||||
names = obj.names
|
||||
backup = backup
|
||||
eval=coco
|
||||
|
File diff suppressed because it is too large
Load diff
|
@ -1,7 +0,0 @@
|
|||
classes= 6
|
||||
train = i_train.txt
|
||||
valid = i_val.txt
|
||||
names = obj.names
|
||||
backup = backup
|
||||
eval=coco
|
||||
|
File diff suppressed because it is too large
Load diff
|
@ -1,44 +0,0 @@
|
|||
---
|
||||
id: BOP_dataset
|
||||
title: script for create BOP dataset
|
||||
---
|
||||
|
||||
## Структура входных данных:
|
||||
```
|
||||
<example_dir>/
|
||||
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 <example_dir>/
|
||||
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 # протокол создания пакетов датасета
|
||||
```
|
|
@ -1,261 +0,0 @@
|
|||
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)
|
Loading…
Add table
Add a link
Reference in a new issue