Hash semantic ids to improve distinguishability of semantic segmentation images

source/isaaclab/isaaclab/renderer/ovrtx_renderer.py

@@ -43,6 +43,7 @@
     create_camera_transforms_kernel,
     extract_depth_tile_from_tiled_buffer_kernel,
     extract_tile_from_tiled_buffer_kernel,
+    compute_crc32_hash_kernel,
     sync_newton_transforms_kernel,
 )
 from .ovrtx_usd import (
@@ -428,6 +429,17 @@ def write_output(
         if src.ptr != output_data.data_ptr():
             wp.copy(dest=wp.from_torch(output_data), src=src)
 
+    def _compute_crc32_hash(self, input_data: wp.array) -> wp.array:
+        """Compute CRC32 hash per pixel for a 2D uint32 buffer; returns (h, w) uint32 array."""
+        output_hash = wp.zeros(shape=input_data.shape, dtype=wp.uint32, device=DEVICE)
+        wp.launch(
+            kernel=compute_crc32_hash_kernel,
+            dim=input_data.shape,
+            inputs=[input_data, output_hash],
+            device=DEVICE,
+        )
+        return output_hash
+
     def _extract_rgba_tiles(
         self,
         tiled_data: wp.array,
@@ -508,13 +520,20 @@ def _process_render_frame(self, frame, output_buffers: dict) -> None:
         if "SemanticSegmentationSD" in frame.render_vars and "semantic_segmentation" in output_buffers:
             with frame.render_vars["SemanticSegmentationSD"].map(device=Device.CUDA) as mapping:
                 tiled_semantic_data = wp.from_dlpack(mapping.tensor)
+
                 if tiled_semantic_data.dtype == wp.uint32:
-                    semantic_torch = wp.to_torch(tiled_semantic_data)
+                    # Hash each semantic ID so nearby IDs (e.g. 1, 2, 3) map to visually distinct colors in the output.
+                    semantic_hash = self._compute_crc32_hash(tiled_semantic_data)
+
+                    semantic_torch = wp.to_torch(semantic_hash)
                     semantic_uint8 = semantic_torch.view(torch.uint8)
+
                     if semantic_torch.dim() == 2:
                         h, w = semantic_torch.shape
                         semantic_uint8 = semantic_uint8.reshape(h, w, 4)
+
                     tiled_semantic_data = wp.from_torch(semantic_uint8, dtype=wp.uint8)
+
                 self._extract_rgba_tiles(
                     tiled_semantic_data, output_buffers, "semantic_segmentation", suffix="semantic"
                 )

source/isaaclab/isaaclab/renderer/ovrtx_renderer_kernels.py

@@ -78,6 +78,44 @@ def extract_depth_tile_from_tiled_buffer_kernel(
     tile_buffer[y, x, 0] = tiled_buffer[src_y, src_x]
 
 
+@wp.func
+def _crc32_byte(crc: wp.uint32, byte: wp.uint32) -> wp.uint32:
+    """One byte step of CRC32 (bit-by-bit, standard reflected polynomial)."""
+    CRC32_POLY = wp.uint32(0xEDB88320)
+
+    crc = crc ^ byte
+
+    for _ in range(8):
+        if (crc & wp.uint32(1)) != wp.uint32(0):
+            crc = (crc >> wp.uint32(1)) ^ CRC32_POLY
+        else:
+            crc = crc >> wp.uint32(1)
+
+    return crc
+
+
+@wp.kernel
+def compute_crc32_hash_kernel(
+    input_data: wp.array(dtype=wp.uint32, ndim=2),  # type: ignore
+    output_hash: wp.array(dtype=wp.uint32, ndim=2),  # type: ignore
+):
+    """Compute a deterministic CRC32 hash per pixel (uint32 -> uint32)."""
+    i, j = wp.tid()
+    value_uint32 = input_data[i, j]
+
+    crc = wp.uint32(0xFFFFFFFF)
+
+    # Extract 4 bytes (little-endian: LSB first, same order as standard CRC32)
+    for _ in range(4):
+        byte = value_uint32 & wp.uint32(0xFF)
+        crc = _crc32_byte(crc, byte)
+        value_uint32 = value_uint32 >> wp.uint32(8)
+
+    crc = crc ^ wp.uint32(0xFFFFFFFF)
+
+    output_hash[i, j] = crc
+
+
 @wp.kernel
 def sync_newton_transforms_kernel(
     ovrtx_transforms: wp.array(dtype=wp.mat44d),  # type: ignore