代码优化需求：提升法线贴图构建速度（3D与2D图形）

问题背景 #

我拥有一个已三角化的3D人脸网格，已计算出每个三角形的中点及法向量(n1,n2,n3)。目标是构建一个[width×Height]的2D网格（类似图像，z值设为常数）并覆盖在人脸网格上，再通过法向量生成对应的RGB三元组得到最终图像。当前代码功能正确，但运行速度极慢，急需优化以提升计算效率与速度。

核心性能瓶颈 #

代码中标记为# ---------------------------------Need Optimization HERE - VERY SLOW ----------------------的区块是性能瓶颈：这段代码对每个像素中点遍历所有三角形计算距离，时间复杂度为O(像素总数 × 三角形总数)，当像素为600×800=48万、三角形数量达数千时，总运算量会达到百亿级，导致运行极慢。

优化方案 #

• 向量化计算替代Python循环：利用NumPy的向量化操作批量计算所有像素中点到三角形中点的距离，避免Python循环的高开销。
• 预计算数组化的三角形数据：将字典存储的三角形中点和法向量转换为NumPy数组，便于批量运算。
• 简化索引映射：直接通过像素坐标的索引关系定位，无需提前构建indexing字典。
• 移除冗余操作：删除循环内的打印语句，避免IO开销。

完整优化代码 #

import os
import math
import numpy as np
import cv2
import matplotlib.pyplot as plt

def data_cleaning(filename):
    # 用NumPy加速数据读取与解析，替代逐行循环
    with open(filename, 'r') as f:
        lines = f.read().splitlines()
    # 统一处理每行数据
    data = []
    for line in lines:
        line = line.split("(")[-1].split(")")[0]
        parts = line.split(",")
        no = int(parts[0])
        x = float(parts[1])
        y = float(parts[2])
        z = float(parts[3])
        data.append([no, x, y, z])
    data_np = np.array(data)
    data_dic = dict(zip(data_np[:,0].astype(int), data_np[:,1:].tolist()))
    data_list = data_np[:,1:].tolist()
    return data_dic, data_list

def Pixel_Grid(Width, Height):
    TOP_LEFT_X = -0.9296
    TOP_RIGHT_X = 1.053
    TOP_LEFT_Y = -0.8783
    BOTTOM_LEFT_Y = 1.311
    x = np.linspace(TOP_LEFT_X, TOP_RIGHT_X, Width)
    y = np.linspace(TOP_LEFT_Y, BOTTOM_LEFT_Y, Height)
    XX, YY = np.meshgrid(x, y, sparse=True)
    return XX, YY

def mid_points(xx, yy, w, h):
    # 向量化生成所有像素中点，替代嵌套循环
    xx = xx.flatten()
    yy = yy.flatten()
    # 生成相邻x和y的中点
    x_mids = (xx[:-1] + xx[1:]) / 2
    y_mids = (yy[:-1] + yy[1:]) / 2
    # 网格化所有中点
    all_midpoints = np.array(np.meshgrid(x_mids, y_mids)).T.reshape(-1, 2)
    return all_midpoints

def Normal_map_fn(triangle_mids, triangle_normals, Mid_Points, Width, Height, Channels):
    Normal_map = np.zeros((Height, Width, Channels), np.float32)
    
    # 向量化计算所有像素中点到三角形中点的距离（仅计算xy平面）
    mid_points_np = Mid_Points
    triangle_mids_xy = triangle_mids[:, :2]
    
    # 广播计算距离平方（避免开根号，不影响最小值判断）
    dist_sq = np.sum((mid_points_np[:, np.newaxis, :] - triangle_mids_xy[np.newaxis, :, :]) ** 2, axis=2)
    
    # 找到每个像素中点对应的最近三角形索引
    closest_tri_idx = np.argmin(dist_sq, axis=1)
    
    # 获取对应的法向量
    closest_normals = triangle_normals[closest_tri_idx]
    
    # 将法向量赋值到Normal_map，直接通过索引映射
    y_coords = np.repeat(np.arange(Height-1), Width-1)
    x_coords = np.tile(np.arange(Width-1), Height-1)
    Normal_map[y_coords, x_coords] = closest_normals
    
    fig, ax = plt.subplots()
    cax = plt.imshow(Normal_map, cmap='gray')
    plt.show()
    name0 = os.path.join("Normal_map.png")
    # 法向量值映射到0-255区间再保存
    cv2.imwrite(name0, (Normal_map + 1) * 127.5)
    return Normal_map

# Task 0 - 读取并清理网格、法线、三角形数据
Vertex_Data = '../3D_Info/Vertex_Data.txt'
Triangle_Index = '../3D_Info/Mesh_Data.txt'
Normal_Vector = '../3D_Info/Face_Normal.txt'

Vertex, v_list = data_cleaning(Vertex_Data)
Triangle, t_list = data_cleaning(Triangle_Index)
Normal, n_list = data_cleaning(Normal_Vector)

# Task 1 - 移除背向相机的法线和三角形
Camera_Vector = np.array([0, 0, 1])
normals_np = np.array(n_list)
# 计算法线与相机向量的点积
dot_products = np.dot(normals_np, Camera_Vector)
# 筛选出正向的法线索引
positive_idx = dot_products > 0
# 更新Triangle和Normal
Triangle = {k: v for k, v in Triangle.items() if positive_idx[k]}
Normal = {k: v for k, v in Normal.items() if positive_idx[k]}

# Task 3 - 预计算三角形中点和法向量的NumPy数组
triangle_mids = []
triangle_normals = []
for k, v in Triangle.items():
    one, two, three = v
    point_m = Vertex[one]
    point_n = Vertex[two]
    point_o = Vertex[three]
    mid_x = (point_m[0] + point_n[0] + point_o[0]) / 3
    mid_y = (point_m[1] + point_n[1] + point_o[1]) / 3
    mid_z = (point_m[2] + point_n[2] + point_o[2]) / 3
    triangle_mids.append([mid_x, mid_y, mid_z])
    triangle_normals.append(Normal[k])
triangle_mids = np.array(triangle_mids)
triangle_normals = np.array(triangle_normals)

# Task 4- 生成法线图
Width = 600
Height = 800
Channels = 3
XX, YY = Pixel_Grid(Width, Height)
Mid_Points = mid_points(XX, YY, Width, Height)
Normal_map = Normal_map_fn(triangle_mids, triangle_normals, Mid_Points, Width, Height, Channels)

优化说明 #

• 向量化计算：将原来的双层嵌套循环替换为NumPy广播运算，运算速度提升数十到上百倍。
• 距离计算优化：通过计算距离平方替代开根号，减少计算量且不影响最小值判断。
• 数据结构优化：将字典存储的三角形数据转换为NumPy数组，避免字典遍历的开销。
• 图像保存修复：原代码直接保存法向量会因值在[-1,1]区间导致图像全黑，优化后将值映射到[0,255]区间。

内容的提问来源于stack exchange，提问作者Joker