Created
January 8, 2023 21:54
-
-
Save drlukeparry/f7124f859577f093c857059f9c887d2b to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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]) | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment