robosim.utils.utils

  1"""utils.py - generic utilities"""
  2from __future__ import annotations
  3from typing import Any
  4from pathlib import Path
  5from dataclasses import dataclass, field
  6from collections import OrderedDict
  7import threading
  8import zlib
  9
 10from loggez import make_logger
 11import numpy as np
 12import raylib as rl
 13
 14logger = make_logger("ROBOSIM", exists_ok=True)
 15
 16Rot3x3  = np.ndarray
 17Pose4x4 = np.ndarray
 18Point2D = np.ndarray
 19Point3D = np.ndarray
 20Point6D = np.ndarray
 21Mat4x4  = np.ndarray
 22Mat3x3  = np.ndarray
 23
 24def fmt(arr: np.ndarray, precision: int=2) -> str:
 25    """formats a numpy array for logging as tuple"""
 26    return "(" + ", ".join(f"{float(x):.{precision}g}" for x in arr) + ")"
 27
 28def get_project_root() -> Path:
 29    """The project root"""
 30    return Path(__file__).parents[2]
 31
 32@dataclass
 33class FPVData:
 34    """The first perso view camera data. Latest frame (as bytes), a lock (for concurrency) and a way to compress it"""
 35    frame: bytes
 36    frame_shape: tuple[int, int, int]
 37    lock: threading.Lock = field(default_factory=threading.Lock)
 38    frame_id: int = 0 # used for de-dup on client side
 39    _frame_compressed: bytes | None = None
 40
 41    @property
 42    def frame_compressed(self) -> bytes:
 43        """compress the frame on demand when the client requests a get_state_with_frame. NOTE: use lock on call."""
 44        if self._frame_compressed is None:
 45            self._frame_compressed = zlib.compress(self.frame, level=1)
 46        return self._frame_compressed
 47
 48# raylib utils
 49
 50RlTexture = Any
 51RlRenderTexture = Any
 52RlModel = Any
 53RlVector3 = Any
 54RlMesh = Any
 55
 56def vec3_from(arr: np.ndarray) -> RlVector3:
 57    """convers the numpy array to a raylib vector"""
 58    return rl.ffi.from_buffer("Vector3 *", np.ascontiguousarray(arr))[0]
 59
 60@dataclass
 61class RlModelWithTexture:
 62    """basic raylib model with an optional texture applied to it"""
 63    model: RlModel
 64    texture: RlTexture | None
 65    _bbox: tuple[np.ndarray, np.ndarray] | None = None
 66
 67    def __post_init__(self):
 68        if self.texture is not None:
 69            self.model.materials.maps[rl.MATERIAL_MAP_ALBEDO].texture = self.texture
 70
 71    @staticmethod
 72    def from_path(model_path: Path, texture_path: Path | None = None) -> RlModelWithTexture:
 73        """builds a RLModelWithTexture from a path"""
 74        assert model_path.exists(), f"Model path: '{model_path}' doesn't exist."
 75        logger.info(f"Loading model from '{model_path}'")
 76        model = rl.LoadModel(str(model_path).encode())
 77        texture = rl.LoadTexture(str(texture_path).encode()) if texture_path is not None else None
 78        return RlModelWithTexture(model=model, texture=texture)
 79
 80    @property
 81    def bbox(self) -> tuple[np.ndarray, np.ndarray]:
 82        """gets the top left and bottom right 3D points of this mesh. Used for collisions. Cahed once after 1st call"""
 83        if self._bbox is None:
 84            mins, maxs = np.float32([1<<31, 1<<31, 1<<31]), np.float32([-1<<31, -1<<31, -1<<31])
 85            for i in range(self.model.meshCount):
 86                mesh: RlMesh = self.model.meshes[i] # of type struct Mesh &
 87                buf = rl.ffi.buffer(mesh.vertices, mesh.vertexCount * 3 * 4)
 88                vertices = np.frombuffer(buf, "float32").reshape(-1, 3)
 89                mins = np.minimum(mins, np.min(vertices, axis=0))
 90                maxs = np.maximum(maxs, np.max(vertices, axis=0))
 91            self._bbox = mins, maxs
 92        return self._bbox
 93
 94    def __del__(self):
 95        if self.model is not None:
 96            rl.UnloadModel(self.model)
 97        if self.texture is not None:
 98            rl.UnloadTexture(self.texture)
 99
100def rl_get_device_and_renderer() -> tuple[str, str]:
101    """returns the device and renderer used by raylib for logging"""
102    gl_renderer = 0x1F01
103    gl_get_string = rl.ffi.cast("unsigned char *(*)(unsigned int)", rl.rlGetProcAddress(b"glGetString"))
104    gstr = lambda e: rl.ffi.string(rl.ffi.cast("char*", gl_get_string(e))).decode()
105    renderer = gstr(gl_renderer)
106    software = ("llvmpipe", "softpipe", "swrast", "lavapipe")
107    device = "CPU (software rasterizer)" if any(s in renderer.lower() for s in software) else "GPU"
108    return device, renderer
109
110# generic utils below
111
112def make_arr(*data):
113    """makes a numpy array from variable data, e.g. make_arr(1,2,3) <=> np.array([1,2,3], "float32")"""
114    return np.float32(data)
115
116class FixedSizeDict(OrderedDict):
117    """An dict with a fixed size. Useful for caching purposes."""
118    def __init__(self, *args, maxlen: int = 0, **kwargs):
119        self._maxlen = maxlen
120        super().__init__(*args, **kwargs)
121
122    def __setitem__(self, key, value):
123        existed_before_add = key in self
124        super().__setitem__(key, value)
125        if not existed_before_add and len(self) > self._maxlen:
126            self.popitem(False)