extrude.py

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