Source code for gs_nyx_plugin.nyx_camera_shared_metadata

"""
Nyx Camera Shared Metadata for Genesis - GPU-accelerated rendering via Nyx renderer.
"""

# External imports
import os
import shutil
import string
import secrets
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from dataclasses import dataclass
import torch

# Internal imports
from genesis.engine.sensors.base_sensor import (
    SharedSensorMetadata,
    KinematicSensorMetadataMixin,
)

# Forward declaure the NyxCameraSensor to avoid circular impots
if TYPE_CHECKING:
    from .nyx_camera_sensor import NyxCameraSensor

__all__ = ["NyxCameraSharedMetadata"]


# ========================== Utils ==========================
def _nyx_random_string(length: int = 16) -> str:
    alphabet = string.ascii_lowercase + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(length))


def _nyx_make_complete_path(path: str) -> str:
    return f"__nyx_cache__/{path}"


# ========================== Shared Metadata ==========================


@dataclass
class NyxCameraSharedMetadata(KinematicSensorMetadataMixin, SharedSensorMetadata):
    """Cross-sensor state shared by every :class:`NyxCameraSensor` in a scene.

    Genesis instantiates one ``SharedSensorMetadata`` per sensor type per
    sensor manager, and stamps it onto every sensor of that type via
    ``self._shared_metadata``. The Nyx plugin uses it as the single source of
    truth for "things there should only ever be one of in the scene":
    the native renderer, the scene exporter, the on-disk scene-description
    cache, and the per-sensor image cache that downstream consumers read.

    Attributes
    ----------
    renderer : NyxPyRenderer or None
        The shared renderer instance. ``None`` until the last sibling sensor
        finishes :meth:`NyxCameraSensor.build`, then constructed once and
        reused for every render across all sensors and environments.
    scene_exporter : NyxSceneExporter or None
        Exporter that produced the on-disk scene-description JSON. Retained
        so :meth:`NyxCameraSensor.pick_pixel` can re-resolve UUIDs back to
        Genesis entities via its
        :attr:`~NyxSceneExporter._entity_uuid_pairs`.
    sensors : list of NyxCameraSensor or None
        Every Nyx camera sensor registered with the manager, in registration
        order. Each sensor's camera index in the shared renderer matches its
        position in this list.
    camera_defs : list of dict or None
        Per-camera configuration dicts passed to the renderer when it
        builds. One entry per sensor, in the same order as :attr:`sensors`.
    image_cache : dict of int to torch.Tensor or None
        ``{sensor_idx: tensor}`` map of pre-allocated ``(B, H, W, 3)``
        ``uint8`` CUDA buffers that the renderer writes into each frame.
        Returned by :meth:`NyxCameraSensor.read` without an extra copy.
    last_render_timestep : int
        Scene timestep at which the renderer last ran. Used for staleness
        checks so multiple sensors querying the same step trigger only one
        render. ``-1`` before the first frame.
    scene_description_export_path : str or None
        Path component (under ``__nyx_cache__/``) where the scene-description
        JSON and any per-vgeom ``.obj`` exports live. Generated once at
        construction; cleaned up by :meth:`destroy`.
    """

    # NyxRenderer instance
    renderer: Optional[Any] = None
    # NyxSceneExporter instance
    scene_exporter: Optional[Any] = None
    sensors: Optional[List["NyxCameraSensor"]] = None
    # List of camera definition dictionaries
    camera_defs: Optional[List[Dict[str, Any]]] = None
    # {sensor_idx: torch.Tensor with shape (B, H, W, 3)}
    image_cache: Optional[Dict[int, torch.Tensor]] = None
    # Track when cameras were last rendered
    last_render_timestep: int = -1
    # Scene description export path (shared across all sensors)
    scene_description_export_path: Optional[str] = None

    # Post init overriding.
    def __post_init__(
        self,
    ):
        # Generate a random string name and create the folder
        self.scene_description_export_path = _nyx_random_string()
        complete_path = _nyx_make_complete_path(self.scene_description_export_path)
        os.makedirs(complete_path, exist_ok=True)



    def destroy(self):
        """Tear down the renderer and remove the on-disk scene cache.

        Called by the Genesis sensor manager when the scene is destroyed.
        Releases the native renderer (so its CUDA buffers are freed),
        deletes the ``__nyx_cache__/<id>/`` directory that holds the exported
        scene-description JSON plus any per-vgeom ``.obj`` files, and clears
        the cached references so the metadata can be garbage-collected.
        Safe to call exactly once; idempotency is not required.
        """
        # Destroy renderer
        if self.renderer is not None:
            self.renderer.unload_scene()
            self.renderer.destroy()
            self.renderer = None

        # Delete the folder and everything within
        complete_path = _nyx_make_complete_path(self.scene_description_export_path)
        shutil.rmtree(complete_path)
        complete_path = None

        # Clear state
        if self.scene_exporter is not None:
            self.scene_exporter.destroy()
        self.scene_exporter = None
        self.sensors = None