diff --git a/.flake8 b/.flake8 index 5174d49..5e52f64 100644 --- a/.flake8 +++ b/.flake8 @@ -1,10 +1,13 @@ [flake8] -exclude = .git +exclude = .git, ./venv-obj2mjcf max-line-length = 80 ignore = - E203, # whitespace before colon (black default) - E501, # line too long ( characters) - W503, # line break before binary operator + # whitespace before colon (black default) + E203, + # line too long ( characters) + E501, + # line break before binary operator + W503, per-file-ignores = */__init__.py: F401 diff --git a/obj2mjcf/_cli.py b/obj2mjcf/_cli.py index 69034d1..39f83aa 100644 --- a/obj2mjcf/_cli.py +++ b/obj2mjcf/_cli.py @@ -9,15 +9,16 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import List, Optional, Sequence +from typing import List, Optional -import mujoco import trimesh import tyro -from lxml import etree from PIL import Image from termcolor import cprint +from obj2mjcf.material import _MTL_COMMENT_CHAR, Material +from obj2mjcf.mjcf_builder import MJCFBuilder + # Find the V-HACD v4.0 executable in the system path. # Note trimesh has not updated their code to work with v4.0 which is why we do not use # their `convex_decomposition` function. @@ -30,24 +31,6 @@ # 2-space indentation for the generated XML. _XML_INDENTATION = " " -# MTL fields relevant to MuJoCo. -_MTL_FIELDS = ( - # Ambient, diffuse and specular colors. - "Ka", - "Kd", - "Ks", - # d or Tr are used for the rgba transparency. - "d", - "Tr", - # Shininess. - "Ns", - # References a texture file. - "map_Kd", -) - -# Character used to denote a comment in an MTL file. -_MTL_COMMENT_CHAR = "#" - class FillMode(enum.Enum): FLOOD = enum.auto() @@ -106,58 +89,6 @@ class Args: """add a free joint to the root body""" -@dataclass -class Material: - name: str - Ka: Optional[str] = None - Kd: Optional[str] = None - Ks: Optional[str] = None - d: Optional[str] = None - Tr: Optional[str] = None - Ns: Optional[str] = None - map_Kd: Optional[str] = None - - @staticmethod - def from_string(lines: Sequence[str]) -> "Material": - """Construct a Material object from a string.""" - attrs = {"name": lines[0].split(" ")[1].strip()} - for line in lines[1:]: - for attr in _MTL_FIELDS: - if line.startswith(attr): - elems = line.split(" ")[1:] - elems = [elem for elem in elems if elem != ""] - attrs[attr] = " ".join(elems) - break - return Material(**attrs) - - def mjcf_rgba(self) -> str: - Kd = self.Kd or "1.0 1.0 1.0" - if self.d is not None: # alpha - alpha = self.d - elif self.Tr is not None: # 1 - alpha - alpha = str(1.0 - float(self.Tr)) - else: - alpha = "1.0" - # TODO(kevin): Figure out how to use Ka for computing rgba. - return f"{Kd} {alpha}" - - def mjcf_shininess(self) -> str: - if self.Ns is not None: - # Normalize Ns value to [0, 1]. Ns values normally range from 0 to 1000. - Ns = float(self.Ns) / 1_000 - else: - Ns = 0.5 - return f"{Ns}" - - def mjcf_specular(self) -> str: - if self.Ks is not None: - # Take the average of the specular RGB values. - Ks = sum(list(map(float, self.Ks.split(" ")))) / 3 - else: - Ks = 0.5 - return f"{Ks}" - - def resize_texture(filename: Path, resize_percent) -> None: """Resize a texture to a percentage of its original size.""" if resize_percent == 1.0: @@ -407,137 +338,16 @@ def process_obj(filename: Path, args: Args) -> None: f.write("".join(lines)) # Build an MJCF. - root = etree.Element("mujoco", model=filename.stem) - - # Add visual and collision default classes. - default_elem = etree.SubElement(root, "default") - visual_default_elem = etree.SubElement(default_elem, "default") - visual_default_elem.attrib["class"] = "visual" - etree.SubElement( - visual_default_elem, - "geom", - group="2", - type="mesh", - contype="0", - conaffinity="0", - ) - collision_default_elem = etree.SubElement(default_elem, "default") - collision_default_elem.attrib["class"] = "collision" - etree.SubElement(collision_default_elem, "geom", group="3", type="mesh") - - # Add assets. - asset_elem = etree.SubElement(root, "asset") - for material in mtls: - if material.map_Kd is not None: - # Create the texture asset. - texture = Path(material.map_Kd) - etree.SubElement( - asset_elem, - "texture", - type="2d", - name=texture.stem, - file=texture.name, - ) - # Reference the texture asset in a material asset. - etree.SubElement( - asset_elem, - "material", - name=material.name, - texture=texture.stem, - specular=material.mjcf_specular(), - shininess=material.mjcf_shininess(), - ) - else: - etree.SubElement( - asset_elem, - "material", - name=material.name, - specular=material.mjcf_specular(), - shininess=material.mjcf_shininess(), - rgba=material.mjcf_rgba(), - ) - - worldbody_elem = etree.SubElement(root, "worldbody") - obj_body = etree.SubElement(worldbody_elem, "body", name=filename.stem) - if args.add_free_joint: - etree.SubElement(obj_body, "freejoint") - - # Add visual geoms. - if isinstance(mesh, trimesh.base.Trimesh): - meshname = Path(f"{filename.stem}.obj") - # Add the mesh to assets. - etree.SubElement(asset_elem, "mesh", file=str(meshname)) - # Add the geom to the worldbody. - if process_mtl: - e_ = etree.SubElement( - obj_body, "geom", material=material.name, mesh=str(meshname.stem) - ) - e_.attrib["class"] = "visual" - else: - e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) - e_.attrib["class"] = "visual" - else: - for i, (name, geom) in enumerate(mesh.geometry.items()): - meshname = Path(f"{filename.stem}_{i}.obj") - # Add the mesh to assets. - etree.SubElement(asset_elem, "mesh", file=str(meshname)) - # Add the geom to the worldbody. - if process_mtl: - e_ = etree.SubElement( - obj_body, "geom", mesh=meshname.stem, material=name - ) - e_.attrib["class"] = "visual" - else: - e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) - e_.attrib["class"] = "visual" - - # Add collision geoms. - if decomp_success: - # Find collision files from the decomposed convex hulls. - collisions = [ - x for x in work_dir.glob("**/*") if x.is_file() and "collision" in x.name - ] - collisions.sort(key=lambda x: int(x.stem.split("_")[-1])) - - for collision in collisions: - etree.SubElement(asset_elem, "mesh", file=collision.name) - e_ = etree.SubElement(obj_body, "geom", mesh=collision.stem) - e_.attrib["class"] = "collision" - else: - # If no decomposed convex hulls were created, use the original mesh as the - # collision mesh. - if isinstance(mesh, trimesh.base.Trimesh): - e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) - e_.attrib["class"] = "collision" - else: - for i, (name, geom) in enumerate(mesh.geometry.items()): - meshname = Path(f"{filename.stem}_{i}.obj") - e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) - e_.attrib["class"] = "collision" - - tree = etree.ElementTree(root) - etree.indent(tree, space=_XML_INDENTATION, level=0) + builder = MJCFBuilder(filename, mesh, mtls, decomp_success=decomp_success) + builder.build() # Compile and step the physics to check for any errors. if args.compile_model: - try: - tmp_path = work_dir / "tmp.xml" - tree.write(tmp_path, encoding="utf-8") - model = mujoco.MjModel.from_xml_path(str(tmp_path)) - data = mujoco.MjData(model) - mujoco.mj_step(model, data) - cprint(f"{filename} compiled successfully!", "green") - except Exception as e: - cprint(f"Error compiling model: {e}", "red") - finally: - if tmp_path.exists(): - tmp_path.unlink() + builder.compile_model() # Dump. if args.save_mjcf: - xml_path = str(work_dir / f"{filename.stem}.xml") - tree.write(xml_path, encoding="utf-8") - logging.info(f"Saved MJCF to {xml_path}") + builder.save_mjcf() def main() -> None: diff --git a/obj2mjcf/material.py b/obj2mjcf/material.py new file mode 100644 index 0000000..a1c83f4 --- /dev/null +++ b/obj2mjcf/material.py @@ -0,0 +1,74 @@ +"""A class for handling MuJoCo material properties.""" + +from dataclasses import dataclass +from typing import Optional, Sequence + +# MTL fields relevant to MuJoCo. +_MTL_FIELDS = ( + # Ambient, diffuse and specular colors. + "Ka", + "Kd", + "Ks", + # d or Tr are used for the rgba transparency. + "d", + "Tr", + # Shininess. + "Ns", + # References a texture file. + "map_Kd", +) + +# Character used to denote a comment in an MTL file. +_MTL_COMMENT_CHAR = "#" + + +@dataclass +class Material: + name: str + Ka: Optional[str] = None + Kd: Optional[str] = None + Ks: Optional[str] = None + d: Optional[str] = None + Tr: Optional[str] = None + Ns: Optional[str] = None + map_Kd: Optional[str] = None + + @staticmethod + def from_string(lines: Sequence[str]) -> "Material": + """Construct a Material object from a string.""" + attrs = {"name": lines[0].split(" ")[1].strip()} + for line in lines[1:]: + for attr in _MTL_FIELDS: + if line.startswith(attr): + elems = line.split(" ")[1:] + elems = [elem for elem in elems if elem != ""] + attrs[attr] = " ".join(elems) + break + return Material(**attrs) + + def mjcf_rgba(self) -> str: + Kd = self.Kd or "1.0 1.0 1.0" + if self.d is not None: # alpha + alpha = self.d + elif self.Tr is not None: # 1 - alpha + alpha = str(1.0 - float(self.Tr)) + else: + alpha = "1.0" + # TODO(kevin): Figure out how to use Ka for computing rgba. + return f"{Kd} {alpha}" + + def mjcf_shininess(self) -> str: + if self.Ns is not None: + # Normalize Ns value to [0, 1]. Ns values normally range from 0 to 1000. + Ns = float(self.Ns) / 1_000 + else: + Ns = 0.5 + return f"{Ns}" + + def mjcf_specular(self) -> str: + if self.Ks is not None: + # Take the average of the specular RGB values. + Ks = sum(list(map(float, self.Ks.split(" ")))) / 3 + else: + Ks = 0.5 + return f"{Ks}" diff --git a/obj2mjcf/mjcf_builder.py b/obj2mjcf/mjcf_builder.py new file mode 100644 index 0000000..1c1070e --- /dev/null +++ b/obj2mjcf/mjcf_builder.py @@ -0,0 +1,257 @@ +""" +MJCFBuilder.py +Description: + This file contains the class that is used to build MJCF files. +""" + +import logging +from pathlib import Path +from typing import Any, List, Union + +import mujoco +import trimesh +from lxml import etree +from termcolor import cprint + +from obj2mjcf.material import Material + +# 2-space indentation for the generated XML. +_XML_INDENTATION = " " + + +class MJCFBuilder: + def __init__( + self, + filename: Path, + mesh: Union[trimesh.base.Trimesh, Any], + materials: List[Material], + work_dir: Path = Path(), + decomp_success: bool = False, + ): + self.filename = filename + self.mesh = mesh + self.materials = materials + self.decomp_success = decomp_success + + self.work_dir = work_dir + if self.work_dir == Path(): + self.work_dir = filename.parent / filename.stem + + # Define variables that will be defined later + self.tree = None + + def add_visual_and_collision_default_classes( + self, + root: etree.Element, + ): + # Define the default element + default_elem = etree.SubElement(root, "default") + + # Define visual defaults + visual_default_elem = etree.SubElement(default_elem, "default") + visual_default_elem.attrib["class"] = "visual" + etree.SubElement( + visual_default_elem, + "geom", + group="2", + type="mesh", + contype="0", + conaffinity="0", + ) + + # Define collision defaults + collision_default_elem = etree.SubElement(default_elem, "default") + collision_default_elem.attrib["class"] = "collision" + etree.SubElement(collision_default_elem, "geom", group="3", type="mesh") + + def add_assets(self, root: etree.Element, mtls: List[Material]) -> etree.Element: + # Define the assets element + asset_elem = etree.SubElement(root, "asset") + + for material in mtls: + if material.map_Kd is not None: + # Create the texture asset. + texture = Path(material.map_Kd) + etree.SubElement( + asset_elem, + "texture", + type="2d", + name=texture.stem, + file=texture.name, + ) + # Reference the texture asset in a material asset. + etree.SubElement( + asset_elem, + "material", + name=material.name, + texture=texture.stem, + specular=material.mjcf_specular(), + shininess=material.mjcf_shininess(), + ) + else: + etree.SubElement( + asset_elem, + "material", + name=material.name, + specular=material.mjcf_specular(), + shininess=material.mjcf_shininess(), + rgba=material.mjcf_rgba(), + ) + + return asset_elem + + def add_visual_geometries( + self, + obj_body: etree.Element, + asset_elem: etree.Element, + ): + # Constants + filename = self.filename + mesh = self.mesh + materials = self.materials + + process_mtl = len(materials) > 0 + + # Add visual geometries to object body + if isinstance(mesh, trimesh.base.Trimesh): + meshname = Path(f"{filename.stem}.obj") + # Add the mesh to assets. + etree.SubElement(asset_elem, "mesh", file=str(meshname)) + # Add the geom to the worldbody. + if process_mtl: + e_ = etree.SubElement( + obj_body, + "geom", + material=materials[0].name, + mesh=str(meshname.stem), + ) + e_.attrib["class"] = "visual" + else: + e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) + e_.attrib["class"] = "visual" + else: + for i, (name, geom) in enumerate(mesh.geometry.items()): + meshname = Path(f"{filename.stem}_{i}.obj") + # Add the mesh to assets. + etree.SubElement(asset_elem, "mesh", file=str(meshname)) + # Add the geom to the worldbody. + if process_mtl: + e_ = etree.SubElement( + obj_body, "geom", mesh=meshname.stem, material=name + ) + e_.attrib["class"] = "visual" + else: + e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) + e_.attrib["class"] = "visual" + + def add_collision_geometries( + self, + obj_body: etree.Element, + asset_elem: etree.Element, + ): + # Constants + filename = self.filename + mesh = self.mesh + decomp_success = self.decomp_success + + work_dir = self.work_dir + + if decomp_success: + # Find collision files from the decomposed convex hulls. + collisions = [ + x + for x in work_dir.glob("**/*") + if x.is_file() and "collision" in x.name + ] + collisions.sort(key=lambda x: int(x.stem.split("_")[-1])) + + for collision in collisions: + etree.SubElement(asset_elem, "mesh", file=collision.name) + e_ = etree.SubElement(obj_body, "geom", mesh=collision.stem) + e_.attrib["class"] = "collision" + else: + # If no decomposed convex hulls were created, use the original mesh as the + # collision mesh. + if isinstance(mesh, trimesh.base.Trimesh): + meshname = Path(f"{filename.stem}.obj") + e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) + e_.attrib["class"] = "collision" + else: + for i, (name, geom) in enumerate(mesh.geometry.items()): + meshname = Path(f"{filename.stem}_{i}.obj") + e_ = etree.SubElement(obj_body, "geom", mesh=meshname.stem) + e_.attrib["class"] = "collision" + + def build( + self, + add_free_joint: bool = False, + ) -> None: + # Constants + filename = self.filename + mtls = self.materials + + # Start assembling xml tree + root = etree.Element("mujoco", model=filename.stem) + + # Add Defaults + Assets + self.add_visual_and_collision_default_classes(root) + asset_elem = self.add_assets(root, mtls) + + # Add Worldbody + worldbody_elem = etree.SubElement(root, "worldbody") + obj_body = etree.SubElement(worldbody_elem, "body", name=filename.stem) + if add_free_joint: + etree.SubElement(obj_body, "freejoint") + + # Add visual and collision geometries to object body + self.add_visual_geometries(obj_body, asset_elem) + self.add_collision_geometries(obj_body, asset_elem) + + # Collect Tree + tree = etree.ElementTree(root) + etree.indent(tree, space=_XML_INDENTATION, level=0) + + self.tree = tree + + def compile_model(self): + # Constants + filename = self.filename + work_dir = self.work_dir + + # Pull up tree if possible + tree = self.tree + if tree is None: + raise ValueError("Tree has not been defined yet.") + + # Create the work directory if it does not exist. + try: + tmp_path = work_dir / "tmp.xml" + tree.write(tmp_path, encoding="utf-8") + model = mujoco.MjModel.from_xml_path(str(tmp_path)) + data = mujoco.MjData(model) + mujoco.mj_step(model, data) + cprint(f"{filename} compiled successfully!", "green") + except Exception as e: + cprint(f"Error compiling model: {e}", "red") + finally: + if tmp_path.exists(): + tmp_path.unlink() + + def save_mjcf( + self, + ): + # Constants + filename = self.filename + work_dir = self.work_dir + + # Input Processing + + # Pull up tree if possible + tree = self.tree + if tree is None: + raise ValueError("Tree has not been defined yet.") + + # Save the MJCF file. + xml_path = str(work_dir / f"{filename.stem}.xml") + tree.write(xml_path, encoding="utf-8") + logging.info(f"Saved MJCF to {xml_path}")