This example is a prototype demonstrating normal mapping in VisPy

normal_mapping_example.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
This example is a prototype demonstrating normal mapping
by comparison between a reference mesh and its flatten
version.
"""

import numpy as np
from vispy import app, gloo
from vispy.geometry import create_plane, MeshData
from vispy.util.transforms import perspective, translate

vert = """
uniform mat4 u_model;
uniform mat4 u_view;
uniform mat4 u_projection;
uniform vec3 u_light;
uniform vec4 u_color;
uniform bool u_flatten;

attribute vec3 a_position;
attribute vec3 a_normal;

varying vec3 v_normal_vec;
varying vec3 v_light_vec;
varying vec3 v_eye_vec;

varying vec4 v_ambientk;
varying vec4 v_light_color;
varying vec4 v_base_color;

void main()
{
    v_ambientk = vec4(0.3, 0.3, 0.3, 1.0);
    v_light_color = vec4(1.0, 1.0, 1.0, 1.0);
    vec3 position = a_position;

    if(u_flatten){
        v_base_color = vec4(a_normal, 1.0);
        position.z = 0;
    }
    else
        v_base_color = u_color;

    v_normal_vec = normalize(a_normal);
    v_light_vec = normalize(u_light);
    v_eye_vec = normalize(vec3(0, 0, 5));
    gl_Position = u_projection * u_view * u_model * vec4(position,1.0);
}
"""


frag = """
varying vec3 v_normal_vec;
varying vec3 v_light_vec;
varying vec3 v_eye_vec;

varying vec4 v_ambientk;
varying vec4 v_light_color;
varying vec4 v_base_color;

void main() {
    float shininess = 1. / 200.;

    //DIFFUSE
    float diffusek = dot(v_light_vec, v_normal_vec);
    // clamp, because 0 < theta < pi/2
    diffusek  = clamp(diffusek, 0.0, 1.0);
    vec4 diffuse_color = v_light_color * diffusek;

    //SPECULAR
    //reflect light wrt normal for the reflected ray, then
    //find the angle made with the eye
    float speculark = 0.0;
    if (shininess > 0.) {
        speculark = dot(reflect(v_light_vec, v_normal_vec), v_eye_vec);
        speculark = clamp(speculark, 0.0, 1.0);
        //raise to the material's shininess, multiply with a
        //small factor for spread
        speculark = 20.0 * pow(speculark, 1.0 / shininess);
    }
    vec4 specular_color = v_light_color * speculark;
    gl_FragColor = v_base_color * (v_ambientk + diffuse_color) + specular_color;
}
"""  # noqa


def mesh(n_segments=5, length=3):
    v, indices, outline = create_plane(width=length,
                                       height=length,
                                       width_segments=n_segments,
                                       height_segments=n_segments)
    vertices = v['position']
    vertices[:, 2] = \
        np.exp(-(vertices[:, 0] ** 2 + vertices[:, 1] ** 2) / 0.2) * 2.0

    return vertices, indices, outline


def compute_normals(vertices, faces):
    md = MeshData(vertices=vertices,
                  faces=faces)
    return md.get_vertex_normals()


def write_normal_map(filename, normals, size):
    from vispy.io import write_png
    data0 = (normals[:, 0] + 1) * 127.5
    data1 = (normals[:, 1] + 1) * 127.5
    data2 = np.clip(normals[:, 2], 0.0, 1.0) * 127.0 + 128.0
    data = np.dstack((data0, data1, data2)).astype(np.uint8)
    write_png(filename, np.reshape(data, size))


def load_normal_map(filename):
    from vispy.io import read_png
    img = read_png(filename)
    img = img.astype(np.float32)
    data0 = img[:, :, 0].flatten() / 127.5 - 1.0
    data1 = img[:, :, 1].flatten() / 127.5 - 1.0
    data2 = (img[:, :, 2].flatten() - 128.0) / 127.0
    normal_map = np.array(list(zip(data0, data1, data2)))
    return normal_map


# -----------------------------------------------------------------------------
class Canvas(app.Canvas):

    def __init__(self):
        app.Canvas.__init__(self, keys='interactive', size=(800, 600))

        self.theta = 0.0
        self.light_distance = 16.0
        self.mesh_color = (0.7, 0.7, 0.7, 1.0)
        self.line_color = (0.0, 0.0, 0.0, 1.0)

        self.filled_buf = gloo.IndexBuffer(filled)
        self.outline_buf = gloo.IndexBuffer(outline)

        self.program = gloo.Program(vert, frag)
        self.buffer = np.array(list(zip(vertices, normals)),
                               dtype=[('a_position', np.float32, 3),
                                      ('a_normal', np.float32, 3)])
        self.program.bind(gloo.VertexBuffer(self.buffer))

        self.view = translate((-length/2.0, 0, -5))
        self.model = np.eye(4, dtype=np.float32)
        gloo.set_viewport(0, 0, self.physical_size[0], self.physical_size[1])
        self.projection = perspective(45.0, self.size[0] /
                                      float(self.size[1]), 2.0, 10.0)

        self.program['u_flatten'] = False
        self.program['u_projection'] = self.projection
        self.program['u_model'] = self.model
        self.program['u_view'] = self.view
        self.program['u_light'] = np.cos(self.theta), np.sin(self.theta), 5

        gloo.set_clear_color('black')
        gloo.set_state('opaque')
        gloo.set_polygon_offset(1, 1)

        self._timer = app.Timer('auto', connect=self.on_timer, start=True)

        self.show()

    # ---------------------------------
    def on_timer(self, event):
        self.theta += .1
        self.program['u_light'] = [self.light_distance*np.cos(self.theta),
                                   self.light_distance*np.sin(self.theta),
                                   self.light_distance]
        self.update()

    # ---------------------------------
    def on_resize(self, event):
        gloo.set_viewport(0, 0, event.physical_size[0], event.physical_size[1])
        self.projection = perspective(45.0, event.size[0] /
                                      float(event.size[1]), 2.0, 10.0)
        self.program['u_projection'] = self.projection

    # ---------------------------------
    def on_draw(self, event):
        gloo.clear()

        self.model = np.eye(4, dtype=np.float32)
        self.program['u_model'] = self.model
        self.program['u_flatten'] = False

        # Filled
        gloo.set_state(blend=False, depth_test=True, polygon_offset_fill=True)
        self.program['u_color'] = self.mesh_color
        self.program.draw('triangles', self.filled_buf)
        # Outline
        gloo.set_state(blend=True, depth_test=True, polygon_offset_fill=False)
        gloo.set_depth_mask(False)
        self.program['u_color'] = self.line_color
        self.program.draw('lines', self.outline_buf)
        gloo.set_depth_mask(True)

        self.model = translate((length, 0, 0))
        self.program['u_model'] = self.model
        self.program['u_flatten'] = True

        # Filled
        gloo.set_state(blend=False, depth_test=True, polygon_offset_fill=True)
        self.program.draw('triangles', self.filled_buf)


if __name__ == '__main__':
    # parameters for mesh generation
    length = 2.5
    n_segments = 128

    # create the reference mesh and compute its vertex normals
    vertices, filled, outline = mesh(n_segments=n_segments, length=length)
    normals = compute_normals(vertices, filled)

    # some helpers functions
    map_width = map_height = n_segments + 1
    write_normal_map('output_normal_map.png', normals,
                     (map_width, map_height, 3))
    normals = load_normal_map('output_normal_map.png')

    c = Canvas()
    app.run()