import math
import os
import numpy as np
import pygltflib
import trimesh
from PIL import Image, ImageFilter
def quaternion_multiply(q1, q2):
x1, y1, z1, w1 = q1
x2, y2, z2, w2 = q2
return [
w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
]
def glb_add_lights(path_input, path_output):
"""
Adds directional lights in the horizontal plane to the glb file.
:param path_input: path to input glb
:param path_output: path to output glb
:return: None
"""
glb = pygltflib.GLTF2().load(path_input)
N = 3 # default max num lights in Babylon.js is 4
angle_step = 2 * math.pi / N
elevation_angle = math.radians(75)
light_colors = [
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
]
lights_extension = {
"lights": [
{"type": "directional", "color": light_colors[i], "intensity": 2.0}
for i in range(N)
]
}
if "KHR_lights_punctual" not in glb.extensionsUsed:
glb.extensionsUsed.append("KHR_lights_punctual")
glb.extensions["KHR_lights_punctual"] = lights_extension
light_nodes = []
for i in range(N):
angle = i * angle_step
pos_rot = [0.0, 0.0, math.sin(angle / 2), math.cos(angle / 2)]
elev_rot = [
math.sin(elevation_angle / 2),
0.0,
0.0,
math.cos(elevation_angle / 2),
]
rotation = quaternion_multiply(pos_rot, elev_rot)
node = {
"rotation": rotation,
"extensions": {"KHR_lights_punctual": {"light": i}},
}
light_nodes.append(node)
light_node_indices = list(range(len(glb.nodes), len(glb.nodes) + N))
glb.nodes.extend(light_nodes)
root_node_index = glb.scenes[glb.scene].nodes[0]
root_node = glb.nodes[root_node_index]
if hasattr(root_node, "children"):
root_node.children.extend(light_node_indices)
else:
root_node.children = light_node_indices
glb.save(path_output)
def extrude_depth_3d(
path_rgb,
path_depth,
output_model_scale=100,
filter_size=3,
coef_near=0.0,
coef_far=1.0,
emboss=0.3,
f_thic=0.05,
f_near=-0.15,
f_back=0.01,
vertex_colors=True,
scene_lights=True,
prepare_for_3d_printing=False,
):
f_far_inner = -emboss
f_far_outer = f_far_inner - f_back
f_near = max(f_near, f_far_inner)
depth_image = Image.open(path_depth)
assert depth_image.mode == "I", depth_image.mode
depth_image = depth_image.filter(ImageFilter.MedianFilter(size=filter_size))
w, h = depth_image.size
d_max = max(w, h)
depth_image = np.array(depth_image).astype(np.double)
z_min, z_max = np.min(depth_image), np.max(depth_image)
depth_image = (depth_image.astype(np.double) - z_min) / (z_max - z_min)
depth_image[depth_image < coef_near] = coef_near
depth_image[depth_image > coef_far] = coef_far
depth_image = emboss * (depth_image - coef_near) / (coef_far - coef_near)
rgb_image = np.array(
Image.open(path_rgb).convert("RGB").resize((w, h), Image.Resampling.LANCZOS)
)
w_norm = w / float(d_max - 1)
h_norm = h / float(d_max - 1)
w_half = w_norm / 2
h_half = h_norm / 2
x, y = np.meshgrid(np.arange(w), np.arange(h))
x = x / float(d_max - 1) - w_half # [-w_half, w_half]
y = -y / float(d_max - 1) + h_half # [-h_half, h_half]
z = -depth_image # -depth_emboss (far) - 0 (near)
vertices_2d = np.stack((x, y, z), axis=-1)
vertices = vertices_2d.reshape(-1, 3)
colors = rgb_image[:, :, :3].reshape(-1, 3) / 255.0
faces = []
for y in range(h - 1):
for x in range(w - 1):
idx = y * w + x
faces.append([idx, idx + w, idx + 1])
faces.append([idx + 1, idx + w, idx + 1 + w])
# OUTER frame
nv = len(vertices)
vertices = np.append(
vertices,
[
[-w_half - f_thic, -h_half - f_thic, f_near], # 00
[-w_half - f_thic, -h_half - f_thic, f_far_outer], # 01
[w_half + f_thic, -h_half - f_thic, f_near], # 02
[w_half + f_thic, -h_half - f_thic, f_far_outer], # 03
[w_half + f_thic, h_half + f_thic, f_near], # 04
[w_half + f_thic, h_half + f_thic, f_far_outer], # 05
[-w_half - f_thic, h_half + f_thic, f_near], # 06
[-w_half - f_thic, h_half + f_thic, f_far_outer], # 07
],
axis=0,
)
faces.extend(
[
[nv + 0, nv + 1, nv + 2],
[nv + 2, nv + 1, nv + 3],
[nv + 2, nv + 3, nv + 4],
[nv + 4, nv + 3, nv + 5],
[nv + 4, nv + 5, nv + 6],
[nv + 6, nv + 5, nv + 7],
[nv + 6, nv + 7, nv + 0],
[nv + 0, nv + 7, nv + 1],
]
)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * 8, axis=0)
# INNER frame
nv = len(vertices)
vertices_left_data = vertices_2d[:, 0] # H x 3
vertices_left_frame = vertices_2d[:, 0].copy() # H x 3
vertices_left_frame[:, 2] = f_near
vertices = np.append(vertices, vertices_left_data, axis=0)
vertices = np.append(vertices, vertices_left_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * h), axis=0)
for i in range(h - 1):
nvi_d = nv + i
nvi_f = nvi_d + h
faces.append([nvi_d, nvi_f, nvi_d + 1])
faces.append([nvi_d + 1, nvi_f, nvi_f + 1])
nv = len(vertices)
vertices_right_data = vertices_2d[:, -1] # H x 3
vertices_right_frame = vertices_2d[:, -1].copy() # H x 3
vertices_right_frame[:, 2] = f_near
vertices = np.append(vertices, vertices_right_data, axis=0)
vertices = np.append(vertices, vertices_right_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * h), axis=0)
for i in range(h - 1):
nvi_d = nv + i
nvi_f = nvi_d + h
faces.append([nvi_d, nvi_d + 1, nvi_f])
faces.append([nvi_d + 1, nvi_f + 1, nvi_f])
nv = len(vertices)
vertices_top_data = vertices_2d[0, :] # H x 3
vertices_top_frame = vertices_2d[0, :].copy() # H x 3
vertices_top_frame[:, 2] = f_near
vertices = np.append(vertices, vertices_top_data, axis=0)
vertices = np.append(vertices, vertices_top_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * w), axis=0)
for i in range(w - 1):
nvi_d = nv + i
nvi_f = nvi_d + w
faces.append([nvi_d, nvi_d + 1, nvi_f])
faces.append([nvi_d + 1, nvi_f + 1, nvi_f])
nv = len(vertices)
vertices_bottom_data = vertices_2d[-1, :] # H x 3
vertices_bottom_frame = vertices_2d[-1, :].copy() # H x 3
vertices_bottom_frame[:, 2] = f_near
vertices = np.append(vertices, vertices_bottom_data, axis=0)
vertices = np.append(vertices, vertices_bottom_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 * w), axis=0)
for i in range(w - 1):
nvi_d = nv + i
nvi_f = nvi_d + w
faces.append([nvi_d, nvi_f, nvi_d + 1])
faces.append([nvi_d + 1, nvi_f, nvi_f + 1])
# FRONT frame
nv = len(vertices)
vertices = np.append(
vertices,
[
[-w_half - f_thic, -h_half - f_thic, f_near],
[-w_half - f_thic, h_half + f_thic, f_near],
],
axis=0,
)
vertices = np.append(vertices, vertices_left_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + h), axis=0)
for i in range(h - 1):
faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
faces.append([nv, nv + 2, nv + 1])
nv = len(vertices)
vertices = np.append(
vertices,
[
[w_half + f_thic, h_half + f_thic, f_near],
[w_half + f_thic, -h_half - f_thic, f_near],
],
axis=0,
)
vertices = np.append(vertices, vertices_right_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + h), axis=0)
for i in range(h - 1):
faces.append([nv, nv + 2 + i, nv + 2 + i + 1])
faces.append([nv, nv + h + 1, nv + 1])
nv = len(vertices)
vertices = np.append(
vertices,
[
[w_half + f_thic, h_half + f_thic, f_near],
[-w_half - f_thic, h_half + f_thic, f_near],
],
axis=0,
)
vertices = np.append(vertices, vertices_top_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + w), axis=0)
for i in range(w - 1):
faces.append([nv, nv + 2 + i, nv + 2 + i + 1])
faces.append([nv, nv + 1, nv + 2])
nv = len(vertices)
vertices = np.append(
vertices,
[
[-w_half - f_thic, -h_half - f_thic, f_near],
[w_half + f_thic, -h_half - f_thic, f_near],
],
axis=0,
)
vertices = np.append(vertices, vertices_bottom_frame, axis=0)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * (2 + w), axis=0)
for i in range(w - 1):
faces.append([nv, nv + 2 + i + 1, nv + 2 + i])
faces.append([nv, nv + 1, nv + w + 1])
# BACK frame
nv = len(vertices)
vertices = np.append(
vertices,
[
[-w_half - f_thic, -h_half - f_thic, f_far_outer], # 00
[w_half + f_thic, -h_half - f_thic, f_far_outer], # 01
[w_half + f_thic, h_half + f_thic, f_far_outer], # 02
[-w_half - f_thic, h_half + f_thic, f_far_outer], # 03
],
axis=0,
)
faces.extend(
[
[nv + 0, nv + 2, nv + 1],
[nv + 2, nv + 0, nv + 3],
]
)
colors = np.append(colors, [[0.5, 0.5, 0.5]] * 4, axis=0)
trimesh_kwargs = {}
if vertex_colors:
trimesh_kwargs["vertex_colors"] = colors
mesh = trimesh.Trimesh(vertices=vertices, faces=faces, **trimesh_kwargs)
mesh.merge_vertices()
current_max_dimension = max(mesh.extents)
scaling_factor = output_model_scale / current_max_dimension
mesh.apply_scale(scaling_factor)
if prepare_for_3d_printing:
rotation_mat = trimesh.transformations.rotation_matrix(np.radians(90), [-1, 0, 0])
mesh.apply_transform(rotation_mat)
path_out_base = os.path.splitext(path_depth)[0].replace("_16bit", "")
path_out_glb = path_out_base + ".glb"
path_out_stl = path_out_base + ".stl"
mesh.export(path_out_glb, file_type="glb")
if scene_lights:
glb_add_lights(path_out_glb, path_out_glb)
mesh.export(path_out_stl, file_type="stl")
return path_out_glb, path_out_stl