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()

@dennis.lynch if I undersood your solution, you modify the semantic_segmentation_labels.json?

But is that really a solution? The actual data remains unchanged. The data[“sementic segmentation”][“data”] holds the pixel values of the segmentations which was set with the original annotation scheme.
Changing the semantic_segmentation_labels.json wont modify your segmented data.

How can we modify the actual IDs of classes? For example, in my case, I output a grayscale image where each pixel in the image is the ID of the class that that pixel is labeled.

so if my image is just the sky then all the pixels are 0 (since {“0”: {“class”: “BACKGROUND”}})

But, now I want to modify this annotation scheme (the ID specifically). say {“7”: {“class”: “BACKGROUND”}}

and now my data[“sementic segmentation”][“data”] will be all pixel with value 7.

how can we modify the ID of a class?

Thank you

found a solution:

just modify CUSTOM_LABELS to map to IDs:

self.CUSTOM_LABELS = {
“unlabelled”:0,
“sphere”:1,
“cube”: 2,
“plane”: 3,
}

this will modify the actuall data. but not the semantic_segmentation_labels.json, which can and should be updated with @dennis.lynch solution.

The data[“sementic segmentation”][“data”] holds the pixel values of the segmentations which was set with the original annotation scheme.
Changing the semantic_segmentation_labels.json wont modify your segmentation data.

Correct, in this example there would need to be a bit more work change ID 2 to ID 4

just modify CUSTOM_LABELS to map to IDs:
self.CUSTOM_LABELS = {
“unlabelled”:0,
“sphere”:1,
“cube”: 2,
“plane”: 3,
}

This would also not quite work with the example above.

The combined semantic label “2”: {“class”: “palette,stack”} is unique

You would need to account for the additional combinations

        self.CUSTOM_LABELS = {
            "unlabelled": (0, 0, 0, 0),
            "sphere": (128, 64, 128, 255),
            "cube": (244, 35, 232, 255),
            "plane": (102, 102, 156, 255),
            "cube,other": (244, 35, 232, 255) # same color as cube
        }

Full code:

import omni.replicator.core as rep
from omni.replicator.core import Writer, BackendDispatch, WriterRegistry


class MyWriter(Writer):
    def __init__(self, output_dir: str):
        self._frame_id = 0
        self.backend = BackendDispatch({"paths": {"out_dir": output_dir}})
        self.annotators = ["rgb","semantic_segmentation"]
        # Dictionary mapping of label to RGBA color
        self.CUSTOM_LABELS = {
            "unlabelled": (0, 0, 0, 0),
            "sphere": (128, 64, 128, 255),
            "cube": (244, 35, 232, 255),
            "plane": (102, 102, 156, 255),
            "cube,other": (244, 35, 232, 255)
        }

    def write(self, data):
        render_products = [k for k in data.keys() if k.startswith("rp_")]
        self._write_rgb(data, "rgb")
        self._write_segmentation(data, "semantic_segmentation")
        self._frame_id += 1

    def _write_rgb(self, data, annotator: str):
        # Save the rgb data under the correct path
        rgb_file_path = f"rgb_{self._frame_id}.png"
        self.backend.write_image(rgb_file_path, data[annotator])

    def _write_segmentation(self, data, annotator: str):
        seg_filepath = f"seg_{self._frame_id}.png"
        semantic_seg_data_colorized = rep.tools.colorize_segmentation(
            data[annotator]["data"], data[annotator]["info"]["idToLabels"], mapping=self.CUSTOM_LABELS
        )
        self.backend.write_image(seg_filepath, semantic_seg_data_colorized)

    def on_final_frame(self):
        self.backend.sync_pending_paths()

# Register new writer
WriterRegistry.register(MyWriter)

# Create a new layer for our work to be performed in.
# This is a good habit to develop for later when working on existing Usd scenes
with rep.new_layer():
    light = rep.create.light(light_type='dome')
    # Create a simple camera with a position and a point to look at
    camera = rep.create.camera(position=(0, 500, 1000), look_at=(0, 0, 0))

    # Create some simple shapes to manipulate
    plane = rep.create.plane(semantics=[("class", "plane")], position=(0, -100, 0), scale=(100, 1, 100))
    torus = rep.create.torus(position=(200, 0, 100)) # Torus will be unlabeled
    sphere = rep.create.sphere(semantics=[("class", "sphere")], position=(0, 0, 100))
    cube = rep.create.cube(semantics=[("class", "cube")], position=(-200, 0, 100))
    cube2 = rep.create.cube(semantics=[("class", "cube"), ('class', 'other')], position=(-200, 0, -200))

    # Randomize position and scale of each object on each frame
    with rep.trigger.on_frame(num_frames=10):
        # Creating a group so that our modify.pose operation works on all the shapes at once
        with rep.create.group([torus, sphere, cube]):
            rep.modify.pose(
                position=rep.distribution.uniform((-300, 0, -300), (300, 0, 300)),
                scale=rep.distribution.uniform(0.1, 2),
            )

# Initialize render product and attach a writer
render_product = rep.create.render_product(camera, (1024, 1024))
writer = rep.WriterRegistry.get("MyWriter")
writer.initialize(output_dir="custom_label")
writer.attach([render_product])
rep.orchestrator.run()    # Run the simulation

This great! thank you

Something I just realized which should be mentioned and fixed in the next version release (unless im missing something, please correct me):

rep.tools.colorize_segmentation() only works if you initialize your writer with colorize_semantic_segmentation = False:

in BasicWriter:

why? because inside rep.tools.colorize_segmentation function:

this only works if the data is an array of labels (data.shape = (width, height)).

if data.shape is (width, height, 4), then the np.uniq function will return a list of uniq r channals, g channals, b channals and a channals which will of course not work. Please correct me if i got this wrong,

Thank you

So in order to be able to use this function and color my segmented images with my own annotation chaloring, i have to set colorize_semantic_segmentation = False, which is counter intuitive and over complicated i believe.

Yes, the reason being: passing colorize_semantic_segmentation = True to the annotator, the renderer will do its own colorization.

You would be trying to colorize something that has already be reshaped and colorized into RGBA.

colorize_segmentation() should be used with colorize_semantic_segmentation = False so that you can pass-in your own mappings and/or colors.