Unable to register Pose Annotator with Replicator and Omniverse Code

Hi,

I’ve been following multiple tutorials and have a version of the “omniverse code extension” tutorial. I’ve modified it to use my own custom Writer, where the only modification is to also write out pose. I’m getting the following error that the pose node could not be found, but I literally copied the code from an isaac sim offline data generation tutorial and parts of the YCBVideoWriter.

Here’s the error:

[omni.replicator.core.scripts.annotators] Attaching pose to render product(s) ['/Render/RenderProduct_Replicator']
2023-06-29 16:40:32 [219,201ms] [Warning] [omni.graph.core.plugin] Could not find node type interface for 'omni.replicator.isaac.Pose'
2023-06-29 16:40:32 [219,202ms] [Error] [omni.graph] Invalid NodeObj object in Py_Node in getAttributeExists
2023-06-29 16:40:32 [219,202ms] [Error] [omni.syntheticdata.scripts.SyntheticData] SyntheticData failed to set node /Render/PostProcess/SDGPipeline/RenderProduct_Replicator_pose static attribute inputs:imageWidth
2023-06-29 16:40:32 [219,202ms] [Error] [omni.graph] Invalid NodeObj object in Py_Node in getAttributeExists

Here’s my custom writer, which works fine until I call the line AnnotatorRegistry.get_annotator("pose", ...)


class CustomWriter(BasicWriter):
    def __init__(self, intrinsic_matrix, class_to_index, *args, semantic_types=None, **kwargs):
        super().__init__(*args, **kwargs, semantic_types=semantic_types)
        self.intrinsic_matrix = intrinsic_matrix
        self.class_to_index = class_to_index

        # Pose Data
        self.annotators.append(
            AnnotatorRegistry.get_annotator("pose", init_params={"semanticTypes": semantic_types})
        )

    def initialize(self, *args, **kwargs):
        super().initialize(*args, **kwargs)

    def write(self, data):
        import ipdb
        ipdb.set_trace()
        super().write(data)

        for annotator in data.keys():
            if annotator.startswith("pose"):
                self._write_pose(data, "pose", annotator)


    def _write_pose(self, data: dict, render_product_path: str, annotator: str):
        """Saves a metadata ".mat" file for the YCB Video Dataset, containing:
           - Class indexes (from a pre-defined mapping) corresponding to each semantically-labeled object in view.
           - A depth image scaling factor.
           - The intrinsic matrix of the camera.
           - Poses from the frame of each semantically-labeled object in view to the world frame, represented as a
             rotation matrix and a translation.
           - The center (in pixel coordinates) of each semantically-labeled object in view. Pixel coordinates are
             expressed relative to the top-left corner of the image, with +x to the right and +y down.

        Args:
            data (dict): A dictionary containing the annotator data for the current frame.
            render_product_path (str): Directory name to save data to, corresponding to a specific render product.
            annotator (str): Annotator name used as a key in the data dictionary, which can also be used to retrieve the
                             annotator from the annotator registry.
        """

        pose_data = data[annotator]["data"]

        n = len(pose_data)

        if n > 0:
            transform_matrices = np.zeros((n, 4, 4))
        else:
            transform_matrices = np.array([[[]]])

        id_to_labels = data[annotator]["info"]["idToLabels"]

        cls_indexes = []
        centers = []

        for i, (semantic_id, pose, center) in enumerate(pose_data):

            # Class indexes
            semantic_labels = id_to_labels[str(semantic_id)]["class"]
            semantic_label = semantic_labels.split(",")[0]
            semantic_index = self.class_to_index[semantic_label]
            cls_indexes.append(semantic_index)

            # Poses
            transform_matrices[i, ...] = pose

            # Centers
            centers.append([center[0], center[1]])

        if n > 0:
            # Make poses have a shape of (3, 4, n)
            poses = np.moveaxis(transform_matrices[:, :-1, :], 0, -1)
        else:
            # Make empty poses have a shape of (1, 1, 0)
            poses = transform_matrices

        meta_dict = {
            "cls_indexes": np.asarray(cls_indexes, dtype=np.uint8),
            "intrinsic_matrix": self.intrinsic_matrix,
            "poses": poses,
            "center": np.array(centers, dtype=np.float64),
        }

        buf = io.BytesIO()
        savemat(buf, meta_dict)

        image_id = "{:06d}".format(self._frame_id)
        file_path = f"{self._output_dir}/{render_product_path}{image_id}-meta.mat"

        self._backend.write_blob(file_path, buf.getvalue())


WriterRegistry.register(CustomWriter)

Here’s the code for registering a pose annotator:


def register_pose_annotator():
    """Register custom pose annotator, specifying its upstream inputs required for computation and its output data
        type.
    """

    NodeConnectionTemplate = SyntheticData.NodeConnectionTemplate

    rep.AnnotatorRegistry.register_annotator_from_node(
        name="PoseSync",
        input_rendervars=[
            NodeConnectionTemplate(
                "PostProcessDispatch", attributes_mapping={"outputs:swhFrameNumber": "inputs:syncValue"}
            ),
            NodeConnectionTemplate(
                "SemanticBoundingBox2DExtentTightSDExportRawArray",
                attributes_mapping={"outputs:exec": "inputs:execIn"},
            ),
            NodeConnectionTemplate(
                "InstanceMappingWithTransforms", attributes_mapping={"outputs:exec": "inputs:execIn"}
            ),
            NodeConnectionTemplate("CameraParams", attributes_mapping={"outputs:exec": "inputs:execIn"}),
        ],
        node_type_id="omni.graph.action.SyncGate",
    )

    rep.AnnotatorRegistry.register_annotator_from_node(
        name="pose",
        input_rendervars=[
            NodeConnectionTemplate("PoseSync", attributes_mapping={"outputs:execOut": "inputs:exec"}),
            NodeConnectionTemplate(
                "SemanticBoundingBox2DExtentTightSDExportRawArray",
                attributes_mapping={"outputs:data": "inputs:data", "outputs:bufferSize": "inputs:bufferSize"},
            ),
            "InstanceMappingWithTransforms",
            "CameraParams",
        ],
        node_type_id="omni.replicator.isaac.Pose",
        init_params={
            "imageWidth": 1024,
            "imageHeight": 1024,
            "cameraRotation": np.array([0,0,0]),
            "getCenters": True,

            # TODO: if we need this True, need to fix weird bug with OgnPose.py that doesn't consider this flag (/home/armstrong/isaac_sim/exts/omni.replicator.isaac/omni/replicator/isaac/ogn/python/nodes/OgnPose.py)
            "includeOccludedPrims": False, 
        },
        output_data_type=np.dtype(
            [
                ("semanticId", "<u4"),
                ("prims_to_desired_camera", "<f4", (4, 4)),
                ("center_coords_image_space", "<f4", (2,)),
            ]
        ),
        output_is_2d=False,
    )

And finally, here’s the part where the writer is created:


        # Initialize and attach writer
        register_pose_annotator()
        writer = rep.WriterRegistry.get("CustomWriter")
        writer.initialize(output_dir=OUTPUT_FOLDER,
                        rgb=True, bounding_box_2d_tight=False, semantic_segmentation=False, bounding_box_3d=True, intrinsic_matrix=None, class_to_index=None)
        writer.attach([render_product, render_product2])

I believe there must be some issue with running code as an extension vs calling the isaac sim python bash script. I’ve encountered other issues where I tried importing core libraries like “from omni.isaac.core import World” and get an error that this couldn’t be found. Maybe I just need to include the python path?

Hi there,

could you provide a short script with the steps in to reproduce the error? If you think you this is caused by import errors, it might be due to omni libraries imported before from omni.isaac.kit import SimulationApp. See 1. Hello World — isaacsim latest documentation

#launch Isaac Sim before any other imports