FreeCAD: Workbench Refactor

simulation/object_detection/obj2Yolov4dataset.py

@@ -0,0 +1,144 @@
+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