mrfspheres.py

#!/bin/env python
"""Generate a glTF like MetalRoughSpheres, but with factors instead of textures.

Output is written to MetalRoughFactorSpheres.gltf and MetalRoughFactorSpheres.bin.
"""

from math import sqrt
import base64
import json
import struct

# http://blog.andreaskahler.com/2009/06/creating-icosphere-mesh-in-code.html
class Icosphere:

    def vert(self, p):
        mag = sqrt(p[0]**2 + p[1]**2 + p[2]**2)
        normalized = [p[0]/mag, p[1]/mag, p[2]/mag]
        self.verts.append(normalized)

    def __init__(self):
        """Build a regular icosahedron."""
        self.verts = []
        self.tris = []
        self.midpoint_cache = {}

        t = (1 + sqrt(5))/2

        self.vert([-1,  t, 0])
        self.vert([ 1,  t, 0])
        self.vert([-1, -t, 0])
        self.vert([ 1, -t, 0])

        self.vert([0, -1,  t])
        self.vert([0,  1,  t])
        self.vert([0, -1, -t])
        self.vert([0,  1, -t])

        self.vert([ t, 0, -1])
        self.vert([ t, 0,  1])
        self.vert([-t, 0, -1])
        self.vert([-t, 0,  1])

        self.tris.append([0, 11, 5])
        self.tris.append([0, 5, 1])
        self.tris.append([0, 1, 7])
        self.tris.append([0, 7, 10])
        self.tris.append([0, 10, 11])
        self.tris.append([1, 5, 9])
        self.tris.append([5, 11, 4])
        self.tris.append([11, 10, 2])
        self.tris.append([10, 7, 6])
        self.tris.append([7, 1, 8])
        self.tris.append([3, 9, 4])
        self.tris.append([3, 4, 2])
        self.tris.append([3, 2, 6])
        self.tris.append([3, 6, 8])
        self.tris.append([3, 8, 9])
        self.tris.append([4, 9, 5])
        self.tris.append([2, 4, 11])
        self.tris.append([6, 2, 10])
        self.tris.append([8, 6, 7])
        self.tris.append([9, 8, 1])

    def refine(self):
        new_tris = []
        midpoint_cache = {}

        def midpoint(v1, v2):
            v1, v2 = sorted((v1, v2))

            if (v1, v2) in midpoint_cache:
                return midpoint_cache[(v1, v2)]

            co1 = self.verts[v1]
            co2 = self.verts[v2]
            mid = [(co1[0] + co2[0])/2, (co1[1] + co2[1])/2, (co1[2] + co2[2])/2]
            self.vert(mid)
            index = len(self.verts) - 1

            midpoint_cache[(v1, v2)] = index

            return index

        for tri in self.tris:
            a = midpoint(tri[0], tri[1])
            b = midpoint(tri[1], tri[2])
            c = midpoint(tri[2], tri[0])

            new_tris.append([tri[0], a, c])
            new_tris.append([tri[1], b, a])
            new_tris.append([tri[2], c, b])
            new_tris.append([a, b, c])

        self.tris = new_tris


class MetalRoughSpheres:

    def add_buffer_data(self, data, align=1):
        """Add new data to the buffer. Returns the offset to the data.

        The buffer is padded with zeros so that the data is aligned to the
        given alignment in bytes.
        """
        off = len(self.buf)
        if off % align != 0:
            self.buf += 'b\0' * ((- off) % align)
        self.buf += data
        return off

    def finish(self, gltf_path, bin_path):
        """Write the output."""
        self.gltf['buffers'] = [{
            'uri': bin_path,
            'byteLength': len(self.buf),
        }]
        with open(gltf_path, 'w+') as fp:
            json.dump(self.gltf, fp, indent=4)
        with open(bin_path, 'wb+') as fp:
            fp.write(self.buf)

    def add_accessor(self, data, component_type, type, count, min=None, max=None, align=1):
        data_off = self.add_buffer_data(data, align=align)
        self.gltf['bufferViews'].append({
            'buffer': 0,
            'byteOffset': data_off,
            'byteLength': len(data),
        })
        buffer_view_idx = len(self.gltf['bufferViews']) - 1
        accessor = {
            'bufferView': buffer_view_idx,
            'type': type,
            'componentType': component_type,
            'count': count,
        }
        if min:
            accessor['min'] = min
        if max:
            accessor['max'] = max
        self.gltf['accessors'].append(accessor)
        return len(self.gltf['accessors']) - 1

    def add_spheres(self):
        # Build an icosphere

        icosphere = Icosphere()
        for _ in range(0, 4):
            icosphere.refine()
        verts = [c for v in icosphere.verts for c in v]
        indices = [i for tri in icosphere.tris for i in tri]

        assert(len(indices) < 2**16) # Make sure we can use u16 for the index data

        # Write buffers/bufferViews/accessors for the icosphere data.
        # We'll reuse one set of accessors for all the different meshes.

        vert_data = struct.pack('<%sf' % len(verts), *verts)
        index_data = struct.pack('<%sH' % len(indices), *indices)
        vert_accessor = self.add_accessor(
            vert_data,
            component_type=5126,
            type='VEC3',
            count=len(verts)/3,
            min=[-1,-1,-1],
            max=[1,1,1],
            align=4,
        )
        index_accessor = self.add_accessor(
            index_data,
            component_type=5123,
            type='SCALAR',
            count=len(indices),
            align=2,
        )

        # Create the grid of spheres; one material/mesh pair for every
        # sphere in the grid.

        num_spheres = 7 * 7 * 2
        self.gltf['scene'] = 0
        self.gltf['scenes'] = [ { 'nodes': [0] } ]
        self.gltf['nodes'] = [{
            'children': list(range(1, 1+num_spheres)),
        }]
        grey = [204/255, 204/255, 204/255, 1]
        yellow = [204/255, 177/255, 29/255, 1]
        for color, z in [(grey, 0), (yellow, 8)]:
            for i in range(0, 7):
                for j in range(0, 7):
                    metal = i / 6
                    rough = j / 6
                    self.gltf['materials'].append({
                        'pbrMetallicRoughness': {
                            'baseColorFactor': color,
                            'metallicFactor': metal,
                            'roughnessFactor': rough,
                        }
                    })
                    material_id = len(self.gltf['materials']) - 1
                    x = -3*j + 9
                    y = 3*i - 9
                    self.gltf['meshes'].append({
                        'primitives': [{
                            'attributes': {
                                'POSITION': vert_accessor,
                                'NORMAL': vert_accessor,
                            },
                            'indices': index_accessor,
                            'material': material_id,
                        }]
                    })
                    mesh_id = len(self.gltf['meshes']) - 1
                    self.gltf['nodes'].append({
                        'mesh': mesh_id,
                        'translation': [x, y, z],
                    })

    def add_labels(self):
        # Create the metal/non-metal and rough/smooth labels.
        # This data was dumped out of MetalRoughSpheres.gltf.
        verts = [
            4.607307, -12.359449, 0,
            10.713691, -11.081038, 0,
            10.713691, -12.359448, 0,
            -11.874729, -12.359450, 0,
            -3.906647, -11.081039, 0,
            -3.906647, -12.359449, 0,
            10.982860, 3.929994, 0,
            12.186588, 10.969941, 0,
            12.186589, 3.929994, 0,
            10.713691, -12.359448, 8,
            4.607307, -11.081039, 8,
            4.607307, -12.359449, 8,
            12.186590, -11.753379, 0,
            10.982860, -3.775491, 0,
            12.186589, -3.775488, 0,
            -3.906648, -12.359449, 8,
            -11.874729, -11.081040, 8,
            -11.874729, -12.359450, 8,
            12.186589, 3.929994, 8,
            10.982859, 10.969941, 8,
            10.982860, 3.929994, 8,
            10.982861, -11.753379, 8,
            12.186589, -3.775491, 8,
            10.982860, -3.775488, 8,
            4.607307, -11.081039, 0,
            -11.874729, -11.081040, 0,
            10.982859, 10.969941, 0,
            10.713691, -11.081038, 8,
            10.982861, -11.753379, 0,
            -3.906648, -11.081039, 8,
            12.186588, 10.969941, 8,
            12.186590, -11.753379, 8,
        ]
        uvs = [
            0.1175, 0.9988,
            0.0205, 0.9785,
            0.0205, 0.9988,
            0.9870, 0.9988,
            0.8604, 0.9785,
            0.8604, 0.9988,
            0.0205, 0.1284,
            0.0014, 0.0166,
            0.0014, 0.1284,
            0.1175, 0.9988,
            0.0205, 0.9785,
            0.0205, 0.9988,
            0.0014, 0.9988,
            0.0205, 0.8721,
            0.0014, 0.8721,
            0.9870, 0.9988,
            0.8604, 0.9785,
            0.8604, 0.9988,
            0.0205, 0.1284,
            0.0014, 0.0166,
            0.0014, 0.1284,
            0.0014, 0.9988,
            0.0205, 0.8721,
            0.0014, 0.8721,
            0.1175, 0.9785,
            0.9870, 0.9785,
            0.0205, 0.0166,
            0.1175, 0.9785,
            0.0205, 0.9988,
            0.9870, 0.9785,
            0.0205, 0.0166,
            0.0205, 0.9988,
        ]
        indices = [
            0, 1, 2,
            3, 4, 5,
            6, 7, 8,
            9, 10, 11,
            12, 13, 14,
            15, 16, 17,
            18, 19, 20,
            21, 22, 23,
            0, 24, 1,
            3, 25, 4,
            6, 26, 7,
            9, 27, 10,
            12, 28, 13,
            15, 29, 16,
            18, 30, 19,
            21, 31, 22,
        ]

        assert(len(indices) < 2**8)

        vert_data = struct.pack('<%sf' % len(verts), *verts)
        uv_data = struct.pack('<%sf' % len(uvs), *uvs)
        index_data = struct.pack('<%sB' % len(indices), *indices)
        vert_accessor = self.add_accessor(
            vert_data,
            component_type=5126,
            type='VEC3',
            count=len(verts)/3,
            min=[-11.874729, -12.35945, 0],
            max=[12.186589, 10.969941, 8],
            align=4,
        )
        uv_accessor = self.add_accessor(
            uv_data,
            component_type=5126,
            type='VEC2',
            count=len(uvs)/2,
            align=4,
        )
        index_accessor = self.add_accessor(
            index_data,
            component_type=5121,
            type='SCALAR',
            count=len(index_data),
            align=1,
        )

        self.gltf['textures'] = [
            {
                'sampler': 0,
                'source': 0,
            },
        ]
        self.gltf['images'] = [
            { 'uri': 'Spheres_BaseColor.png' }
        ]
        self.gltf['samplers'] = [
            {
                'magFilter': 9729,
                'minFilter': 9986,
                'wrapS': 33071,
                'wrapT': 33071
            }
        ]

        self.gltf['materials'].append({
            'name': 'Labels Material',
            'pbrMetallicRoughness': {
                'baseColorTexture': { 'index': 0 },
                'roughnessFactor': 0.5,
                'metallicFactor': 0,
            }
        })
        material_id = len(self.gltf['materials']) - 1
        self.gltf['meshes'].append({
            'name': 'Labels',
            'primitives': [{
                'attributes': {
                    'POSITION': vert_accessor,
                    'TEXCOORD_0': uv_accessor,
                },
                'indices': index_accessor,
                'material': material_id,
            }]
        })
        mesh_id = len(self.gltf['meshes']) - 1
        self.gltf['nodes'].append({
            'name': 'Labels',
            'mesh': mesh_id,
        })
        node_id = len(self.gltf['nodes']) - 1
        self.gltf['nodes'][0]['children'].append(node_id)

    def __init__(self):
        self.gltf = {
            'asset': { 'version' : '2.0' },
            'nodes': [],
            'meshes': [],
            'materials': [],
            'accessors': [],
            'bufferViews': [],
            'buffers': [],
        }
        self.buf = b''

        self.add_spheres()
        self.add_labels()


mrs = MetalRoughSpheres()
mrs.finish('MetalRoughFactorSpheres.gltf', 'MetalRoughFactorSpheres.bin')