drlukeparry/gpu-3d-printer-slicer-pyslm.py

## gpu-3d-printer-slicer-pyslm.py
"""
Copyright (c) 2023 Luke Parry <ukeparry.uk>. All rights reserved. l

This work is licensed under the terms of the MIT license.
For a copy, see <https://opensource.org/licenses/MIT>.
"""

from matplotlib import pyplot as plt
import numpy as np

import trimesh

from vispy import app, gloo
from vispy.util.transforms import translate, rotate, ortho
from vispy.gloo.util import _screenshot

import pyslm

# Note the backend is sometimes required to ensure Vispy correctly renders using the QT backend.
#app.use_app('pyqt5')  # Set backend

vert = """
uniform   mat4 u_model; // Model transform matrix
uniform   mat4 u_view;
uniform   mat4 u_projection;

uniform  vec2 bounds; // Z bounds
uniform  float frac;  // Z fraction (0 to 1)
uniform  float aspect; // Aspect ratio

attribute vec3 a_position;

#define EPSILON 0.001

void main() {

    vec3 pos = a_position;

    // Ensure the bottom of the part is positioned to z=0 using the bottom bounding box
    pos.z -= bounds[0];

    // Scale the so that it fits within the clipping range (-1.0 < z < 1.0)
    pos.z *= -2.0/(bounds[1]-bounds[0]);

    // Adjust
    pos.z -= frac;
    gl_Position = u_projection * u_view * u_model * vec4(pos, 1.0);
    gl_Position.z += 1.0 - EPSILON;

}
"""

frag = """
void main()
{
    gl_FragColor = vec4(1);
}
"""

# This is the re-implemented Vispy Canvas used for the render infrastructure
class Canvas(app.Canvas):

    @property
    def resolution(self):
        return self._resolution

    @resolution.setter
    def resolution(self, res: float):
        self._resolution = res

    @property
    def visSize(self):
        return self._visSize

    @property
    def mesh(self) -> trimesh.Trimesh:
        return self._mesh

    @property
    def bbox(self) -> np.ndarray:
        return self._bbox

    def setMesh(self, mesh: trimesh.Trimesh):

        self.vertices = np.copy(mesh.vertices).astype(np.float32)
        self.filled = np.copy(mesh.faces).astype(np.int32)
        self.verticesColor = np.copy(mesh.vertices).astype(np.float32)

        self._bbox = mesh.bounds
        self._mesh = mesh

    def __enter__(self):

        self._backend._vispy_warmup()
        return self

    def __init__(self, mesh: trimesh.Trimesh = None,
                 rasterResolution: float = 0.05,
                 flipDir = False, bbox = None):

        self.flipDir = flipDir
        self.rgb = None
        self.vertices = None
        self.filled = None
        self.verticesColor = None
        self._resolution = rasterResolution
        self._mesh = None
        self._bbox = None

        self._z = 0.0

        if mesh:
            self.setMesh(mesh)

        if bbox is not None:
            self._bbox = bbox

        meshExtents = np.diff(self.bbox, axis=0)

        self._visSize = (meshExtents / self.resolution).flatten()

        app.Canvas.__init__(self, 'interactive', show=False, resizable=True, autoswap=False, decorate=False,
                            size=(self.visSize[0], self.visSize[1]))

        print('size', self._visSize)
        print('dpi', self.dpi)

        self.filled = self.filled.astype(np.uint32).flatten()
        self.filled_buf = gloo.IndexBuffer(self.filled)

        vertex_data = np.zeros(self.vertices.shape[0], dtype=[('a_position', np.float32, 3),
                                                              ('a_color', np.float32, 3)])

        vertex_data['a_position'] = self.vertices.astype(np.float32)
        vertex_data['a_color'] = self.vertices.astype(np.float32)

        self.program = gloo.Program(vert, frag)
        self.program.bind(gloo.VertexBuffer(vertex_data))

        avg = np.mean(self.bbox, axis=0)

        if flipDir:
            self.view = rotate(0, [1, 0, 0])
        else:
            self.view = np.dot(np.dot(translate((-avg[0], -avg[1], -avg[2])),
                              rotate(-180, [1,0,0])),
                              translate((avg[0], avg[1], avg[2])))


        self.model = np.eye(4, dtype=np.float32)
        shape = int(self._visSize[1]), int(self._visSize[0])

        # Create the render texture
        self._rendertex = gloo.Texture2D((shape + (4,)), format='rgba', internalformat='rgba32f')
        # self._colorBuffer = gloo.RenderBuffer(self.shape, format='color')

        self._stencilRenderBuffer = gloo.RenderBuffer(shape, format='stencil')
        self._stencilRenderBuffer.resize(shape, format=gloo.gl.GL_STENCIL_INDEX8)

        #self._depthRenderBuffer = gloo.RenderBuffer(shape, format='depth')
        #self._depthRenderBuffer.resize(shape, format=gloo.gl.GL_DEPTH_COMPONENT16)

        # Create FBO, attach the color buffer and depth buffer
        self._fbo = gloo.FrameBuffer(self._rendertex)#, self._depthRenderBuffer)
        self._fbo.stencil_buffer=self._stencilRenderBuffer
        gloo.set_viewport(0, 0, self.physical_size[0], self.physical_size[1])
        gloo.set_viewport(0, 0, self._visSize[0], self._visSize[1])

        self.projection = ortho(self.bbox[1, 0], self.bbox[0, 0], self.bbox[1, 1], self.bbox[0, 1], 2, 40)
        self.model = np.eye(4, dtype=np.float32)

        # Set MVP variables for shaders
        self.program['u_projection'] = self.projection
        self.program['u_model'] = self.model
        self.program['u_view'] = self.view

        gloo.set_clear_color('white')
        gloo.set_state('opaque')

        self.update()

    def on_resize(self, event):

        gloo.set_viewport(0, 0, self._visSize[0], self._visSize[1])
        self.finalSize = (event.physical_size[0], event.physical_size[1])
        #print('event physical size', event.physical_size[0], event.physical_size[1])

        self.projection = ortho(self.bbox[0, 0], self.bbox[1, 0],
                                self.bbox[0, 1], self.bbox[1, 1],
                                1.0,-1.0)
                                #-0.5*(self.bbox[1, 2] - self.bbox[0, 2]),0.5*(self.bbox[1, 2] - self.bbox[0, 2]))

        self.program['u_projection'] = self.projection

    @property
    def slicePosition(self):
        return self._z

    @slicePosition.setter
    def slicePosition(self, value):
        self._z = value

    def on_draw(self, event):

        with self._fbo:
            # Set of the corret GL State
            gloo.set_state(blend=False, stencil_test=True, depth_test=False, polygon_offset_fill=False, cull_face=False)

            gloo.set_viewport(0, 0, self._visSize[0], self._visSize[1])
            gloo.set_clear_stencil(0)
            gloo.set_clear_color((0.0, 0.0, 0.0, 0.0))
            gloo.clear()

            self.program['bounds'] = self.bbox[0,2], self.bbox[1,2]
            self.program['aspect'] = self.physical_size[1] /  self.physical_size[0]

            # The position of the slice position passed to the GLSL shader
            self.program['frac'] = self._z * 2.0 # - (self.bbox[1,2]-self.bbox[0,2])

            # Draw twice, adding and subtracting values in the stencil buffer

            # Render Pass 1 (Increment Stencil Buffers)
            gloo.set_stencil_func('always', 0, 0xff)
            gloo.set_stencil_op('keep', 'keep', 'incr', 'back')
            gloo.set_stencil_op('keep', 'keep', 'keep', 'front')
            self.program.draw('triangles', self.filled_buf)

            # Render Pass 2 (Decrease Stencil Buffers)
            gloo.set_stencil_op('keep', 'keep', 'decr', 'front')
            gloo.set_stencil_op('keep', 'keep', 'keep', 'back')
            self.program.draw('triangles', self.filled_buf)

            # Clear only the color buffer
            gloo.clear(color=True, depth=False, stencil=False)

            # Render Pass 3
            gloo.set_stencil_func('notequal', 0, 0xff)
            gloo.set_stencil_op('keep', 'keep', 'keep')

            # Store the final framebuffer
            self.program.draw('triangles', self.filled_buf)

            self.rgb = _screenshot((0, 0, self._visSize[0], self._visSize[1]))
            #print('physical size', self.physical_size)
            #print('image size', self.rgb.shape)


def projectHeightMap(mesh: trimesh.Trimesh,
                     resolution: float = 0.05,
                     flipDir: bool = False,
                     bbox: np.ndarray = None):

    c = Canvas(mesh, resolution, flipDir, bbox)

    c.show(visible=True)
    c.render()
    c.close()
    return c.rgb[:, :, 1]

# Imports the part and sets the geometry to  an STL file using the PySLM infrastructre
solidPart = pyslm.Part('aerofoilk')
solidPart.setGeometry('../../models/aerofoil.stl')

solidPart.origin[0] = 5.0
solidPart.origin[1] = 2.5
solidPart.scaleFactor = 1.0
solidPart.rotation = [0, 0.0, 0]
solidPart.dropToPlatform(5.0)

plt.show()
resolution = 0.02 # 20 microns per pixel
c = Canvas(solidPart.geometry, resolution, True, solidPart.geometry.bounds)

# Slice the whole part using a 10 micron layer resolution in the first 1 mm
slices =[]

for z in np.arange(0, 1.01, 0.01):
    c.slicePosition = 0.01
    c.slicePosition = z
    c.show(visible=True)
    c.render()
    slices.append(c.rgb[:,:,0])

# Close the Vispy Canvas Window
c.close()

# Show the Image
plt.imshow(slices[100])
	"""
	Copyright (c) 2023 Luke Parry <ukeparry.uk>. All rights reserved. l

	This work is licensed under the terms of the MIT license.
	For a copy, see <https://opensource.org/licenses/MIT>.
	"""

	from matplotlib import pyplot as plt
	import numpy as np

	import trimesh

	from vispy import app, gloo
	from vispy.util.transforms import translate, rotate, ortho
	from vispy.gloo.util import _screenshot

	import pyslm

	# Note the backend is sometimes required to ensure Vispy correctly renders using the QT backend.
	#app.use_app('pyqt5') # Set backend

	vert = """
	uniform mat4 u_model; // Model transform matrix
	uniform mat4 u_view;
	uniform mat4 u_projection;

	uniform vec2 bounds; // Z bounds
	uniform float frac; // Z fraction (0 to 1)
	uniform float aspect; // Aspect ratio

	attribute vec3 a_position;

	#define EPSILON 0.001

	void main() {

	vec3 pos = a_position;

	// Ensure the bottom of the part is positioned to z=0 using the bottom bounding box
	pos.z -= bounds[0];

	// Scale the so that it fits within the clipping range (-1.0 < z < 1.0)
	pos.z *= -2.0/(bounds[1]-bounds[0]);

	// Adjust
	pos.z -= frac;
	gl_Position = u_projection * u_view * u_model * vec4(pos, 1.0);
	gl_Position.z += 1.0 - EPSILON;

	}
	"""

	frag = """
	void main()
	{
	gl_FragColor = vec4(1);
	}
	"""

	# This is the re-implemented Vispy Canvas used for the render infrastructure
	class Canvas(app.Canvas):

	@property
	def resolution(self):
	return self._resolution

	@resolution.setter
	def resolution(self, res: float):
	self._resolution = res

	@property
	def visSize(self):
	return self._visSize

	@property
	def mesh(self) -> trimesh.Trimesh:
	return self._mesh

	@property
	def bbox(self) -> np.ndarray:
	return self._bbox

	def setMesh(self, mesh: trimesh.Trimesh):

	self.vertices = np.copy(mesh.vertices).astype(np.float32)
	self.filled = np.copy(mesh.faces).astype(np.int32)
	self.verticesColor = np.copy(mesh.vertices).astype(np.float32)

	self._bbox = mesh.bounds
	self._mesh = mesh

	def __enter__(self):

	self._backend._vispy_warmup()
	return self

	def __init__(self, mesh: trimesh.Trimesh = None,
	rasterResolution: float = 0.05,
	flipDir = False, bbox = None):

	self.flipDir = flipDir
	self.rgb = None
	self.vertices = None
	self.filled = None
	self.verticesColor = None
	self._resolution = rasterResolution
	self._mesh = None
	self._bbox = None

	self._z = 0.0

	if mesh:
	self.setMesh(mesh)

	if bbox is not None:
	self._bbox = bbox

	meshExtents = np.diff(self.bbox, axis=0)

	self._visSize = (meshExtents / self.resolution).flatten()

	app.Canvas.__init__(self, 'interactive', show=False, resizable=True, autoswap=False, decorate=False,
	size=(self.visSize[0], self.visSize[1]))

	print('size', self._visSize)
	print('dpi', self.dpi)

	self.filled = self.filled.astype(np.uint32).flatten()
	self.filled_buf = gloo.IndexBuffer(self.filled)

	vertex_data = np.zeros(self.vertices.shape[0], dtype=[('a_position', np.float32, 3),
	('a_color', np.float32, 3)])

	vertex_data['a_position'] = self.vertices.astype(np.float32)
	vertex_data['a_color'] = self.vertices.astype(np.float32)

	self.program = gloo.Program(vert, frag)
	self.program.bind(gloo.VertexBuffer(vertex_data))

	avg = np.mean(self.bbox, axis=0)

	if flipDir:
	self.view = rotate(0, [1, 0, 0])
	else:
	self.view = np.dot(np.dot(translate((-avg[0], -avg[1], -avg[2])),
	rotate(-180, [1,0,0])),
	translate((avg[0], avg[1], avg[2])))


	self.model = np.eye(4, dtype=np.float32)
	shape = int(self._visSize[1]), int(self._visSize[0])

	# Create the render texture
	self._rendertex = gloo.Texture2D((shape + (4,)), format='rgba', internalformat='rgba32f')
	# self._colorBuffer = gloo.RenderBuffer(self.shape, format='color')

	self._stencilRenderBuffer = gloo.RenderBuffer(shape, format='stencil')
	self._stencilRenderBuffer.resize(shape, format=gloo.gl.GL_STENCIL_INDEX8)

	#self._depthRenderBuffer = gloo.RenderBuffer(shape, format='depth')
	#self._depthRenderBuffer.resize(shape, format=gloo.gl.GL_DEPTH_COMPONENT16)

	# Create FBO, attach the color buffer and depth buffer
	self._fbo = gloo.FrameBuffer(self._rendertex)#, self._depthRenderBuffer)
	self._fbo.stencil_buffer=self._stencilRenderBuffer
	gloo.set_viewport(0, 0, self.physical_size[0], self.physical_size[1])
	gloo.set_viewport(0, 0, self._visSize[0], self._visSize[1])

	self.projection = ortho(self.bbox[1, 0], self.bbox[0, 0], self.bbox[1, 1], self.bbox[0, 1], 2, 40)
	self.model = np.eye(4, dtype=np.float32)

	# Set MVP variables for shaders
	self.program['u_projection'] = self.projection
	self.program['u_model'] = self.model
	self.program['u_view'] = self.view

	gloo.set_clear_color('white')
	gloo.set_state('opaque')

	self.update()

	def on_resize(self, event):

	gloo.set_viewport(0, 0, self._visSize[0], self._visSize[1])
	self.finalSize = (event.physical_size[0], event.physical_size[1])
	#print('event physical size', event.physical_size[0], event.physical_size[1])

	self.projection = ortho(self.bbox[0, 0], self.bbox[1, 0],
	self.bbox[0, 1], self.bbox[1, 1],
	1.0,-1.0)
	#-0.5(self.bbox[1, 2] - self.bbox[0, 2]),0.5(self.bbox[1, 2] - self.bbox[0, 2]))

	self.program['u_projection'] = self.projection

	@property
	def slicePosition(self):
	return self._z

	@slicePosition.setter
	def slicePosition(self, value):
	self._z = value

	def on_draw(self, event):

	with self._fbo:
	# Set of the corret GL State
	gloo.set_state(blend=False, stencil_test=True, depth_test=False, polygon_offset_fill=False, cull_face=False)

	gloo.set_viewport(0, 0, self._visSize[0], self._visSize[1])
	gloo.set_clear_stencil(0)
	gloo.set_clear_color((0.0, 0.0, 0.0, 0.0))
	gloo.clear()

	self.program['bounds'] = self.bbox[0,2], self.bbox[1,2]
	self.program['aspect'] = self.physical_size[1] / self.physical_size[0]

	# The position of the slice position passed to the GLSL shader
	self.program['frac'] = self._z * 2.0 # - (self.bbox[1,2]-self.bbox[0,2])

	# Draw twice, adding and subtracting values in the stencil buffer

	# Render Pass 1 (Increment Stencil Buffers)
	gloo.set_stencil_func('always', 0, 0xff)
	gloo.set_stencil_op('keep', 'keep', 'incr', 'back')
	gloo.set_stencil_op('keep', 'keep', 'keep', 'front')
	self.program.draw('triangles', self.filled_buf)

	# Render Pass 2 (Decrease Stencil Buffers)
	gloo.set_stencil_op('keep', 'keep', 'decr', 'front')
	gloo.set_stencil_op('keep', 'keep', 'keep', 'back')
	self.program.draw('triangles', self.filled_buf)

	# Clear only the color buffer
	gloo.clear(color=True, depth=False, stencil=False)

	# Render Pass 3
	gloo.set_stencil_func('notequal', 0, 0xff)
	gloo.set_stencil_op('keep', 'keep', 'keep')

	# Store the final framebuffer
	self.program.draw('triangles', self.filled_buf)

	self.rgb = _screenshot((0, 0, self._visSize[0], self._visSize[1]))
	#print('physical size', self.physical_size)
	#print('image size', self.rgb.shape)


	def projectHeightMap(mesh: trimesh.Trimesh,
	resolution: float = 0.05,
	flipDir: bool = False,
	bbox: np.ndarray = None):

	c = Canvas(mesh, resolution, flipDir, bbox)

	c.show(visible=True)
	c.render()
	c.close()
	return c.rgb[:, :, 1]

	# Imports the part and sets the geometry to an STL file using the PySLM infrastructre
	solidPart = pyslm.Part('aerofoilk')
	solidPart.setGeometry('../../models/aerofoil.stl')

	solidPart.origin[0] = 5.0
	solidPart.origin[1] = 2.5
	solidPart.scaleFactor = 1.0
	solidPart.rotation = [0, 0.0, 0]
	solidPart.dropToPlatform(5.0)

	plt.show()
	resolution = 0.02 # 20 microns per pixel
	c = Canvas(solidPart.geometry, resolution, True, solidPart.geometry.bounds)

	# Slice the whole part using a 10 micron layer resolution in the first 1 mm
	slices =[]

	for z in np.arange(0, 1.01, 0.01):
	c.slicePosition = 0.01
	c.slicePosition = z
	c.show(visible=True)
	c.render()
	slices.append(c.rgb[:,:,0])

	# Close the Vispy Canvas Window
	c.close()

	# Show the Image
	plt.imshow(slices[100])