6D Pose Estimation dataset generation tools for BOP format
This commit is contained in:
parent
184ac7df88
commit
90945c0430
3 changed files with 305 additions and 0 deletions
0
PoseEstimation/.gitkeep
Normal file
0
PoseEstimation/.gitkeep
Normal file
44
PoseEstimation/BOPdataset.md
Normal file
44
PoseEstimation/BOPdataset.md
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
---
|
||||
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 # протокол создания пакетов датасета
|
||||
```
|
||||
261
PoseEstimation/objs2BOPdataset.py
Normal file
261
PoseEstimation/objs2BOPdataset.py
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
import blenderproc as bproc
|
||||
"""
|
||||
objs2BOPdataset
|
||||
Общая задача: распознавание 6D позы объекта (6D pose estimation)
|
||||
Реализуемая функция: создание датасета в формате BOP для серии заданных объектов (*.obj) в заданной сцене (*.blend)
|
||||
Используется модуль blenderproc
|
||||
|
||||
29.08.2023 @shalenikol release 0.1
|
||||
12.10.2023 @shalenikol release 0.2
|
||||
"""
|
||||
import sys
|
||||
import numpy as np
|
||||
import argparse
|
||||
import random
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
|
||||
Not_Categories_Name = True # наименование категории в аннотации отсутствует
|
||||
|
||||
def convert2relative(height, width, bbox):
|
||||
"""
|
||||
YOLO format use relative coordinates for annotation
|
||||
"""
|
||||
x, y, w, h = bbox
|
||||
x += w/2
|
||||
y += h/2
|
||||
return x/width, y/height, w/width, h/height
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('scene', nargs='?', default="resources/sklad.blend", help="Path to the scene object.")
|
||||
parser.add_argument('obj_path', nargs='?', default="resources/in_obj", help="Path to the object files.")
|
||||
parser.add_argument('output_dir', nargs='?', default="output", help="Path to where the final files, will be saved")
|
||||
parser.add_argument('vhacd_path', nargs='?', default="blenderproc_resources/vhacd", help="The directory in which vhacd should be installed or is already installed.")
|
||||
parser.add_argument('-single_object', nargs='?', type= bool, default=True, help="One object per frame.")
|
||||
parser.add_argument('--imgs', default=2, type=int, help="The number of times the objects should be rendered.")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.isdir(args.obj_path):
|
||||
print(f"{args.obj_path} : no object directory")
|
||||
sys.exit()
|
||||
|
||||
if not os.path.isdir(args.output_dir):
|
||||
os.mkdir(args.output_dir)
|
||||
|
||||
single_object = args.single_object
|
||||
|
||||
bproc.init()
|
||||
|
||||
# ? загрузим свет из сцены
|
||||
#cam = bproc.loader.load_blend(args.scene, data_blocks=["cameras"])
|
||||
#lights = bproc.loader.load_blend(args.scene, data_blocks=["lights"])
|
||||
|
||||
# загрузим объекты
|
||||
list_files = os.listdir(args.obj_path)
|
||||
obj_names = []
|
||||
obj_filenames = []
|
||||
all_meshs = []
|
||||
nObj = 0
|
||||
for f in list_files:
|
||||
if (os.path.splitext(f))[1] == ".obj":
|
||||
f = os.path.join(args.obj_path, f) # путь к файлу объекта
|
||||
if os.path.isfile(f):
|
||||
obj = bproc.loader.load_obj(f)
|
||||
all_meshs += obj
|
||||
obj_names += [obj[0].get_name()]
|
||||
obj_filenames += [f]
|
||||
nObj += 1
|
||||
|
||||
if nObj == 0:
|
||||
print("Objects not found")
|
||||
sys.exit()
|
||||
|
||||
for i,obj in enumerate(all_meshs):
|
||||
#print(f"{i} *** {obj}")
|
||||
obj.set_cp("category_id", i+1)
|
||||
|
||||
# загрузим сцену
|
||||
scene = bproc.loader.load_blend(args.scene, data_blocks=["objects"])
|
||||
|
||||
# найдём объекты коллизии (пол и т.д.)
|
||||
obj_type = ["floor", "obj"]
|
||||
collision_objects = []
|
||||
#floor = None
|
||||
for o in scene:
|
||||
o.set_cp("category_id", 999)
|
||||
s = o.get_name()
|
||||
for type in obj_type:
|
||||
if s.find(type) >= 0:
|
||||
collision_objects += [o]
|
||||
o.enable_rigidbody(False, collision_shape='BOX')
|
||||
if not collision_objects:
|
||||
print("Collision objects not found in the scene")
|
||||
sys.exit()
|
||||
|
||||
#floor.enable_rigidbody(False, collision_shape='BOX')
|
||||
|
||||
for obj in all_meshs:
|
||||
# Make the object actively participate in the physics simulation
|
||||
obj.enable_rigidbody(active=True, collision_shape="COMPOUND")
|
||||
# Also use convex decomposition as collision shapes
|
||||
obj.build_convex_decomposition_collision_shape(args.vhacd_path)
|
||||
|
||||
objs = all_meshs + scene
|
||||
|
||||
with open(os.path.join(args.output_dir,"res.txt"), "w") as fh:
|
||||
# fh.write(str(type(scene[0]))+"\n")
|
||||
i = 0
|
||||
for o in objs:
|
||||
i += 1
|
||||
loc = o.get_location()
|
||||
euler = o.get_rotation_euler()
|
||||
fh.write(f"{i} : {o.get_name()} {loc} {euler} category_id = {o.get_cp('category_id')}\n")
|
||||
|
||||
# define a light and set its location and energy level
|
||||
light = bproc.types.Light()
|
||||
light.set_type("POINT")
|
||||
light.set_location([5, -5, 5])
|
||||
#light.set_energy(900)
|
||||
#light.set_color([0.7, 0.7, 0.7])
|
||||
|
||||
light1 = bproc.types.Light(name="light1")
|
||||
light1.set_type("SUN")
|
||||
light1.set_location([0, 0, 0])
|
||||
light1.set_rotation_euler([-0.063, 0.6177, -0.1985])
|
||||
#light1.set_energy(7)
|
||||
light1.set_color([1, 1, 1])
|
||||
|
||||
# define the camera intrinsics
|
||||
bproc.camera.set_intrinsics_from_blender_params(1, 640, 480, lens_unit="FOV")
|
||||
|
||||
# add segmentation masks (per class and per instance)
|
||||
bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"])
|
||||
#bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name", "bop_dataset_name"],
|
||||
# default_values={"category_id": 0, "bop_dataset_name": None})
|
||||
|
||||
# activate depth rendering
|
||||
bproc.renderer.enable_depth_output(activate_antialiasing=False)
|
||||
|
||||
res_dir = os.path.join(args.output_dir, "bop_data")
|
||||
if os.path.isdir(res_dir):
|
||||
shutil.rmtree(res_dir)
|
||||
# Цикл рендеринга
|
||||
n_cam_location = 3 #5 # количество случайных локаций камеры
|
||||
n_cam_poses = 3 #3 # количество сэмплов для каждой локации камеры
|
||||
# Do multiple times: Position the shapenet objects using the physics simulator and render X images with random camera poses
|
||||
for r in range(args.imgs):
|
||||
# один случайный объект в кадре / все заданные объекты
|
||||
meshs = [random.choice(all_meshs)] if single_object else all_meshs[:]
|
||||
|
||||
# Randomly set the color and energy
|
||||
light.set_color(np.random.uniform([0.5, 0.5, 0.5], [1, 1, 1]))
|
||||
light.set_energy(random.uniform(500, 1000))
|
||||
light1.set_energy(random.uniform(3, 11))
|
||||
|
||||
for i,o in enumerate(meshs): #objs
|
||||
mat = o.get_materials()[0]
|
||||
mat.set_principled_shader_value("Specular", random.uniform(0, 1))
|
||||
mat.set_principled_shader_value("Roughness", random.uniform(0, 1))
|
||||
mat.set_principled_shader_value("Base Color", np.random.uniform([0, 0, 0, 1], [1, 1, 1, 1]))
|
||||
mat.set_principled_shader_value("Metallic", random.uniform(0, 1))
|
||||
|
||||
# Clear all key frames from the previous run
|
||||
bproc.utility.reset_keyframes()
|
||||
|
||||
# Define a function that samples 6-DoF poses
|
||||
def sample_pose(obj: bproc.types.MeshObject):
|
||||
obj.set_location(np.random.uniform([-1, -1.5, 0.2], [1, 2, 1.2])) #[-1, -1, 0], [1, 1, 2]))
|
||||
obj.set_rotation_euler(bproc.sampler.uniformSO3())
|
||||
|
||||
# Sample the poses of all shapenet objects above the ground without any collisions in-between
|
||||
#bproc.object.sample_poses(meshs, objects_to_check_collisions = meshs + [floor], sample_pose_func = sample_pose)
|
||||
bproc.object.sample_poses(meshs, objects_to_check_collisions = meshs + collision_objects, sample_pose_func = sample_pose)
|
||||
|
||||
# Run the simulation and fix the poses of the shapenet objects at the end
|
||||
bproc.object.simulate_physics_and_fix_final_poses(min_simulation_time=4, max_simulation_time=20, check_object_interval=1)
|
||||
|
||||
# Find point of interest, all cam poses should look towards it
|
||||
poi = bproc.object.compute_poi(meshs)
|
||||
|
||||
coord_max = [0.1, 0.1, 0.1]
|
||||
coord_min = [0., 0., 0.]
|
||||
|
||||
with open(os.path.join(args.output_dir,"res.txt"), "a") as fh:
|
||||
fh.write("*****************\n")
|
||||
fh.write(f"{r}) poi = {poi}\n")
|
||||
i = 0
|
||||
for o in meshs:
|
||||
i += 1
|
||||
loc = o.get_location()
|
||||
euler = o.get_rotation_euler()
|
||||
fh.write(f" {i} : {o.get_name()} {loc} {euler}\n")
|
||||
for j in range(3):
|
||||
if loc[j] < coord_min[j]:
|
||||
coord_min[j] = loc[j]
|
||||
if loc[j] > coord_max[j]:
|
||||
coord_max[j] = loc[j]
|
||||
|
||||
# Sample up to X camera poses
|
||||
#an = np.random.uniform(0.78, 1.2) #1. #0.35
|
||||
for i in range(n_cam_location):
|
||||
# Sample location
|
||||
location = bproc.sampler.shell(center=[0, 0, 0],
|
||||
radius_min=1.1,
|
||||
radius_max=2.2,
|
||||
elevation_min=5,
|
||||
elevation_max=89)
|
||||
# координата, по которой будем сэмплировать положение камеры
|
||||
j = random.randint(0, 2)
|
||||
# разовый сдвиг по случайной координате
|
||||
d = (coord_max[j] - coord_min[j]) / n_cam_poses
|
||||
if location[j] < 0:
|
||||
d = -d
|
||||
for k in range(n_cam_poses):
|
||||
# Compute rotation based on vector going from location towards poi
|
||||
rotation_matrix = bproc.camera.rotation_from_forward_vec(poi - location, inplane_rot=np.random.uniform(-0.7854, 0.7854))
|
||||
# Add homog cam pose based on location an rotation
|
||||
cam2world_matrix = bproc.math.build_transformation_mat(location, rotation_matrix)
|
||||
bproc.camera.add_camera_pose(cam2world_matrix)
|
||||
location[j] -= d
|
||||
#world_matrix = bproc.math.build_transformation_mat([2.3, -0.4, 0.66], [1.396, 0., an])
|
||||
#bproc.camera.add_camera_pose(world_matrix)
|
||||
#an += 0.2
|
||||
|
||||
# render the whole pipeline
|
||||
data = bproc.renderer.render()
|
||||
# Write data to bop format
|
||||
bproc.writer.write_bop(res_dir,
|
||||
target_objects = all_meshs, # Optional[List[MeshObject]] = None
|
||||
depths = data["depth"],
|
||||
depth_scale = 1.0,
|
||||
colors = data["colors"],
|
||||
color_file_format='JPEG',
|
||||
append_to_existing_output = (r>0),
|
||||
save_world2cam = False) # world coords are arbitrary in most real BOP datasets
|
||||
# dataset="robo_ds",
|
||||
"""
|
||||
!!! categories -> name берётся из category_id !!!
|
||||
см.ниже
|
||||
blenderproc.python.writer : BopWriterUtility.py
|
||||
class _BopWriterUtility
|
||||
def calc_gt_coco
|
||||
...
|
||||
CATEGORIES = [{'id': obj.get_cp('category_id'), 'name': str(obj.get_cp('category_id')), 'supercategory':
|
||||
dataset_name} for obj in dataset_objects]
|
||||
|
||||
поэтому заменим наименование категории в аннотации
|
||||
"""
|
||||
if Not_Categories_Name:
|
||||
coco_file = os.path.join(res_dir,"train_pbr/000000/scene_gt_coco.json")
|
||||
with open(coco_file, "r") as fh:
|
||||
data = json.load(fh)
|
||||
cats = data["categories"]
|
||||
#print(f"type(cat) = {type(cat)} cat : {cat}")
|
||||
i = 0
|
||||
for cat in cats:
|
||||
cat["name"] = obj_names[i]
|
||||
i += 1
|
||||
#print(cat)
|
||||
with open(coco_file, "w") as fh:
|
||||
json.dump(data, fh, indent=0)
|
||||
Loading…
Add table
Add a link
Reference in a new issue