modify write

face3d/mesh/io.py

@@ -23,13 +23,13 @@ def write_asc(path, vertices):
     else:
         np.savetxt(path + '.asc', vertices)
 
-def write_obj_with_colors(obj_name, vertices, colors, triangles):
-    ''' Save 3D face model
+def write_obj_with_colors(obj_name, vertices, triangles, colors):
+    ''' Save 3D face model with texture represented by colors.
     Args:
         obj_name: str
         vertices: shape = (nver, 3)
-        colors: shape = (nver, 3)
         triangles: shape = (ntri, 3)
+        colors: shape = (nver, 3)
     '''
     triangles = triangles.copy()
     triangles += 1 # meshlab start with 1
@@ -54,8 +54,60 @@ def write_obj_with_colors(obj_name, vertices, colors, triangles):
             f.write(s)
 
 ## TODO: c++ version
-def write_obj_with_texture(obj_name, vertices, colors, triangles, texture, uv_coords):
+def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords):
+    ''' Save 3D face model with texture represented by texture map.
+    Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
+    Args:
+        obj_name: str
+        vertices: shape = (nver, 3)
+        triangles: shape = (ntri, 3)
+        texture: shape = (256,256,3)
+        uv_coords: shape = (nver, 3) max value<=1
+    '''
+    if obj_name.split('.')[-1] != 'obj':
+        obj_name = obj_name + '.obj'
+    mtl_name = obj_name.replace('.obj', '.mtl')
+    texture_name = obj_name.replace('.obj', '_texture.png')
+
+    triangles = triangles.copy()
+    triangles += 1 # mesh lab start with 1
+
+    # write obj
+    with open(obj_name, 'w') as f:
+        # first line: write mtlib(material library)
+        s = "mtllib {}\n".format(os.path.abspath(mtl_name))
+        f.write(s)
+
+        # write vertices
+        for i in range(vertices.shape[0]):
+            s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2])
+            f.write(s)
+
+        # write uv coords
+        for i in range(uv_coords.shape[0]):
+            s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1])
+            f.write(s)
+
+        f.write("usemtl FaceTexture\n")
+
+        # write f: ver ind/ uv ind
+        for i in range(triangles.shape[0]):
+            s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0])
+            f.write(s)
+
+    # write mtl
+    with open(mtl_name, 'w') as f:
+        f.write("newmtl FaceTexture\n")
+        s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
+        f.write(s)
+
+    # write texture as png
+    imsave(texture_name, texture)
+
+
+def write_obj_with_colors_texture(obj_name, vertices, colors, triangles, texture, uv_coords):
     ''' Save 3D face model with texture. 
+    Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
     Args:
         obj_name: str
         vertices: shape = (nver, 3)