Repository: mikedh/trimesh
Pull request #272 (2018-12-15 01:15:43)

Support loading ShapeNet models: support Kd in mtl and some fixes in 'vt' and 'f'
With this change, trimesh now supports loading textures and colors for models in ShapeNetSem.v0.
(This information was the key for uv_to_color with uv out of range from [0, 1]: https://shapenet.org/qaforum/index.php?qa=15&qa_1=why-is-the-texture-coordinate-in-the-obj-file-not-in-the-range)

<img src='https://drive.google.com/uc?id=1ANxzreOjd-q3m82KkYczaSyfiinNuKNf' width='50%' />

```python
import pathlib
import numpy as np
import trimesh

# trimesh.util.attach_to_log()

scene = trimesh.Scene()

models_dir = pathlib.Path('ShapeNetSemModels')
for cad_file in models_dir.glob('*.obj'):
    print(cad_file)

    cad = trimesh.load(str(cad_file))

    # centerize
    xyz_min = np.min(cad.vertices, axis=0)
    xyz_max = np.max(cad.vertices, axis=0)
    xyz_center = (xyz_max + xyz_min) / 2.
    cad.vertices -= xyz_center

    max_size = (xyz_max - xyz_min).max()
    cad.vertices /= max_size  # normalize

    print(cad.bounding_box.primitive.extents)

    # shift for visualization
    x = len(scene.geometry) % 4
    y = len(scene.geometry) // 4
    transform = trimesh.transformations.translation_matrix(
        [1.5 * x, 1.5 * y, 0]
    )
    scene.add_geometry(cad, transform=transform)

    if len(scene.geometry) == 16:
        break
scene.show()
```

Of course, it works with YCB object models:

<img src="https://drive.google.com/uc?id=1AQ73VP2MPfJBKihN7gY9CrTm8W0rLm7f" width="50%" />

```python
import pathlib
import numpy as np
import trimesh

# trimesh.util.attach_to_log()

scene = trimesh.Scene()

models_dir = pathlib.Path('YCB_Video_Models')
for obj_dir in sorted(models_dir.iterdir()):
    cad_file = obj_dir / 'textured.obj'
    print(cad_file)

    cad = trimesh.load(str(cad_file))

    # centerize
    xyz_min = np.min(cad.vertices, axis=0)
    xyz_max = np.max(cad.vertices, axis=0)
    xyz_center = (xyz_max + xyz_min) / 2.
    cad.vertices -= xyz_center

    max_size = (xyz_max - xyz_min).max()
    cad.vertices /= max_size  # normalize

    print(cad.bounding_box.primitive.extents)

    # shift for visualization
    x = len(scene.geometry) % 4
    y = len(scene.geometry) // 4
    transform = trimesh.transformations.translation_matrix(
        [1.5 * x, 1.5 * y, 0]
    )
    scene.add_geometry(cad, transform=transform)

    if len(scene.geometry) == 16:
        break
scene.show()
```


Patch:
diff --git a/trimesh/io/wavefront.py b/trimesh/io/wavefront.py
index 3dd0594a..84c3565e 100644
--- a/trimesh/io/wavefront.py
+++ b/trimesh/io/wavefront.py
@@ -1,3 +1,4 @@
+import collections
 import os.path
 
 import numpy as np
@@ -12,7 +13,7 @@ from .. import util
 from .. import visual
 
 
-def _get_vertex_colors(uv, mtllibs, resolver):
+def _get_vertex_colors(uv, mtllib, resolver):
     """
     Turn loaded UV and materials into vertex colors.
 
@@ -20,7 +21,7 @@ def _get_vertex_colors(uv, mtllibs, resolver):
     ------------
     uv : (n, 2) float
       Texture coordinates
-    mttlibs : list of dict
+    mttlib : dict
       Material data including image names
     resolver : trimesh.visual.resolvers.Resolver
       Object to load data referenced by name
@@ -34,21 +35,21 @@ def _get_vertex_colors(uv, mtllibs, resolver):
         log.warning('pillow is required to load texture')
         return
 
-    success = False
-    vertex_colors = np.zeros((len(uv), 4), dtype=np.uint8)
+    if np.isnan(uv).all() and mtllib['Kd']:
+        color = np.array(mtllib['Kd'], dtype=np.float64)
+        if len(color) == 3:
+            color = np.hstack((color, [1]))
+        color = (color * 255).astype(np.uint8)
+        vertex_colors = np.full((len(uv), 4), color, dtype=np.uint8)
+        return vertex_colors
+
+    # get the data for the image
+    texture_data = resolver.get(mtllib['map_Kd'])
+    # load the actual image
+    texture = Image.open(util.wrap_as_stream(texture_data))
+    # append the loaded colors
+    vertex_colors = visual.uv_to_color(uv, texture)
 
-    for mtl in mtllibs:
-        # get the data for the image
-        texture_data = resolver.get(mtl['map_Kd'])
-        # load the actual image
-        texture = Image.open(util.wrap_as_stream(texture_data))
-        # append the loaded colors
-        vertex_colors += visual.uv_to_color(uv, texture)
-        success = True
-
-    # nothing actually loaded
-    if not success:
-        return
     return vertex_colors
 
 
@@ -63,16 +64,15 @@ def parse_mtl(mtl):
 
     Returns
     ------------
-    data : dict
-      Has keys map_Kd
+    mtllibs : list of dict
+      Each dict has keys: newmtl, map_Kd, Kd
     """
     # decode bytes if necessary
     if hasattr(mtl, 'decode'):
         mtl = mtl.decode('utf-8')
 
-    # initial data layout
-    data = {'newmtl': None,
-            'map_Kd': None}
+    mtllib = None
+    mtllibs = []
 
     # use universal newline splitting
     for line in str.splitlines(str(mtl).strip()):
@@ -82,10 +82,21 @@ def parse_mtl(mtl):
         if len(line_split) <= 1:
             continue
         # store the keys
-        if line_split[0] in data:
-            data[line_split[0]] = line_split[1]
-
-    return data
+        key = line_split[0]
+        if key == 'newmtl':
+            if mtllib:
+                mtllibs.append(mtllib)
+            mtllib = {'newmtl': line_split[1],
+                      'map_Kd': None,
+                      'Kd': None}
+        elif key == 'map_Kd':
+            mtllib[key] = line_split[1]
+        elif key == 'Kd':
+            mtllib[key] = [float(x) for x in line_split[1:]]
+    if mtllib:
+        mtllibs.append(mtllib)
+
+    return mtllibs
 
 
 def load_wavefront(file_obj, resolver=None, **kwargs):
@@ -172,16 +183,14 @@ def load_wavefront(file_obj, resolver=None, **kwargs):
 
             # handle vertex texture
             if len(current['vt']) > 0:
-                texture = np.array(current['vt'], dtype=np.float64)
-                # make sure vertex texture is the right shape
-                # AKA (len(vertices), dimension)
-                try:
-                    texture = texture.reshape((len(vertices), -1))
-                    # save vertex texture with correct ordering
-                    loaded['metadata']['vertex_texture'] = texture[vert_order]
-                except ValueError:
-                    log.warning('texture information seems broken'.format(
-                                file_obj.name))
+                texture = np.full((len(current['vt_given']), 3),
+                                  np.nan,
+                                  dtype=np.float64)
+                texture[current['vt_given']] = current['vt']
+                texture = texture[vert_order]
+                texture = texture[~np.any(np.isnan(texture), axis=1)]
+                # save vertex texture with correct ordering
+                loaded['metadata']['vertex_texture'] = texture
 
             # build face groups information
             # faces didn't move around so we don't have to reindex
@@ -192,24 +201,36 @@ def load_wavefront(file_obj, resolver=None, **kwargs):
                     face_groups[start_f:] = idx
                 loaded['metadata']['face_groups'] = face_groups
 
-            if len(current['usemtl']) > 0:
-                current_mtllibs = [mtllibs[k] for k in current['usemtl']
-                                   if k in mtllibs]
-                try:
-                    vertex_colors = _get_vertex_colors(
-                        uv=loaded['metadata']['vertex_texture'],
-                        mtllibs=current_mtllibs,
-                        resolver=resolver)
-                    loaded['vertex_colors'] = vertex_colors
-                    loaded['metadata']['mtllibs'] = current_mtllibs
-                except BaseException:
-                    log.error('failed to load texture!', exc_info=True)
+            if len(current['usemtl']) > 0 and np.sum(current['vt_given']) > 0:
+                texture = np.full((len(current['vt_given']), 3),
+                                  np.nan,
+                                  dtype=np.float64)
+                texture[current['vt_given']] = current['vt']
+
+                vertex_colors = np.zeros((0, 4), dtype=np.uint8)
+                for usemtl in current['usemtl']:
+                    findices = usemtl_to_findices[usemtl]
+                    uv = texture[findices]
+                    try:
+                        vertex_colors_i = _get_vertex_colors(
+                            uv=uv, mtllib=mtllibs[usemtl], resolver=resolver)
+                    except BaseException:
+                        log.error('failed to load texture: {}'.format(usemtl),
+                                  exc_info=True)
+                        vertex_colors_i = None
+                    if vertex_colors_i is None:
+                        vertex_colors_i = np.full(
+                            (len(uv), 4), visual.DEFAULT_COLOR, dtype=np.uint8)
+                    vertex_colors = np.r_[vertex_colors, vertex_colors_i]
+                loaded['vertex_colors'] = vertex_colors[vert_order]
 
             # this mesh is done so append the loaded mesh kwarg dict
             meshes.append(loaded)
 
     attribs = {k: [] for k in ['v', 'vt', 'vn']}
-    current = {k: [] for k in ['v', 'vt', 'vn', 'f', 'g', 'usemtl']}
+    current = {k: [] for k in ['v', 'vt', 'vn', 'f', 'g', 'usemtl',
+                               'vt_given', 'vn_given']}
+    usemtl_to_findices = collections.defaultdict(list)  # usemtl to 'f' indices
     mtllibs = {}
     # remap vertex indexes {str key: int index}
     remap = {}
@@ -224,8 +245,12 @@ def load_wavefront(file_obj, resolver=None, **kwargs):
             # v, vt, or vn
             # vertex, vertex texture, or vertex normal
             # only parse 3 values, ignore colors
-            attribs[line_split[0]].append([float(x)
-                                           for x in line_split[1:4]])
+            value = [float(x) for x in line_split[1:4]]
+            # vt: u [v] [w]  # v, w is optional and default value is 0
+            if line_split[0] == 'vt' and len(value) != 3:
+                for _ in range(3 - len(value)):
+                    value.append(0)
+            attribs[line_split[0]].append(value)
         elif line_split[0] == 'f':
             # a face
             ft = line_split[1:]
@@ -246,15 +271,25 @@ def load_wavefront(file_obj, resolver=None, **kwargs):
                     if len(f_split) > 1 and f_split[1] != '':
                         current['vt'].append(
                             attribs['vt'][int(f_split[1]) - 1])
+                        current['vt_given'].append(True)
+                    else:
+                        current['vt_given'].append(False)
                     if len(f_split) > 2:
                         current['vn'].append(
                             attribs['vn'][int(f_split[2]) - 1])
+                        current['vn_given'].append(True)
+                    else:
+                        current['vn_given'].append(False)
+                    if len(current['usemtl']) > 0:
+                        usemtl_to_findices[current['usemtl'][-1]].append(
+                            len(current['vt']) - 1)
                 current['f'].append(remap[f])
         elif line_split[0] == 'o':
             # defining a new object
             append_mesh()
             # reset current to empty lists
             current = {k: [] for k in current.keys()}
+            usemtl_to_findices = collections.defaultdict(list)
             remap = {}
             next_idx = 0
             group_idx = 0
@@ -274,10 +309,9 @@ def load_wavefront(file_obj, resolver=None, **kwargs):
             try:
                 # fetch bytes containing MTL data
                 mtl_data = resolver.get(mtl_name)
-                # load into a dict
-                mtllib = parse_mtl(mtl_data)
-                # save new materials
-                if 'newmtl' in mtllib:
+                # load into a list of dict
+                for mtllib in parse_mtl(mtl_data):
+                    # save new materials
                     mtllibs[mtllib['newmtl']] = mtllib
             except BaseException:
                 log.error('unable to load material: {}'.format(mtl_name),
diff --git a/trimesh/visual/texture.py b/trimesh/visual/texture.py
index 20b8212d..c7e723d2 100644
--- a/trimesh/visual/texture.py
+++ b/trimesh/visual/texture.py
@@ -19,10 +19,18 @@ def uv_to_color(uv, image):
     colors : (n, 4) float
       RGBA color at each of the UV coordinates
     """
-    x = (uv[:, 0] * image.width).round().astype(int)
-    y = ((1 - uv[:, 1]) * image.height).round().astype(int)
+    # https://shapenet.org/qaforum/index.php?qa=15&qa_1=why-is-the-texture-coordinate-in-the-obj-file-not-in-the-range  # NOQA
+    if not (uv.min() >= 0 and uv.max() <= 1):
+        uv = uv - np.floor(uv)
+
+    x = (uv[:, 0] * (image.width - 1)).round().astype(int)
+    y = ((1 - uv[:, 1]) * (image.height - 1)).round().astype(int)
 
     image = np.asarray(image)
     colors = image[y, x]
-    colors = to_rgba(colors)
+    if colors.ndim == 1:  # gray scale
+        colors = np.repeat(colors[:, None], 3, axis=1)
+    # now ndim == 2
+    if colors.shape[1] == 3:
+        colors = to_rgba(colors)  # rgb -> rgba
     return colors


Test patch:
diff --git a/tests/test_color.py b/tests/test_color.py
index 5f9cff85..7ad91840 100644
--- a/tests/test_color.py
+++ b/tests/test_color.py
@@ -215,7 +215,7 @@ class VisualTest(g.unittest.TestCase):
         uv = g.np.array([[0.25, 0.2], [0.4, 0.5]], dtype=float)
         texture = g.np.arange(96, dtype=g.np.uint8).reshape(8, 4, 3)
         colors = g.trimesh.visual.uv_to_color(uv, PIL.Image.fromarray(texture))
-        colors_expected = [[75, 76, 77, 255], [54, 55, 56, 255]]
+        colors_expected = [[75, 76, 77, 255], [51, 52, 53, 255]]
 
         g.np.testing.assert_allclose(colors, colors_expected, rtol=0, atol=0)