Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
- Migrate `wp.array(dtype=X)` type annotations to `wp.array[X]` bracket syntax (Warp 1.12+).
- Align articulated `State.body_qd` / FK / IK / Jacobian / mass-matrix linear velocity with COM-referenced motion. If you were comparing `body_qd[:3]` against finite-differenced body-origin motion, recover origin velocity via `v_origin = v_com - omega x r_com_world`. Descendant `FREE` / `DISTANCE` `joint_qd` remains parent-frame and `joint_f` remains a world-frame COM wrench.
- Pin `mujoco` and `mujoco-warp` dependencies to `~=3.6.0`
- Store `Model.shape_color` in linear RGB instead of display/sRGB, honor USD-authored color-space metadata for imported material/display colors and textures, and let `SensorTiledCamera` keep packed color/albedo outputs in linear bytes when `RenderConfig.encode_output_srgb=False`. If you write colors directly into `Model.shape_color`, convert any display/sRGB values to linear first with `newton.utils.srgb_to_linear_rgb()`.

### Deprecated

Expand Down
20 changes: 19 additions & 1 deletion newton/_src/geometry/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,8 @@ def __init__(
self.color = color
# Store texture lazily: strings/paths are kept as-is, arrays are normalized
self._texture = _normalize_texture_input(texture)
self.texture_color_space: str = "auto"
"""Source color space of :attr:`texture`: ``"auto"``, ``"srgb"``, or ``"raw"``."""
self._roughness = roughness
self._metallic = metallic
self.is_solid = is_solid
Expand Down Expand Up @@ -689,6 +691,7 @@ def copy(
roughness=self._roughness,
metallic=self._metallic,
)
m.texture_color_space = self.texture_color_space
if not recompute_inertia:
m.inertia = self.inertia
m.mass = self.mass
Expand Down Expand Up @@ -841,7 +844,7 @@ def uvs(self):

@property
def color(self) -> Vec3 | None:
"""Optional display RGB color with values in [0, 1]."""
"""Optional linear RGB color with values in [0, 1]."""
return self._color

@color.setter
Expand All @@ -857,9 +860,24 @@ def texture(self) -> str | np.ndarray | None:
def texture(self, value: str | np.ndarray | None):
# Store texture lazily: strings/paths are kept as-is, arrays are normalized
self._texture = _normalize_texture_input(value)
self.texture_color_space = "auto"
self._texture_hash = None
self._cached_hash = None

@property
def texture_color_space(self) -> str:
"""Source color space of the assigned texture.

One of ``"auto"`` (assume sRGB for the raytracer), ``"raw"`` (linear /
data), or ``"srgb"`` (explicit sRGB). Reset to ``"auto"`` whenever
:attr:`texture` is reassigned.
"""
return self._texture_color_space

@texture_color_space.setter
def texture_color_space(self, value: str):
self._texture_color_space = value

@property
def texture_hash(self) -> int:
"""Content-based hash of the assigned texture.
Expand Down
24 changes: 19 additions & 5 deletions newton/_src/sensors/sensor_tiled_camera.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,17 +37,26 @@ class SensorTiledCamera(metaclass=_SensorTiledCameraMeta):

Renders up to five image channels per (world, camera) pair:

- **color** -- RGBA shaded image (``uint32``).
- **color** -- RGBA shaded image packed into ``uint32``. By default these
bytes are display/sRGB-encoded; set
``SensorTiledCamera.RenderConfig.encode_output_srgb=False`` to keep them
linear.
- **depth** -- ray-hit distance [m] (``float32``); negative means no hit.
- **normal** -- surface normal at hit point (``vec3f``).
- **albedo** -- unshaded surface color (``uint32``).
- **albedo** -- unshaded surface color packed into ``uint32`` using the
same output encoding convention as **color**.
- **shape_index** -- shape id per pixel (``uint32``).

All output arrays have shape ``(world_count, camera_count, height, width)``. Use the ``flatten_*`` helpers to
rearrange them into tiled RGBA buffers for display, with one tile per (world, camera) pair laid out in a grid.

Shapes without the ``VISIBLE`` flag are excluded.

Shape colors and texture lighting are evaluated in linear RGB internally.
Mesh textures authored as sRGB are converted to linear before shading, and
the packed ``color``/``albedo`` outputs are optionally encoded back to
display/sRGB at the end.

Example:
::

Expand Down Expand Up @@ -119,7 +128,9 @@ def __init__(self, model: Model, *, config: Config | RenderConfig | None = None,
config: Rendering configuration. Pass a :class:`RenderConfig` to
control raytrace settings directly, or ``None`` to use
defaults. The legacy :class:`Config` dataclass is still
accepted but deprecated.
accepted but deprecated. Use
``RenderConfig.encode_output_srgb`` to control whether packed
``color``/``albedo`` outputs are display-encoded or left linear.
load_textures: Load texture data from the model. Set to ``False``
to skip texture loading when textures are not needed.
"""
Expand Down Expand Up @@ -196,11 +207,14 @@ def update(
camera_transforms: Camera-to-world transforms, shape ``(camera_count, world_count)``.
camera_rays: Camera-space rays from :meth:`compute_pinhole_camera_rays`, shape
``(camera_count, height, width, 2)``.
color_image: Output for RGBA color. None to skip.
color_image: Output for packed RGBA color. The bytes are sRGB by
default, or linear when
``render_config.encode_output_srgb=False``. None to skip.
depth_image: Output for ray-hit distance [m]. None to skip.
shape_index_image: Output for per-pixel shape id. None to skip.
normal_image: Output for surface normals. None to skip.
albedo_image: Output for unshaded surface color. None to skip.
albedo_image: Output for packed RGBA albedo. Uses the same output
encoding convention as ``color_image``. None to skip.
refit_bvh: Refit the BVH before rendering.
clear_data: Values to clear output buffers with.
See :attr:`DEFAULT_CLEAR_DATA`, :attr:`GRAY_CLEAR_DATA`.
Expand Down
11 changes: 9 additions & 2 deletions newton/_src/sensors/warp_raytrace/render.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import warp as wp

from ...geometry import Gaussian, GeoType
from ...utils.color import linear_to_srgb_wp
from . import lighting, raytrace, textures, tiling
from .types import MeshData, RenderOrder, TextureData

Expand Down Expand Up @@ -194,7 +195,10 @@ def render_megakernel(
albedo_color = wp.cw_mul(albedo_color, tex_color)

if wp.static(state.render_albedo):
out_albedo[out_index] = tiling.pack_rgba_to_uint32(albedo_color, 1.0)
packed_albedo = albedo_color
if wp.static(config.encode_output_srgb):
packed_albedo = linear_to_srgb_wp(packed_albedo)
out_albedo[out_index] = tiling.pack_rgba_to_uint32(packed_albedo, 1.0)

if not wp.static(state.render_color):
return
Expand Down Expand Up @@ -243,6 +247,9 @@ def render_megakernel(
)
shaded_color = shaded_color + albedo_color * light_contribution

out_color[out_index] = tiling.pack_rgba_to_uint32(shaded_color, 1.0)
packed_color = shaded_color
if wp.static(config.encode_output_srgb):
packed_color = linear_to_srgb_wp(packed_color)
out_color[out_index] = tiling.pack_rgba_to_uint32(packed_color, 1.0)

return render_megakernel
9 changes: 6 additions & 3 deletions newton/_src/sensors/warp_raytrace/render_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from ...geometry import Gaussian, GeoType, Mesh, ShapeFlags
from ...sim import Model, State
from ...utils import load_texture, normalize_texture
from ...utils.color import texture_color_space_to_id
from .bvh import (
compute_bvh_group_roots,
compute_particle_bvh_bounds,
Expand Down Expand Up @@ -697,7 +698,8 @@ def __load_texture_and_mesh_data(self, model: Model, load_textures: bool):
for shape in model.shape_source:
if isinstance(shape, Mesh):
if shape.texture is not None and load_textures:
if shape.texture_hash not in texture_hashes:
texture_key = (shape.texture_hash, shape.texture_color_space)
if texture_key not in texture_hashes:
pixels = load_texture(shape.texture)
if pixels is None:
raise ValueError(f"Failed to load texture: {shape.texture}")
Expand All @@ -707,7 +709,7 @@ def __load_texture_and_mesh_data(self, model: Model, load_textures: bool):
if pixels.dtype != np.uint8:
pixels = pixels.astype(np.uint8, copy=False)

texture_hashes[shape.texture_hash] = len(self.__texture_data)
texture_hashes[texture_key] = len(self.__texture_data)

data = TextureData()
data.texture = wp.Texture2D(
Expand All @@ -720,9 +722,10 @@ def __load_texture_and_mesh_data(self, model: Model, load_textures: bool):
device=self.device,
)
data.repeat = wp.vec2f(1.0, 1.0)
data.color_space = texture_color_space_to_id(shape.texture_color_space)
self.__texture_data.append(data)

texture_data_ids.append(texture_hashes[shape.texture_hash])
texture_data_ids.append(texture_hashes[texture_key])
else:
texture_data_ids.append(-1)

Expand Down
6 changes: 5 additions & 1 deletion newton/_src/sensors/warp_raytrace/textures.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import warp as wp

from ...geometry import GeoType
from ...utils.color import TEXTURE_COLOR_SPACE_RAW_ID, srgb_to_linear_wp
from .types import MeshData, TextureData


Expand All @@ -15,7 +16,10 @@ def flip_v(uv: wp.vec2f) -> wp.vec2f:
@wp.func
def sample_texture_2d(uv: wp.vec2f, texture_data: TextureData) -> wp.vec3f:
color = wp.texture_sample(texture_data.texture, uv, dtype=wp.vec4f)
return wp.vec3f(color[0], color[1], color[2])
rgb = wp.vec3f(color[0], color[1], color[2])
if texture_data.color_space != TEXTURE_COLOR_SPACE_RAW_ID:
rgb = srgb_to_linear_wp(rgb)
return rgb


@wp.func
Expand Down
3 changes: 2 additions & 1 deletion newton/_src/sensors/warp_raytrace/tiling.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,8 @@ def tid_to_coord_view_priority(tid: wp.int32, camera_count: wp.int32, width: wp.

@wp.func
def pack_rgba_to_uint32(rgb: wp.vec3f, alpha: wp.float32) -> wp.uint32:
"""Pack RGBA values into a single uint32 for efficient memory access."""
"""Pack RGB bytes plus alpha into a single ``uint32``."""

return (
(wp.clamp(wp.uint32(alpha * 255.0), wp.uint32(0), wp.uint32(255)) << wp.uint32(24))
| (wp.clamp(wp.uint32(rgb[2] * 255.0), wp.uint32(0), wp.uint32(255)) << wp.uint32(16))
Expand Down
10 changes: 10 additions & 0 deletions newton/_src/sensors/warp_raytrace/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,14 @@ class RenderConfig:
enable_backface_culling: bool = True
"""Cull back-facing triangles."""

encode_output_srgb: bool = True
"""Encode packed color/albedo outputs to display/sRGB.

When ``False``, :class:`SensorTiledCamera` writes packed linear RGB bytes
instead, which can be useful for training pipelines that want to avoid
display transfer functions.
"""

render_order: int = RenderOrder.PIXEL_PRIORITY
"""Render traversal order (see :class:`RenderOrder`)."""

Expand Down Expand Up @@ -115,7 +123,9 @@ class TextureData:
Attributes:
texture: 2D Texture as ``wp.Texture2D``.
repeat: UV tiling factors along U and V axes.
color_space: ``0`` for raw/linear textures, ``1`` for sRGB textures.
"""

texture: wp.Texture2D
repeat: wp.vec2f
color_space: wp.int32
15 changes: 12 additions & 3 deletions newton/_src/sensors/warp_raytrace/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,10 @@ def create_color_image_output(self, width: int, height: int, camera_count: int =
camera_count: Number of cameras.

Returns:
Array of shape ``(world_count, camera_count, height, width)``, dtype ``uint32``.
Array of shape ``(world_count, camera_count, height, width)``,
dtype ``uint32``. Each pixel stores packed RGBA bytes in either
display/sRGB or linear space depending on
``render_config.encode_output_srgb``.
"""
return wp.zeros(
(self.__render_context.world_count, camera_count, height, width),
Expand Down Expand Up @@ -228,7 +231,10 @@ def create_albedo_image_output(self, width: int, height: int, camera_count: int
camera_count: Number of cameras.

Returns:
Array of shape ``(world_count, camera_count, height, width)``, dtype ``uint32``.
Array of shape ``(world_count, camera_count, height, width)``,
dtype ``uint32``. Each pixel stores packed RGBA bytes in either
display/sRGB or linear space depending on
``render_config.encode_output_srgb``.
"""
return wp.zeros(
(self.__render_context.world_count, camera_count, height, width),
Expand Down Expand Up @@ -338,7 +344,10 @@ def flatten_color_image_to_rgba(
Arranges ``(world_count * camera_count)`` tiles in a grid. Each tile shows one camera's view of one world.

Args:
image: Color output from :meth:`~SensorTiledCamera.update`, shape ``(world_count, camera_count, height, width)``.
image: Color output from :meth:`~SensorTiledCamera.update`, shape
``(world_count, camera_count, height, width)``. The packed
bytes are copied as-is; no additional color-space conversion is
performed here.
out_buffer: Pre-allocated RGBA buffer. If None, allocates a new one.
worlds_per_row: Tiles per row in the grid. If None, picks a square-ish layout.
"""
Expand Down
Loading
Loading