296 lines
No EOL
11 KiB
Python
296 lines
No EOL
11 KiB
Python
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") |