robolib.serialization

 1"""
 2entites.py Helpers to make various entities, like cubes, robots, meshes
 3Import level: 3
 4"""
 5from pathlib import Path
 6from typing import Any
 7import numpy as np
 8from microecs import World, EntityId, ComponentType, Entity
 9from robolib.utils import RlModelWithTexture, get_project_root, make_radius
10from robolib.components import HasModel, HasPose, HasFPV, ROBOLIB_COMPONENTS
11
12def world_to_dict(world: World) -> dict[str, Any]:
13    """Serialize the world. Goes through all the entities and their components and converts the serializables to dict"""
14    res = {"entities": [], "components": world.component_names, "extra_metadata": world.extra_metadata}
15    for eid in world.live_entities.keys():
16        res["entities"].append({"entity_id": eid, **world.get_entity(eid).to_dict(serialization_field="serializable")})
17    return res
18
19def world_from_dict(data: dict[str, Any]) -> World:
20    """Creates a world from a serialized representation e.g. from world_to_dict()"""
21    components = [{c.__name__: c for c in ROBOLIB_COMPONENTS}[c] for c in data["components"]]
22    world = World(components=components, extra_metadata=data["extra_metadata"])
23    for entity in data["entities"]:
24        components = [world.component_name_to_type[c] for c in entity["components"]]
25        add_entity(world, components, **entity["data"])
26    world.update()
27    return world
28
29def make_nonserializable_data(components: list[ComponentType], data: dict[str, Any]) -> dict[str, np.ndarray]:
30    """
31    Given the serializable data (e.g. from entity_to_dict()), create the live non-serializable (e.g. gpu model).
32    This is the missing link that allows us to spawn new entities (via world.add_entity).
33    """
34    res = {}
35    for component in components:
36        if component == HasModel:
37            model = RlModelWithTexture.from_path(data["model_path"], data["texture_path"])
38            top_left, bottom_right = model.bbox
39            res["model"] = np.array([model], "object")
40            res["model_bbox"] = np.float32([top_left, bottom_right])
41            res["model_radius"] = np.float32([make_radius(size=bottom_right - top_left)])
42        elif component == HasPose:
43            res["candidate_pose"] = data["pose"].copy()
44        elif component == HasFPV:
45            res["fpv_camera"] = np.array([None], "object")
46            res["fpv_data"] = np.array([None], "object")
47            res["fpv_texture"] = np.array([None], "object")
48    return res
49
50def add_entity(world: World, components: list[ComponentType | str], **kwargs) -> EntityId:
51    """Calls world.make_entity() but only after calling make_nonserializable_data. Bridges robosim with microecs."""
52    components: list[ComponentType] = [world.component_name_to_type[c] if isinstance(c, str) else c for c in components]
53
54    for key in ("model_path", "texture_path"): # store portable, root-relative strings
55        if kwargs.get(key) is not None:
56            p = Path(kwargs[key])
57            kwargs[key] = str(p.relative_to(get_project_root()) if p.is_absolute() else p)
58    return world.add_entity(components, **kwargs, **make_nonserializable_data(components, kwargs))
59
60def add_component(entity: Entity, component: ComponentType, **kwargs):
61    """Small wrapper on top of entity.add_component handling the non serializable data"""
62    entity.add_component(component, **kwargs, **make_nonserializable_data([component], kwargs))