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