Basic wavefront obj loader

wavefront.py

import logging
import os.path

import numpy as np
from PIL import Image


logger = logging.getLogger(__name__)


def load(file):
    vertices = []
    vertex_color = []
    normals = []
    texcoords = []
    face_v_indices = []
    face_t_indices = []
    face_n_indices = []
    texture = None

    if not hasattr(file, "read"):
        f = open(file, "r")
    else:
        f = file

    for line in f:
        if line.startswith("#"):
            continue
        values = line.split()
        if not values:
            continue
        if values[0] == "v":
            v = [float(v) for v in values[1:4]]
            vertices.append(v)
            if len(values) == 7:
                vertex_color.append([float(c) for c in values[4:7]])
        elif values[0] == "vn":
            v = [float(v) for v in values[1:4]]
            normals.append(v)
        elif values[0] == "vt":
            texcoords.append([float(v) for v in values[1:3]])
        elif values[0] == "f":
            vert = []
            text = []
            norm = []
            for v in values[1:]:
                w = v.split("/")
                vert.append(int(w[0]) - 1)
                if len(w) > 1 and w[1]:
                    text.append(int(w[1]) - 1)
                if len(w) > 2 and w[2]:
                    norm.append(int(w[2]) - 1)
            if vert:
                face_v_indices.append(vert)
            if text:
                face_t_indices.append(text)
            if norm:
                face_n_indices.append(norm)

    if texcoords:
        try:
            texture_path = file.replace(".obj", ".png")
            logger.debug("trying to load texture {}".format(texture_path))
            texture = np.array(Image.open(texture_path))
        except IOError:
            logger.debug("texture not found")
            pass

    faces = np.array(face_v_indices, dtype=int) if face_v_indices else None
    t_indices = np.array(face_t_indices, dtype=int) if face_t_indices else None
    n_indices = np.array(face_n_indices, dtype=int) if face_n_indices else None

    data = {
        "vertices": np.array(vertices) if vertices else None,
        "vertex_color": np.array(vertex_color) if vertex_color else None,
        "faces": faces,
        "normals": np.array(normals) if normals else None,
        "texcoords": np.array(texcoords) if texcoords else None,
        "texcoords_indices": t_indices,
        "normals_indices": n_indices,
        "texture": texture,
    }

    return data


def _interleave_columns(*arrays):
    """Interleave the columns of a sequence of arrays with equal dimensions.

    Returns
    -------
    array of shape  (M, K * N)
        Array with the columns of the K input arrays interleaved.
    """
    n_rows = arrays[0].shape[0]
    n_cols = len(arrays) * arrays[0].shape[1]
    return np.vstack(arrays).reshape(n_rows, n_cols, order="F")


MTL_CONTENT = """\
newmtl material_0
Ka 0.7 0.7 0.7
Kd 1.0 1.0 1.0
Ks 1.0 1.0 1.0
Tr 0.0
illum 2
Ns 0.0
"""


def _write_mtl(path, texture_filename=None):
    content = MTL_CONTENT
    if texture_filename is not None:
        content += "map_Kd {}".format(texture_filename)

    with open(path, "w") as f:
        f.write(content)


def _save_texture(path, texture):
    Image.fromarray(texture).save(path)


def save(path, vertices=None, faces=None, normals=None, texcoords=None,
         texcoords_indices=None, normals_indices=None,
         texture=None, vertex_color=None):
    if not path.endswith(".obj"):
        path = path + ".obj"
    mtl_path = path.replace(".obj", ".mtl")
    mtl_filename = os.path.basename(mtl_path)
    texture_path = path.replace(".obj", ".png")
    texture_filename = os.path.basename(texture_path)

    if texture is not None:
        _write_mtl(mtl_path, texture_filename=texture_filename)
        _save_texture(texture_path, texture)
    else:
        _write_mtl(mtl_path)

    with open(path, "w") as f:
        f.write("mtllib {}\n".format(mtl_filename))
        f.write("usemtl material_0\n")
        if vertices is not None and vertex_color is None:
            np.savetxt(f, vertices, fmt="v %f %f %f", delimiter=" ")
        elif vertices is not None and vertex_color is not None:
            # vc = vertex_color.astype(float) / 255
            v_c = np.hstack((vertices, vertex_color))
            np.savetxt(f, v_c, fmt="v %f %f %f %f %f %f", delimiter=" ")
        if normals is not None:
            np.savetxt(f, normals, fmt="vn %f %f %f", delimiter=" ")
        if texcoords is not None:
            np.savetxt(f, texcoords, fmt="vt %f %f", delimiter=" ")
        if faces is not None:
            if texcoords_indices is None and normals_indices is None:

                np.savetxt(f, faces + 1, fmt="f %d %d %d", delimiter=" ")
            elif texcoords_indices is not None \
                    and normals_indices is None:
                f_t = _interleave_columns(faces, texcoords_indices)
                f_t += 1
                np.savetxt(f, f_t, fmt="f %d/%d %d/%d %d/%d", delimiter=" ")
            elif texcoords_indices is not None \
                    and normals_indices is not None:
                f_t_n = _interleave_columns(faces, texcoords_indices,
                                            normals_indices)
                f_t_n += 1
                np.savetxt(f, f_t_n, fmt="f %d/%d/%d %d/%d/%d %d/%d/%d",
                           delimiter=" ")
            elif texcoords_indices is None \
                    and normals_indices is not None:
                f_n = _interleave_columns(faces, normals_indices)
                f_n += 1
                np.savetxt(f, f_n, fmt="f %d//%d %d//%d %d//%d",
                           delimiter=" ")