Converts a PSD or PSB file to a layered shadeless GLTF file

psd_to_gltf.py

# Depends on psd-tools pygltflib

from psd_tools import PSDImage
import pygltflib
from pygltflib import GLTF2
import sys
import os
import numpy as np


SIZE = 10
PER_LAYER_Z_OFFSET = 0.5


def make_plane_mesh(material):
    gltf.meshes.append(pygltflib.Mesh(
        primitives=[
                pygltflib.Primitive(
                    attributes=pygltflib.Attributes(POSITION=1, TEXCOORD_0=2),
                    indices=0,
                    material=material,
                )
            ]
    ))
    return len(gltf.meshes) - 1


def make_material(file_name):
    gltf.images.append(pygltflib.Image(
        mimeType="image/png",
        name=file_name,
        uri=file_name,
    ))
    gltf.textures.append(pygltflib.Texture(
        source=len(gltf.images) - 1,
        # sampler=0,
    ))
    gltf.materials.append(pygltflib.Material(
        alphaMode=pygltflib.BLEND,
        doubleSided=True,
        name=file_name,
        pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
            baseColorTexture=pygltflib.TextureInfo(index=len(gltf.textures) - 1)
        )
    ))
    return len(gltf.materials) - 1


def make_single_mesh():
    triangles = np.array(
        [
            [0, 1, 2],
            [2, 1, 3],
        ],
        dtype="uint8",
    )
    triangles_binary_blob = triangles.flatten().tobytes()
    points = np.array(
        [
            [0, 0, 0],
            [1, 0, 0],
            [0, -1, 0],
            [1, -1, 0],
        ],
        dtype="float32",
    )
    points_binary_blob = points.tobytes()
    tex_points = np.array(
        [
            [0, 0],
            [1, 0],
            [0, 1],
            [1, 1],
        ],
        dtype="float32",
    )
    tex_points_binary_blob = tex_points.tobytes()
    gltf.accessors.append(pygltflib.Accessor(
        bufferView=0,
        componentType=pygltflib.UNSIGNED_BYTE,
        count=triangles.size,
        type=pygltflib.SCALAR,
        # max=[int(triangles.max())],
        # min=[int(triangles.min())],
    ))
    gltf.accessors.append(pygltflib.Accessor(
        bufferView=1,
        componentType=pygltflib.FLOAT,
        count=len(points),
        type=pygltflib.VEC3,
        max=points.max(axis=0).tolist(),
        min=points.min(axis=0).tolist(),
    ))
    gltf.accessors.append(pygltflib.Accessor(
        bufferView=2,
        componentType=pygltflib.FLOAT,
        count=len(tex_points),
        type=pygltflib.VEC2,
        # max=tex_points.max(axis=0).tolist(),
        # min=tex_points.min(axis=0).tolist(),
    ))
    gltf.bufferViews.append(pygltflib.BufferView(
        buffer=0,
        byteLength=len(triangles_binary_blob),
        target=pygltflib.ELEMENT_ARRAY_BUFFER,
    ))
    gltf.bufferViews.append(pygltflib.BufferView(
        buffer=0,
        byteOffset=len(triangles_binary_blob),
        byteLength=len(points_binary_blob),
        target=pygltflib.ARRAY_BUFFER,
    ))
    gltf.bufferViews.append(pygltflib.BufferView(
        buffer=0,
        byteOffset=len(triangles_binary_blob) + len(points_binary_blob),
        byteLength=len(tex_points_binary_blob),
        target=pygltflib.ARRAY_BUFFER,
    ))

    binary_blob = triangles_binary_blob + points_binary_blob + tex_points_binary_blob

    gltf.buffers.append(pygltflib.Buffer(
        byteLength=len(binary_blob)
    ))

    gltf.set_binary_blob(binary_blob)


if __name__ == "__main__":
    to_load = sys.argv[1]
    to_save = sys.argv[2]
    print("LOADING", repr(to_load))

    psd = PSDImage.open(to_load)

    scene = pygltflib.Scene()
    gltf = pygltflib.GLTF2(
        scene=0,
        scenes=[scene]
    )
    make_single_mesh()

    saved_images = set()

    (img_left, img_top, img_right, img_bottom) = psd.viewbox
    img_width = img_right - img_left
    img_height = img_bottom - img_top
    scale = SIZE / max(img_width, img_height)

    print("SCALE", scale)

    for layer_idx, layer in enumerate(psd.descendants()):
        file_name = f"{layer.name}.png"
        i = 1
        while file_name in saved_images:
            i += 1
            file_name = f"{layer.name} {i}.png"
        saved_images.add(file_name)

        material_idx = make_material(file_name)
        mesh_idx = make_plane_mesh(material_idx)
        gltf.nodes.append(pygltflib.Node(
            mesh=mesh_idx,
            name=layer.name,
            translation=[layer.left * scale, -layer.top * scale, layer_idx * PER_LAYER_Z_OFFSET],
            scale=[layer.width * scale, layer.height * scale, 1],
        ))
        scene.nodes.append(len(gltf.nodes) - 1)

        print(layer)
        file_path = os.path.join(to_save, file_name)
        if not os.path.exists(file_path):
            layer_image = layer.composite()
            print("  Generated")
            layer_image.save(file_path)
            print("  Written")
        else:
            print("  Skipped")


    os.makedirs(to_save, exist_ok=True)
    gltf.save(os.path.join(to_save, "layers.gltf"))