Custom writer produces black semantic images

Hello,

I am trying to create a custom writers that saves semantic images. If I set colorize = True I get a single channel array with all of the values 4278190080 and data type uint32. If I don’t do colorize then it is a single channel array of all ones still uint32. How do I get a colored 3 or 4 channel array from this?

The semantic labels and the rgb image work perfectly
For replicator.yaml I found something called groundtruth_visuals. Is there a python equivalent?

class CustomWriter(Writer):
    def __init__(self, 
                 output_dir,
                 rgb = False,
                 bbox_tight = False,
                 bbox_loose = False,
                 bbox_3d = False,                   
                 depth = False,
                 semantic = False,
                 instance = False,
                 image_output_format="png"
                 ):
        
        self.output_dir = output_dir

        self._backend = BackendDispatch({"paths": {"out_dir": output_dir}})

        self._frame_id = 0
        self.rgb = rgb
        self.depth = depth
        self.semantic = semantic
        self.instance = instance
        self.bbox_tight = bbox_tight
        self.bbox_loose = bbox_loose
        self.bbox_3d = bbox_3d


        if rgb:
            self.annotators.append(AnnotatorRegistry.get_annotator("rgb"))
        if semantic:
            self.annotators.append(omni.replicator.core.annotators.get("semantic_segmentation", init_params={"colorize": True}))


    def write(self, data):
        if self.rgb:
            filepath = f"rgb_images/rgb-{self._frame_id}.{self._image_output_format}"
            self._backend.write_image(filepath, data["rgb"])
        
        if self.semantic:
            filepath = f"semantic_images/semantic_segmentation-{self._frame_id}.{self._image_output_format}"
            self._backend.write_image(filepath, data["semantic_segmentation"]['data'])

            filepath = f"semantic_images/semantic_segmentation_labels-{self._frame_id}.json"
            
            buf = io.BytesIO()
            buf.write(json.dumps(data["semantic_segmentation"]['info']['idToLabels']).encode())
            self._backend.write_blob(filepath, buf.getvalue())

Hi!

You could use colorize_segmentation from writer tools.
PYTHON API — omni_replicator 1.10.10 documentation (nvidia.com)

Modified part:

        if self.semantic:
            filepath = f"semantic_images/semantic_segmentation-{self._frame_id}.{self._image_output_format}"
            semantic_seg_data, palette, color_2_labels = rep.tools.colorize_segmentation(
                data["semantic_segmentation"]["data"],
                data["semantic_segmentation"]["info"]["idToLabels"]
            )
            self._backend.write_image(filepath, semantic_seg_data)

or you can try the logic from BasicWriter

        if self.semantic:
            filepath = f"semantic_images/semantic_segmentation--{self._frame_id}.{self._image_output_format}"

            semantic_seg_data = data["semantic_segmentation"]["data"]
            height, width = semantic_seg_data.shape[:2]
            semantic_seg_data = semantic_seg_data.view(np.uint8).reshape(height, width, -1)

            self._backend.write_image(filepath, semantic_seg_data)

Both have worked for me, I would recommend to check the BasicWriter class for inspiration.

2 Likes

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.