144 lines
5.3 KiB
Python
144 lines
5.3 KiB
Python
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
|