diff --git a/source/isaaclab/isaaclab/renderer/ovrtx_renderer.py b/source/isaaclab/isaaclab/renderer/ovrtx_renderer.py index c6a402442fee..1a59a67a1d3b 100644 --- a/source/isaaclab/isaaclab/renderer/ovrtx_renderer.py +++ b/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" ) diff --git a/source/isaaclab/isaaclab/renderer/ovrtx_renderer_kernels.py b/source/isaaclab/isaaclab/renderer/ovrtx_renderer_kernels.py index 124ee1f2b39a..811fc9a2a2a9 100644 --- a/source/isaaclab/isaaclab/renderer/ovrtx_renderer_kernels.py +++ b/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