Multiple labels per SemanticId

Hello,

When generating synthetic data using the replicator, I realized that some semantic ids reference multiple label. Here is an example:

{“0”: {“class”: “chair”}, “1”: {“class”: “stack”}, “2”: {“class”: “palette,stack”}, “3”: {“class”: “rack”}, “4”: {“class”: “palette”}}

My current use case requires to have a single label per id, and as such, id 2 is not valid anymore. Is there a way to enforce the annotator to output only a single label per id (potentially the label given to the deepest prim/mesh of the usd asset)

Thank you in advance,

Best
Anthony

Hello,

There is currently no way to do this with the default built-in API. The “flattening” of labels on prims is intended behavior for prims with multiple labels and/or nested child-labels.

Forunutately, you can use the existing writers and filter the JSON data yourself post-render, or if you write a Custom Writer that would filter the data how you like.

You can take a look at how annotators are structured and access, the label data exists in a dictionary:

id_to_labels = data['semantic_segmentation']["info"]["idToLabels"]

Currently this dictionary is written as-is to JSON

If you are using BasicWriter, a minimal code example would be to create a custom writer that inherits from BasicWriter, but overwrites the _write_semantic_segmentation method

Here is a full example of creating a custom writer that inherits from BasicWriter, but modifies the method that writes the semantic information to JSON to only write a single class label:

import io
import json

import numpy as np

import omni.replicator.core as rep
from omni.syntheticdata import SyntheticData


class MyWriter(rep.BasicWriter):
    # Overwriting default BasicWriter behavior
    # We will copy the original method, but modify it to work how we need it to be
    def _write_semantic_segmentation(self, data: dict, render_product_path: str, annotator: str):
        semantic_seg_data = data[annotator]["data"]
        height, width = semantic_seg_data.shape[:2]

        file_path = (
            f"{render_product_path}semantic_segmentation_{self._sequence_id}{self._frame_id:0{self._frame_padding}}.png"
        )
        if self.colorize_semantic_segmentation:
            semantic_seg_data = semantic_seg_data.view(np.uint8).reshape(height, width, -1)
            self._backend.write_image(file_path, semantic_seg_data)
        else:
            semantic_seg_data = semantic_seg_data.view(np.uint32).reshape(height, width)
            self._backend.write_image(file_path, semantic_seg_data)

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

        ###
        ###  Here is where you would modify id_to_labels dictionary to only have 1 label
        ###
        # The data looks like:     "(255, 197, 25, 255)": {"class": "ball,sphere"}
        for k, v in id_to_labels.items():
            v["class"] = v["class"].split(",")[0]

        file_path = f"{render_product_path}semantic_segmentation_labels_{self._sequence_id}{self._frame_id:0{self._frame_padding}}.json"
        buf = io.BytesIO()
        buf.write(json.dumps({str(k): v for k, v in id_to_labels.items()}).encode())
        self._backend.write_blob(file_path, buf.getvalue())


# Register new writer with Replicator
rep.WriterRegistry.register(MyWriter)

## Scene creation
with rep.new_layer():
    rep.create.light()  # Default light
    camera = rep.create.camera(position=(0, 500, 1000), look_at=(0, 0, 0))

    # Create simple shapes to manipulate
    plane = rep.create.plane(semantics=[("class", "plane")], position=(0, -100, 0), scale=(100, 1, 100))
    cubes = rep.create.cube(
        semantics=[("class", "cube")],
        position=rep.distribution.uniform((-300, 0, -300), (300, 0, 300)),
        count=6,
    )
    # Sphere has 2 semantic labels
    spheres = rep.create.sphere(
        semantics=[("class", "sphere"), ("class", "ball")],
        position=rep.distribution.uniform((-300, 0, -300), (300, 0, 300)),
        count=6,
    )

    with rep.trigger.on_frame(num_frames=10):
        with cubes:
            rep.randomizer.color(colors=rep.distribution.normal((0.2, 0.2, 0.2), (1.0, 1.0, 1.0)))
        with spheres:
            rep.randomizer.color(colors=rep.distribution.normal((0.2, 0.2, 0.2), (1.0, 1.0, 1.0)))


render_product = rep.create.render_product(camera, (512, 512))

writer = rep.WriterRegistry.get("MyWriter")
writer.initialize(
    output_dir="custom_writer_semantics",
    rgb=True,
    semantic_segmentation=True,
    colorize_semantic_segmentation=True,
)

writer.attach([render_product])
rep.orchestrator.run()