Feature/keypoints (#4)

* Add keypoints functionality and instance id

* Update integration test

.github/workflows/integration_test.yaml

@@ -54,6 +54,8 @@ jobs:
           "./output/*/main_camera_annotations/pointcloud/*metadata.yaml"
           "./output/*/main_camera_annotations/object_volume/*.npz"
           "./output/*/main_camera_annotations/object_volume/*metadata.yaml"
+          "./output/*/main_camera_annotations/keypoints/*.json"
+          "./output/*/main_camera_annotations/keypoints/*metadata.yaml"
           "./output/*/object_positions/*.json"
           "./output/*/object_positions/*metadata.yaml"
         )

docs/docs/index.md

@@ -33,6 +33,7 @@ Syclops supports a variety of annotated outputs for different use cases. The fol
 |**Bounding Boxes**|Bounding boxes for each object in the scene|
 |**Object Positions**|3D position of each object in the scene|
 |**Point Cloud**|3D location of each pixel in camera space|
+|**Keypoints**|Location of keypoints in camera space|
 
 # 📣 Terminology
 |Term|Description|

pyproject.toml

@@ -28,6 +28,7 @@ syclops_plugin_object = "syclops.blender.plugins.object:Object"
 syclops_sensor_camera = "syclops.blender.sensors.camera:Camera"
 [project.entry-points."syclops.outputs"]
 syclops_output_rgb = "syclops.blender.sensor_outputs.rgb:RGB"
+syclops_output_keypoints = "syclops.blender.sensor_outputs.keypoints:Keypoints"
 syclops_output_pixel_annotation = "syclops.blender.sensor_outputs.pixel_annotation:PixelAnnotation"
 syclops_output_object_positions = "syclops.blender.sensor_outputs.object_positions:ObjectPositions"
 [project.entry-points."syclops.postprocessing"]

syclops/__example_assets__/example_job.syclops.yaml

@@ -63,12 +63,6 @@ scene:
       class_id_offset:
         Stem: 1
       seed: 1
-      clumps:
-        ratio: 0.3
-        size: 3
-        size_std: 2
-        position_std: 0.02
-        scale_std: 0.4
     - name: "Weed Scatter"
       models: [Example Assets/Plain Weeds]
       floor_object: Ground
@@ -129,6 +123,8 @@ sensor:
             id: main_cam_rgb
         syclops_output_object_positions:
           - id: main_cam_object_positions
+        syclops_output_keypoints:
+          - id: main_cam_keypoints
         syclops_output_pixel_annotation:
           - semantic_segmentation:
               id: main_cam_semantic

syclops/__example_assets__/test_job.syclops.yaml

@@ -128,6 +128,8 @@ sensor:
             id: main_cam_rgb
         syclops_output_object_positions:
           - id: main_cam_object_positions
+        syclops_output_keypoints:
+          - id: main_cam_keypoints
         syclops_output_pixel_annotation:
           - semantic_segmentation:
               id: main_cam_semantic

syclops/blender/sensor_outputs/keypoints.py

@@ -0,0 +1,137 @@
+import json
+import logging
+from pathlib import Path
+
+import bpy
+import bpy_extras
+from mathutils import Vector
+import mathutils
+from syclops import utility
+from syclops.blender.sensor_outputs.output_interface import OutputInterface
+import numpy as np
+from typing import List
+
+META_DESCRIPTION = {
+    "type": "KEYPOINTS",
+    "format": "JSON",
+    "description": "Keypoints of objects in camera space.",
+}
+
+
+class Keypoints(OutputInterface):
+
+    def generate_output(self, parent_class: object):
+        """Calculate the 3D keypoint positions and transform to camera space."""
+        with utility.RevertAfter():
+            # Update dependencies
+            self._update_depsgraph()
+            self.check_debug_breakpoint()
+
+            obj_dict = {}
+            depsgraph = bpy.context.view_layer.depsgraph
+            scene = bpy.context.scene
+            camera = scene.camera
+            # Camera matrices
+            modelview_matrix = camera.matrix_world.inverted()
+            projection_matrix = camera.calc_matrix_camera(
+                bpy.context.evaluated_depsgraph_get(), x=bpy.context.scene.render.resolution_x, y=bpy.context.scene.render.resolution_y
+            )
+            render_scale = scene.render.resolution_percentage / 100
+            width = int(scene.render.resolution_x * render_scale)
+            height = int(scene.render.resolution_y * render_scale)
+
+
+            for object_instance in depsgraph.object_instances:
+                if object_instance.object.type == "MESH":
+                    parent = object_instance.parent
+                    class_id = None
+                    is_instance = (
+                        object_instance.object.is_from_instancer and parent is not None
+                    )
+                    object_visible = utility.render_visibility(object_instance.object)
+                    if is_instance:
+                        parent_visible = utility.render_visibility(parent)
+                        if parent_visible:
+                            class_id = object_instance.object.get("class_id")
+                    elif object_visible:
+                        class_id = object_instance.object.get("class_id")
+
+                    if class_id is not None:
+                        if "keypoints" in object_instance.object:
+                            location = object_instance.matrix_world.translation
+                            location = [round(x, 4) for x in location]
+                            instance_id = self._calculate_instance_id(location)
+                            for keypoint, pos in object_instance.object["keypoints"].items():
+                                vec = mathutils.Vector((pos['x'], pos['y'], pos['z']))
+                                vec = object_instance.matrix_world @ vec
+                                co_2d = bpy_extras.object_utils.world_to_camera_view(scene, camera, vec)
+
+                                pixel_x = round(co_2d.x * width)
+                                pixel_y = height - round(co_2d.y * height)
+
+                                if pixel_x < 0 or pixel_y < 0 or pixel_x > width or pixel_y > height:
+                                    continue
+
+                                # Add to obj_dict
+                                if instance_id not in obj_dict:
+                                    obj_dict[instance_id] = {"class_id": class_id}
+                                obj_dict[instance_id][keypoint] = {
+                                    "x": pixel_x,
+                                    "y": pixel_y,
+                                }
+
+            # Save output
+            output_path = self._prepare_output_folder(parent_class.config["name"])
+            json_file = output_path / f"{bpy.context.scene.frame_current:04}.json"
+            with open(json_file, "w") as f:
+                json.dump(obj_dict, f)
+
+            # Write meta output
+            self.write_meta_output_file(json_file, parent_class.config["name"])
+            logging.info(f"Writing keypoints output to {json_file}")
+
+    def _update_depsgraph(self):
+        """Update the dependency graph."""
+        bpy.context.view_layer.update()
+        utility.refresh_modifiers()
+
+    def _prepare_output_folder(self, sensor_name):
+        """Prepare the output folder and return its path."""
+        output_folder = utility.append_output_path(f"{sensor_name}_annotations/keypoints/")
+        utility.create_folder(output_folder)
+        return output_folder
+
+
+    def write_meta_output_file(self, file: Path, sensor_name: str):
+        """Write the metadata output to a YAML file."""
+        output_path = file.parent
+
+        with utility.AtomicYAMLWriter(str(output_path / "metadata.yaml")) as writer:
+            writer.data.update(META_DESCRIPTION)
+            writer.add_step(
+                step=bpy.context.scene.frame_current,
+                step_dicts=[
+                    {
+                        "type": META_DESCRIPTION["type"],
+                        "path": str(file.name),
+                    },
+                ],
+            )
+            writer.data["expected_steps"] = utility.get_job_conf()["steps"]
+            writer.data["sensor"] = sensor_name
+            writer.data["id"] = self.config["id"]
+
+    @staticmethod
+    def _calculate_instance_id(location: List[float]) -> int:
+        """Calculate the instance id from the location.
+
+        Args:
+            location: The location.
+
+        Returns:
+            int: The instance id.
+        """
+        # Convert x, y, z to mm and round to integer
+        location = np.round(np.array(location) * 1000)
+
+        return utility.hash_vector(location)

syclops/blender/sensor_outputs/object_positions.py

@@ -1,11 +1,12 @@
 import json
 import logging
 from pathlib import Path
-
+import numpy as np
 import bpy
 from syclops import utility
 from syclops.blender.sensor_outputs.output_interface import OutputInterface
-
+from typing import List
+import mathutils
 # META DESCRIPTION
 meta_description_object_positions = {
     "type": "OBJECT_POSITIONS",
@@ -51,7 +52,18 @@ def generate_output(self, parent_class: object = None):
                             "loc": location,
                             "rot": rotation,
                             "scl": scale,
+                            "id": self._calculate_instance_id(location),
                         }
+                        if "keypoints" in object_instance.object:
+                            pose_dict["keypoints"] = {}
+                            for keypoint, pos in object_instance.object["keypoints"].items():
+                                vec = mathutils.Vector((pos['x'], pos['y'], pos['z']))
+                                world_position = object_instance.object.matrix_world @ vec
+                                # Add keypoint to pose_dict
+                                pose_dict["keypoints"][keypoint] = {
+                                    "loc": [round(x, 4) for x in world_position]
+                                }
+
                         if class_id not in obj_dict:
                             obj_dict[class_id] = []
                         obj_dict[class_id].append(pose_dict)
@@ -68,6 +80,21 @@ def generate_output(self, parent_class: object = None):
             self.write_meta_output_file(Path(json_file))
             logging.info("Wrote object positions to %s", json_file)
 
+    @staticmethod
+    def _calculate_instance_id(location: List[float]) -> int:
+        """Calculate the instance id from the location.
+
+        Args:
+            location: The location.
+
+        Returns:
+            int: The instance id.
+        """
+        # Convert x, y, z to mm and round to integer
+        location = np.round(np.array(location) * 1000)
+
+        return utility.hash_vector(location)
+
     def write_meta_output_file(self, file: Path):
         """Write the meta output file"""
         # Get the output folder

syclops/blender/sensor_outputs/pixel_annotation.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import pickle
 from pathlib import Path
 
 import bpy
@@ -40,23 +39,29 @@
 }
 
 
+
 # Output postprocessing functions
 def postprocess_functions(img, file):
     if "semantic_segmentation" in file:
         img = np.round(img).astype(np.int32)
 
     elif "instance_segmentation" in file:
         try:
-            # Convert x,y,z to mm
-            img = img * 1000
-            # Round to integer
-            img = np.round(img)
-            # Calculate unique x,y,z values and assign new index
-            _, index = np.unique(
-                img.reshape(-1, img.shape[2]), axis=0, return_inverse=True
-            )
-            # Reshape to original shape
-            img = index.reshape(img.shape[:2])
+            # Convert x, y, z to mm and round to integer
+            img_mm = np.round(img * 1000)
+
+            # Calculate unique x, y, z values and assign new index
+            values, index = np.unique(img_mm.reshape(-1, img_mm.shape[2]), axis=0, return_inverse=True)
+
+            # Hash the unique values to get the instance id
+            vectorized_hash = np.vectorize(utility.hash_vector, signature='(n)->()')
+            instance_id = vectorized_hash(values)
+
+            # Create instance segmentation mask
+            img_mask = instance_id[index]
+            img = img_mask.reshape(img_mm.shape[0], img_mm.shape[1])
+            if values.shape[0] != instance_id.shape[0]:
+                logging.warning("Hashing of instance ids created collisions")
         except IndexError:
             logging.warning("Instance segmentation mask is empty")
             img = np.zeros(img.shape[:2], dtype=np.int32)

syclops/blender/sensor_outputs/schema/keypoints.schema.yaml

@@ -0,0 +1,14 @@
+description: Positions of keypoints in image space.
+type: array
+items:
+  type: object
+  properties:
+    id:
+      description: Unique identifier of the output
+      type: string
+    debug_breakpoint:
+      description: Wether to break and open Blender before rendering. Only works if scene debugging is active.
+      type: boolean
+  required: [id]
+minItems: 1
+maxItems: 1

syclops/blender/sensors/schema/camera.schema.yaml

@@ -67,6 +67,8 @@ items:
           $ref: "#/definitions/syclops_output_pixel_annotation"
         syclops_output_object_positions:
           $ref: "#/definitions/syclops_output_object_positions"
+        syclops_output_keypoints:
+          $ref: "#/definitions/syclops_output_keypoints"
       additionalProperties: False
   required: [name, frame_id, resolution, focal_length, sensor_width, outputs]
 

syclops/utility/__init__.py

@@ -26,7 +26,7 @@
 
 
 from .general_utils import (AtomicYAMLWriter, create_folder,
-                            find_class_id_mapping,get_site_packages_path, get_module_path)
+                            find_class_id_mapping,get_site_packages_path, get_module_path, hash_vector)
 from .postprocessing_utils import (crawl_output_meta, filter_type, create_module_instances_pp)
 
 from .setup_utils import (download_file, extract_zip, extract_tar, install_blender, get_or_create_install_folder)

syclops/utility/blender_utils.py

@@ -696,7 +696,7 @@ def create_clumps(collection: bpy.types.Collection, config: dict) -> list:
     Returns:
         list: List of clumps objects.
     """
-    all_objects = collection.all_objects
+    all_objects = [obj for obj in collection.all_objects if obj.type == "MESH"]
     new_clumps = []
     num_objects = len(all_objects)
     num_clumps = _get_num_clumps(num_objects, config["ratio"])
@@ -714,6 +714,7 @@ def create_clumps(collection: bpy.types.Collection, config: dict) -> list:
         clump_items = []
         for obj in instance:
             clump_item = obj.copy()
+
             clump_item.data = obj.data.copy()
             random_transform_object(
                 clump_item, config["position_std"], config["scale_std"]

syclops/utility/general_utils.py

@@ -7,6 +7,21 @@
 from filelock import FileLock, Timeout
 from ruamel import yaml
 import importlib.util
+import hashlib
+import struct
+
+def hash_vector(vector):
+    # Convert the 3D vector into bytes
+    packed_vector = struct.pack('fff', *vector)
+
+    # Use SHA-256 and truncate to 64 bits
+    hash_object = hashlib.sha256(packed_vector)
+    hash_digest = hash_object.digest()[:8]
+
+    # Convert the first 8 bytes of the hash to a 64-bit integer
+    hash_value = int.from_bytes(hash_digest, byteorder='big')
+
+    return hash_value
 
 
 def create_folder(path: str) -> None: