feat: Added image extraction

src/readyplayerme/meshops/image.py

@@ -1,8 +1,8 @@
-from readyplayerme.meshops.types import Color, ColorMode
-from readyplayerme.meshops.types import Image as IMG_type
+"""Module for dealing with colors and image manipulation."""
+from readyplayerme.meshops.types import Color, ColorMode, Image
 
 
-def get_image_color_mode(image: IMG_type) -> ColorMode:
+def get_image_color_mode(image: Image) -> ColorMode:
     """
     Determine the color mode of an image.
 
@@ -52,7 +52,7 @@ def get_color_array_color_mode(color_array: Color) -> ColorMode:
             raise ValueError(msg)
 
 
-def blend_images(image1: IMG_type, image2: IMG_type, mask: IMG_type) -> IMG_type:
+def blend_images(image1: Image, image2: Image, mask: Image) -> Image:
     """
     Blend two images using a mask.
 

src/readyplayerme/meshops/material.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from trimesh.visual.material import PBRMaterial
+
+from readyplayerme.meshops.types import Color, Image
+
+
+@dataclass()
+class Material:
+    """Structural type for a PBR material class.
+
+    Any class considered a material must structurally be compatible with this class.
+
+    Attributes are named after the GLTF 2.0 spec.
+    """
+
+    name: str
+    baseColorTexture: Image | None  # noqa: N815
+    baseColorFactor: Color  # noqa: N815
+    doubleSided: bool  # noqa: N815
+    emissiveTexture: Image | None  # noqa: N815
+    emissiveFactor: npt.NDArray[np.float64]  # noqa: N815  # Shape (3,)
+    metallicRoughnessTexture: Image | None  # noqa: N815  # Can have occlusion in red channel.
+    metallicFactor: float  # noqa: N815
+    roughnessFactor: float  # noqa: N815
+    normalTexture: Image | None  # noqa: N815
+    occlusionTexture: Image | None  # noqa: N815
+
+    @staticmethod
+    def from_trimesh_material(material: PBRMaterial) -> "Material":
+        """Create a Material object from a trimesh Material object."""
+        base_color_texture = np.array(material.baseColorTexture) if material.baseColorTexture else None
+        emissive_texture = np.array(material.emissiveTexture) if material.emissiveTexture else None
+        metallic_roughness_texture = (
+            np.array(material.metallicRoughnessTexture) if material.metallicRoughnessTexture else None
+        )
+        normal_texture = np.array(material.normalTexture) if material.normalTexture else None
+        occlusion_texture = np.array(material.occlusionTexture) if material.occlusionTexture else None
+
+        return Material(
+            name=material.name,
+            baseColorTexture=base_color_texture,
+            baseColorFactor=material.baseColorFactor or np.array([255, 255, 255, 255], dtype=np.uint8),
+            doubleSided=material.doubleSided or False,
+            emissiveTexture=emissive_texture,
+            emissiveFactor=material.emissiveFactor or np.array([0.0, 0.0, 0.0], dtype=np.float64),
+            metallicRoughnessTexture=metallic_roughness_texture,
+            metallicFactor=material.metallicFactor or (0.0 if metallic_roughness_texture is None else 1.0),
+            roughnessFactor=material.roughnessFactor or 1.0,
+            normalTexture=normal_texture,
+            occlusionTexture=occlusion_texture,
+        )

src/readyplayerme/meshops/mesh.py

@@ -1,12 +1,28 @@
 """Functions to handle mesh data and read it from file."""
 from collections.abc import Callable
+from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
 import trimesh
 from scipy.spatial import cKDTree
 
-from readyplayerme.meshops.types import Color, Edges, IndexGroups, Indices, Mesh, PixelCoord, UVs, Vertices
+from readyplayerme.meshops.material import Material
+from readyplayerme.meshops.types import Color, Edges, Faces, IndexGroups, Indices, PixelCoord, UVs, Vertices
+
+
+@dataclass()
+class Mesh:
+    """Mesh data type.
+
+    This class serves as an abstraction for loading mesh data from different file formats.
+    """
+
+    vertices: Vertices
+    uv_coords: UVs | None
+    edges: Edges
+    faces: Faces
+    material: Material | None
 
 
 def read_mesh(filename: str | Path) -> Mesh:
@@ -33,7 +49,27 @@ def read_gltf(filename: str | Path) -> Mesh:
     :param filename: The path to the glTF file to be loaded.
     :return: The loaded mesh object.
     """
-    return trimesh.load(filename, process=False, force="mesh")
+    try:
+        loaded = trimesh.load(filename, process=False, force="mesh")
+    except ValueError as error:
+        msg = f"Error loading {filename}: {error}"
+        raise OSError(msg) from error
+    # Convert the loaded trimesh into a Mesh object for abstraction.
+    try:
+        uvs = loaded.visual.uv  # Fails if it has ColorVisuals instead of TextureVisuals.
+    except AttributeError:
+        uvs = None
+    try:
+        material = Material.from_trimesh_material(loaded.visual.material)
+    except AttributeError:
+        material = None
+    return Mesh(
+        vertices=np.array(loaded.vertices),
+        uv_coords=uvs,
+        faces=np.array(loaded.faces),
+        edges=loaded.edges,
+        material=material,
+    )
 
 
 def get_boundary_vertices(edges: Edges) -> Indices:

src/readyplayerme/meshops/types.py

@@ -1,6 +1,6 @@
 """Custom types for meshops."""
 from enum import Enum
-from typing import Protocol, TypeAlias
+from typing import TypeAlias
 
 import numpy as np
 import numpy.typing as npt
@@ -22,19 +22,10 @@ class ColorMode(Enum):
     RGBA = 4
 
 
-Image: TypeAlias = npt.NDArray[np.float32] | npt.NDArray[np.float64]  # Shape (h, w, c)
+# The Image type is based on numpy arrays for compatibility with skimage. Floats are used to allow NANs,
+# which are not supported by uint8, but the range of floats is supposed to be [0, 255] for colors and [0, 1] for masks.
+Image: TypeAlias = npt.NDArray[np.uint8] | npt.NDArray[np.float32] | npt.NDArray[np.float64]  # Shape (h, w, c)
 
 
 UVs: TypeAlias = npt.NDArray[np.float32] | npt.NDArray[np.float64]  # Shape (i, 2)
 PixelCoord: TypeAlias = npt.NDArray[np.uint16]  # Shape (i, 2)
-
-
-class Mesh(Protocol):
-    """Structural type for a mesh class.
-
-    Any class considered a mesh must structurally be compatible with this protocol.
-    """
-
-    vertices: Vertices
-    edges: Edges
-    faces: Faces

tests/readyplayerme/meshops/conftest.py

@@ -1,5 +1,6 @@
 """Pytest fixtures for meshops tests."""
 import pytest
+import skimage as ski
 
 
 @pytest.fixture
@@ -10,3 +11,24 @@ def gltf_simple_file():
     import tests.mocks
 
     return files(tests.mocks).joinpath("uv-seams.glb")
+
+
+@pytest.fixture
+def gltf_file_with_basecolor_texture():
+    """Return a path to a glTF file that contains a baseColorTexture."""
+    from importlib.resources import files
+
+    import tests.mocks
+
+    return files(tests.mocks).joinpath("input-mesh.glb")
+
+
+@pytest.fixture
+def mock_image():
+    """Return an image as a numpy array."""
+    from importlib.resources import files
+
+    import tests.mocks
+
+    filepath = files(tests.mocks).joinpath("input-img.png")
+    return ski.io.imread(filepath)

tests/readyplayerme/meshops/integration/test_integration.py

@@ -1,6 +1,8 @@
 from pathlib import Path
+from typing import Any
 
 import numpy as np
+import numpy.typing as npt
 
 import readyplayerme.meshops.mesh as mops
 
@@ -13,3 +15,16 @@ def test_boundary_vertices_from_file(gltf_simple_file: str | Path):
     assert np.array_equiv(
         np.sort(boundary_vertices), [0, 2, 4, 6, 7, 9, 10]
     ), "The vertices returned by get_border_vertices do not match the expected vertices."
+
+
+def test_get_basecolor_texture(gltf_file_with_basecolor_texture: str | Path, mock_image: npt.NDArray[Any]):
+    """Test the integration of extracting an image from GLB."""
+    local_mesh = mops.read_mesh(gltf_file_with_basecolor_texture)
+    extracted_image = local_mesh.material.baseColorTexture
+    assert np.array_equal(extracted_image, mock_image), "The extracted image does not match the expected image."
+
+
+def test_access_material_should_fail(gltf_simple_file: str | Path):
+    """Test the simple gltf does not have a material."""
+    mesh = mops.read_mesh(gltf_simple_file)
+    assert mesh.material is None, "Mesh should not have a material."

tests/readyplayerme/meshops/unit/test_mesh.py

@@ -8,7 +8,7 @@
 import pytest
 
 from readyplayerme.meshops import mesh
-from readyplayerme.meshops.types import IndexGroups, Indices, Mesh, PixelCoord, UVs, Vertices
+from readyplayerme.meshops.types import IndexGroups, Indices, PixelCoord, UVs, Vertices
 
 
 class TestReadMesh:
@@ -18,7 +18,7 @@ class TestReadMesh:
         "filepath, expected",
         [("test.glb", mesh.read_gltf), (Path("test.glb"), mesh.read_gltf), ("test.gltf", mesh.read_gltf)],
     )
-    def test_get_mesh_reader_glb(self, filepath: str | Path, expected: Callable[[str | Path], Mesh]):
+    def test_get_mesh_reader_glb(self, filepath: str | Path, expected: Callable[[str | Path], mesh.Mesh]):
         """Test the get_mesh_reader function with a .glb file path."""
         reader = mesh.get_mesh_reader(filepath)
         assert callable(reader), "The mesh reader should be a callable function."
@@ -34,23 +34,26 @@ def test_read_mesh_gltf(self, gltf_simple_file: str | Path):
         """Test the read_gltf function returns the expected type."""
         result = mesh.read_gltf(gltf_simple_file)
         # Check the result has all the expected attributes and of the correct type.
-        for attr in Mesh.__annotations__:
+        for attr in mesh.Mesh.__annotations__:
             assert hasattr(result, attr), f"The mesh class should have a '{attr}' attribute."
             # Find the type so we can use it as a second argument to isinstance.
-            tp = get_origin(Mesh.__annotations__[attr])
+            tp = get_origin(mesh.Mesh.__annotations__[attr])
             if tp is Union or tp is types.UnionType:
-                tp = get_args(Mesh.__annotations__[attr])
+                tp = get_args(mesh.Mesh.__annotations__[attr])
                 # Loop through the types in the union and check if the result is compatible with any of them.
-                assert any(
-                    isinstance(getattr(result, attr), get_origin(t)) for t in tp
-                ), f"The '{attr}' attribute should be compatible with {tp}."
+                matched = False
+                for type_ in tp:
+                    # Check original type definition. isinstance doesn't work with None, but with NoneType.
+                    origin = o if (o := get_origin(type_)) is not None else type_
+                    matched |= isinstance(getattr(result, attr), origin)
+                assert matched, f"The '{attr}' attribute should be compatible with {tp}."
             else:
                 assert isinstance(
                     getattr(result, attr), tp
-                ), f"The '{attr}' attribute should be compatible with {Mesh.__annotations__[attr]}."
+                ), f"The '{attr}' attribute should be compatible with {mesh.Mesh.__annotations__[attr]}."
 
 
-def test_get_boundary_vertices(mock_mesh: Mesh):
+def test_get_boundary_vertices(mock_mesh: mesh.Mesh):
     """Test the get_boundary_vertices function returns the expected indices."""
     boundary_vertices = mesh.get_boundary_vertices(mock_mesh.edges)