robosim.entities

  1"""
  2entites.py Helpers to make various entities, like cubes, robots, meshes
  3Import level: 3
  4"""
  5from pathlib import Path
  6import raylib as rl
  7import numpy as np
  8from microecs import World, EntityId
  9from robosim.constants import FPV_HEIGHT, FPV_WIDTH
 10from robosim.utils import RlModelWithTexture, get_project_root, Pose4x4, Point3D, make_radius, FPVData, make_arr
 11from robosim.components import (ColliderKinds, HasModel, HasPose, HasTypeAndArgs,
 12                                    HasCollision, HasFPV, HasVelocity6DoF, HasAcceleration6DoF,
 13                                    HasMotionInput, HasPhysicsLevel1, HasPhysicsLevel2)
 14from robosim.camera import Camera
 15
 16def make_cube(world: World, pose: Pose4x4, size: Point3D, scale: float,
 17              texture_path: Path | None = None) -> EntityId:
 18    """helper function to create a cube in the ECS world"""
 19
 20    mesh = rl.GenMeshCube(size[0], size[1], size[2])
 21    rl_model = rl.LoadModelFromMesh(mesh)
 22    rel_texture_path = str(texture_path.relative_to(get_project_root())) if texture_path is not None else None
 23    rl_texture = rl.LoadTexture(str(rel_texture_path).encode()) if texture_path is not None else None
 24    model = RlModelWithTexture(rl_model, rl_texture)
 25    components = [HasModel, HasPose, HasTypeAndArgs, HasCollision]
 26    top_left, bottom_right = model.bbox
 27
 28    return world.add_entity(
 29        components=components,
 30        _type=np.array([{"type": "cube"}], "object"),
 31        _args=np.array([{"size": size.tolist(), "texture_path": rel_texture_path, "scale": scale}], "object"),
 32        model=np.array([model], "object"),
 33        scale=np.float32([scale]),
 34        model_bbox=np.float32([top_left, bottom_right]),
 35        model_radius=np.float32([make_radius(size=bottom_right - top_left)]),
 36        pose=pose,
 37        collider_kind=np.int32([ColliderKinds.AABB]),
 38        candidate_pose=pose,
 39    )
 40
 41def make_mesh(world: World, model_path: Path, pose: Pose4x4, scale: float,
 42              texture_path: Path | None = None, collidable: bool = False) -> EntityId:
 43    """helper function to create a mesh in the ECS world"""
 44    rel_model_path = model_path.relative_to(get_project_root())
 45    rel_texture_path = str(texture_path.relative_to(get_project_root())) if texture_path is not None else None
 46    model = RlModelWithTexture.from_path(rel_model_path, rel_texture_path)
 47    top_left, bottom_right = model.bbox
 48    components = [HasModel, HasPose, HasTypeAndArgs]
 49    kwargs = {}
 50    if collidable:
 51        components.append(HasCollision)
 52        kwargs = {
 53            "collider_kind": np.array([ColliderKinds.AABB], "int32"),
 54            "candidate_pose": pose,
 55        }
 56
 57    return world.add_entity(
 58        components=components,
 59        model=np.array([model], "object"),
 60        scale=np.float32([scale]),
 61        model_bbox=np.float32([top_left, bottom_right]),
 62        model_radius=np.float32([make_radius(size=bottom_right - top_left)]),
 63        pose=pose,
 64        _type=np.array([{"type": "mesh"}], "object"),
 65        _args=np.array([
 66            {"model_path": str(rel_model_path), "texture_path": rel_texture_path, "scale": scale}], "object"),
 67        **kwargs,
 68    )
 69
 70def make_robot(world: World, model_path: Path, pose: Pose4x4, scale: float, physics_type: str,
 71               velocity: np.ndarray | None = None, acceleration: np.ndarray | None = None, **kwargs) -> EntityId:
 72    """helper function to create a robot in the ECS world"""
 73    fpv_camera = Camera(make_arr(0, 0, 0), make_arr(0, 0, 0), make_arr(0, 0, 0), 60.0, rl.CAMERA_PERSPECTIVE)
 74    fpv_data = FPVData(frame=np.zeros((FPV_HEIGHT * FPV_WIDTH * 4, ), "uint8"), frame_shape=(FPV_HEIGHT, FPV_WIDTH, 4))
 75    fpv_texture = rl.LoadRenderTexture(FPV_WIDTH, FPV_HEIGHT) # TODO: how do we deallocate?
 76
 77    rel_model_path = model_path.relative_to(get_project_root())
 78    model = RlModelWithTexture.from_path(rel_model_path)  # TODO: how do we deallocate?
 79    top_left, bottom_right = model.bbox
 80    velocity = velocity if velocity is not None else np.zeros((6, ), "float32")
 81    acceleration = acceleration if acceleration is not None else np.zeros((6, ), "float32")
 82
 83    components = [HasTypeAndArgs, HasModel, HasPose, HasFPV, HasVelocity6DoF,
 84                  HasAcceleration6DoF, HasMotionInput, HasCollision]
 85    if physics_type == "level1":
 86        components.append(HasPhysicsLevel1)
 87    else:
 88        components.append(HasPhysicsLevel2)
 89
 90    return world.add_entity(
 91        components=components,
 92        _type=np.array([{"type": "robot"}], "object"),
 93        _args=np.array([{"model_path": str(rel_model_path), "scale": scale}], "object"),
 94        model=np.array([model], "object"),
 95        scale=np.float32([scale]),
 96        model_bbox=np.float32([top_left, bottom_right]),
 97        model_radius=np.float32([make_radius(size=bottom_right - top_left)]),
 98        pose=pose,
 99        velocity=velocity,
100        acceleration=acceleration,
101        fpv_camera=np.array([fpv_camera], "object"),
102        fpv_data=np.array([fpv_data], "object"),
103        fpv_texture=np.array([fpv_texture], "object"),
104        candidate_pose=pose,
105        collider_kind=np.array([ColliderKinds.SPHERE], "int32"),
106        **kwargs
107    )