runtime/env_manager/env_manager/env_manager/models/objects/mesh.py

import numpy as np
import trimesh
from gym_gz.scenario import model_wrapper
from gym_gz.utils.scenario import get_unique_model_name
from scenario import core as scenario_core
from scenario import gazebo as scenario_gazebo
from trimesh import Trimesh

from env_manager.utils.types import Point

InertiaTensor = tuple[float, float, float, float, float, float]


class Mesh(model_wrapper.ModelWrapper):
    """ """

    def __init__(
        self,
        world: scenario_gazebo.World,
        name: str = "object",
        type: str = "mesh",
        relative_to: str = "world",
        position: tuple[float, float, float] = (0, 0, 0),
        orientation: tuple[float, float, float, float] = (1, 0, 0, 0),
        color: tuple[float, float, float, float] | None = None,
        static: bool = False,
        texture: list[float] = [],
        mass: float = 0.0,
        inertia: InertiaTensor | None = None,
        cm: Point | None = None,
        collision: str = "",
        visual: str = "",
        friction: float = 0.0,
        density: float = 0.0,
        **kwargs,
    ):
        # Get a unique model name
        model_name = get_unique_model_name(world, name)

        # Initial pose
        initial_pose = scenario_core.Pose(position, orientation)

        # Calculate inertia parameters for provided mesh file if not exist
        if not inertia and collision and density:
            mass, cm, inertia = self.calculate_inertia(collision, density)
        else:
            raise RuntimeError(
                f"Please provide collision mesh filepath for model {name} to calculate inertia"
            )

        if not color:
            color = tuple(np.random.uniform(0.0, 1.0, (3,)))
            color = (color[0], color[1], color[2], 1.0)

        # Create SDF string for the model
        sdf = self.get_sdf(
            name=name,
            static=static,
            collision=collision,
            visual=visual,
            mass=mass,
            inertia=inertia,
            friction=friction,
            color=color,
            center_of_mass=cm,
            gui_only=True,
        )

        # Insert the model
        ok_model = world.to_gazebo().insert_model_from_string(
            sdf, initial_pose, model_name
        )
        if not ok_model:
            raise RuntimeError(f"Failed to insert {model_name}")
        # Get the model
        model = world.get_model(model_name)

        # Initialize base class
        super().__init__(model=model)

    def calculate_inertia(self, file_path, density):
        mesh = trimesh.load(file_path)

        if not isinstance(mesh, Trimesh):
            raise RuntimeError("Please provide correct stl mesh filepath")

        volume = mesh.volume
        mass: float = volume * density
        center_of_mass: Point = tuple(mesh.center_mass)
        inertia = mesh.moment_inertia
        eigenvalues = np.linalg.eigvals(inertia)
        inertia_tensor: InertiaTensor = (
            inertia[0, 0],
            inertia[0, 1],
            inertia[0, 2],
            inertia[1, 1],
            inertia[1, 2],
            inertia[2, 2],
        )
        return mass, center_of_mass, inertia_tensor

    @classmethod
    def get_sdf(
        cls,
        name: str,
        static: bool,
        collision: str,
        visual: str,
        mass: float,
        inertia: InertiaTensor,
        friction: float,
        color: tuple[float, float, float, float],
        center_of_mass: Point,
        gui_only: bool,
    ) -> str:
        """
        Generates the SDF string for the box model.

        Args:
            - mesh_args (MeshPureData): Object that contain data of provided mesh data

        Returns:
            The SDF string that defines the box model in Gazebo.
        """
        return f'''<sdf version="1.7">
                <model name="{name}">
                    <static>{"true" if static else "false"}</static>
                    <link name="{name}_link">
                        {
                        f"""
                        <collision name="{name}_collision">
                            <geometry>
                                <mesh>
                                    <uri>{collision}</uri>
                                </mesh>
                            </geometry>
                            <surface>
                                <friction>
                                    <ode>
                                        <mu>{friction}</mu>
                                        <mu2>{friction}</mu2>
                                        <fdir1>0 0 0</fdir1>
                                        <slip1>0.0</slip1>
                                        <slip2>0.0</slip2>
                                    </ode>
                                </friction>
                            </surface>
                        </collision>
                        """ if collision else ""
                        }
                        {
                        f"""
                        <visual name="{name}_visual">
                            <geometry>
                                <mesh>
                                    <uri>{visual}</uri>
                                </mesh>
                            </geometry>
                            <material>
                                <ambient>{color[0]} {color[1]} {color[2]} {color[3]}</ambient>
                                <diffuse>{color[0]} {color[1]} {color[2]} {color[3]}</diffuse>
                                <specular>{color[0]} {color[1]} {color[2]} {color[3]}</specular>
                            </material>
                            <transparency>{1.0 - color[3]}</transparency>
                            {'<visibility_flags>1</visibility_flags> <cast_shadows>false</cast_shadows>' if gui_only else ''}
                        </visual>
                        """ if visual else ""
                        }
                        <inertial>
                            <mass>{mass}</mass>
                                <pose>{center_of_mass[0]} {center_of_mass[1]} {center_of_mass[2]} 0 0 0</pose>
                            <inertia>
                                <ixx>{inertia[0]}</ixx>
                                <ixy>{inertia[1]}</ixy>
                                <ixz>{inertia[2]}</ixz>
                                <iyy>{inertia[3]}</iyy>
                                <iyz>{inertia[4]}</iyz>
                                <izz>{inertia[5]}</izz>
                            </inertia>
                        </inertial>
                    </link>
                </model>
            </sdf>'''