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