蛇形通道网格沿骨架离散点尺寸提取与标量映射优化求助

我有一个带渐变蛇形通道的表面网格（STL文件）：

需要完成以下操作：

• 获取沿通道骨架（长度方向）离散点处的通道尺寸（宽度、高度）
• 根据宽度和高度计算标量值，公式为 calc_scalar(width, height) = 1/width + 1/height
• 将标量值映射回体素表示或切片中

当前实现存在的问题：
我使用PyVista、Scikit-Image形态学模块、Scipy ndimage和Matplotlib完成了部分工作，但仅提取了宽度（未获取高度），且将标量映射到2D切片的过程速度极慢、效果不完善。相关代码如下：

#!/usr/bin/env python
# coding: utf-8

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

import pyvista as pv

import skimage as ski
from skimage.morphology import skeletonize, medial_axis

from scipy.ndimage import distance_transform_edt
import copy

# The function to calculate the scalar from the channel width and height
def calc_scalar(width, height):
    return 1/width + 1/height

# Some constants
mesh_res = 2400
figsize_mult = 15
image_type = 'png'
savefig_dpi = 1200
thresh = .95
thresh_dist = 9
atol = 16

# Colors maps
cmap_greys = cmap = plt.get_cmap("Greys", 2)
cmap_coolwarm = copy.copy(mpl.cm.coolwarm)
cmap_coolwarm.set_bad(color='black', alpha=1.)

# Import the STL mesh
mesh = pv.read("serpent.stl")
mesh_length = mesh.bounds[1] - mesh.bounds[0]
mesh_width = mesh.bounds[3] - mesh.bounds[2]
print(f"Mesh's bounding box dimensions:\n  Length x Width: {mesh_length:.3f} mm x {mesh_width:.3f} mm")

# Make voxels (volume) out of the surface mesh
voxels = pv.voxelize(mesh, density=mesh.length/mesh_res)
voxels['dummy'] = np.zeros(voxels.GetNumberOfCells())

# Slice the voxels
data = voxels.ctp().slice('z', generate_triangles=True)
tri = data.faces.reshape((-1,4))[:,1:]
u = data.active_scalars

# Plot the slice and save as a raster image
fig, ax = plt.subplots(figsize=(figsize_mult*mesh_length/25.4, figsize_mult*mesh_width/25.4))
ax.tricontourf(data.points[:,0], data.points[:,1], tri, u, cmap=cmap_greys, vmin=0, vmax=1)
ax = plt.gca()
ax.set_aspect('equal')
ax.set_facecolor('k')
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.set_ylim(-.55*mesh_width, .55*mesh_width)
plt.tight_layout()
plt.savefig(f'serpent bw.{image_type}', dpi=savefig_dpi, bbox_inches='tight', facecolor='black')

# Load back the raster image using scikit.image
# Do some clean-up (Gaussian filter, threshold)
im = ski.io.imread(f'serpent bw.{image_type}', as_gray=True)
blurred_im = ski.filters.gaussian(im, sigma=1.0)
binary_mask = blurred_im > thresh

# Get the skeletons and the distance
# Somehow, the different functions yield somewhat different skeletons
skel, distance = medial_axis(binary_mask, return_distance=True)
skeleton = skeletonize(binary_mask)
dist_on_skel = distance * skel

# Yet another, cleaner, skeleton from Scipy ndimage
im_edt, indices = distance_transform_edt(im, return_indices=True)
im_edt_skel = im_edt * skeleton
sc_im_edt = np.where(skeleton > 0, calc_scalar(im_edt, 1), 0)

# Scan the array searching for the skeleton.
# Compute the scalar value at each skeleton point.
# Fill the channel's width with the same value.
# Super slow...
sc_im_edt_cp = np.full_like(sc_im_edt, np.nan)
for i, r in enumerate(sc_im_edt):
    for j, p in enumerate(r):
        if ~np.isnan(p) and p >=.1:
            y = indices[0][i][j]
            x = indices[1][i][j]
            sc_im_edt_cp[i][j] = p
            if x == j:
                # print(f"(j, i): ({j}, {i}); (x, y): ({x}, {y})")
                for y1 in np.arange(np.abs(y - i)):
                    # print(f"(x1, y1) = ({x}, {y + y1})")
                    try:
                        sc_im_edt_cp[i + y1][j] = p
                    except IndexError:
                        pass
                    try:
                        sc_im_edt_cp[i - y1][j] = p
                    except IndexError:
                        pass
            elif y == i:
                # print(f"(j, i): ({j}, {i}); (x, y): ({x}, {y})")
                for x1 in np.arange(np.abs(x - j)):
                    # print(f"(x1, y1) = ({x + x1}, {y})")
                    try:
                        sc_im_edt_cp[i][j + x1] = p
                    except IndexError:
                        pass
                    try:
                        sc_im_edt_cp[i][j - x1] = p
                    except IndexError:
                        pass
            else:
                m = np.nan
                ## in which cadran the boundary is located
                if x < j:
                    m = (i - y) / (j - x)
                elif x > j:
                    m = (y - i) / (x - j)
                b = y - m * x
                # print(f"(j, i): ({j}, {i}); (x, y): ({x}, {y}), (m, b): ({m}, {b})")
                for y1 in np.arange(np.abs(y - i)):
                    for x1 in np.arange(np.abs(x - j)):
                        if m < 0:
                            if np.isclose(i + y1, m * (j - x1) + b, atol=atol):
                                # print(f"(x', y') = ({j + x1}, {i + y1}), m.x1+b: {m*(j+x1)+b}")
                                try:
                                    sc_im_edt_cp[i + y1][j - x1] = p
                                except IndexError:
                                    pass
                            if np.isclose(i - y1, m * (j + x1) + b, atol=atol):
                                # print(f"(x', y') = ({j - x1}, {i - y1}), m.x1+b: {m*(j-x1)+b}")
                                try:
                                    sc_im_edt_cp[i - y1][j + x1] = p
                                except IndexError:
                                    pass
                        else:
                            if np.isclose(i + y1, m * (j + x1) + b, atol=atol):
                                # print(f"(x', y') = ({j + x1}, {i + y1}), m.x1+b: {m*(j+x1)+b}")
                                try:
                                    sc_im_edt_cp[i + y1][j + x1] = p
                                except IndexError:
                                    pass
                            if np.isclose(i - y1, m * (j - x1) + b, atol=atol):
                                # print(f"(x', y') = ({j - x1}, {i - y1}), m.x1+b: {m*(j-x1)+b}")
                                try:
                                    sc_im_edt_cp[i - y1][j - x1] = p
                                except IndexError:
                                    pass

# Plot the result and save the raster image
fig, ax = plt.subplots(figsize=(12, 8))
ax.imshow(sc_im_edt_cp, cmap=cmap_coolwarm)
ax.contour(im, [0.5], colors='gray')
ax.axis('off')
plt.tight_layout()
plt.savefig(f'serpent sc.{image_type}', dpi=savefig_dpi, bbox_inches='tight', facecolor='black')

当前效果不完善的输出结果：

我肯定是在重复造轮子，这类问题应该有更高效的标准实现方法，求各位给出优化建议。

内容的提问来源于stack exchange，提问作者François