floor_plan_meshproject增加角度正则损失

数据结构

data.x  #(128, 16)  16 = triangles(9) + confidence(7)
"""
每个元素
[x1, y1, z1, x2, y2, z2, x3, y3, z3, 
c_v1, c_v2, c_v3, c_e1, c_e2, c_e3, c_f]
"""
data.y  #(128)


encoded_x #（128, 576)
encoded_x_conv # (2, 96, 576)
decoded_x_conv # (2, 96, 4)
decoded_x # (128, 4)

Angle_loss 计算

\[L_{\text{angle}} = p_{\text{merge}}^{\alpha} \times \min_{t \in \{90, 180, 270, 360\}} | \text{Angle}_{\text{merge}} - t | \]

\[Loss= w_1 L_{\text{classification}} + w_2 L_{\text{angle}} \]

1、朴素实现

data.x  #(num_triangles, 16) , 16 = triangles(9) + confidence(7)
# 每个元素
[x1, y1, z1, x2, y2, z2, x3, y3, z3, 
c_v1, c_v2, c_v3, c_e1, c_e2, c_e3, c_f]

data.y  #(128)

# 每个 triangle 的信息（添加了角）
data.x_with_angle
{
    face_idx: i,
    vertices: [(x1, y1, z1), (x2, y2, z2), (x3, y3, z3)],
    angles: [α1, α2, α3]
}

# 获取所有 room triangles

#取它们的顶点(去重)

(x1, y1, z1): {
    vertices: (x1, y1, z1),
    faces_idx: [1, 2, 3, ...], # 包含该顶点的room面idx
    angle: [α1, α2, α3, ...], # 该顶点在该面的angle
    decoded_x: [[p_11, p_12, p_13, p_14], [p_21, p_22, p_23, p_24], [p_31, p_32, p_33, p_34], ...]
}

# 计算angle_merge, p_merge

# loss 计算

朴素实现

def angle_loss_torch(self, data, decoded_x):
    """
    doc: 计算角度正则化损失
    return: total_angle_loss: 角度正则化损失
    """

    # 1. 选出所有最大分类概率为room的三角形(及其索引)
    softmax_decoded_x = torch.nn.functional.softmax(decoded_x, dim=1)  # 使用softmax得到不同分类的概率
    max_probs, max_indices = torch.max(softmax_decoded_x, dim=1)
    room_indices = torch.where(max_indices == 2)[0]

    # 获取每个面的三个角 data.x_with_angle
    vertices = data.x[:, :9].reshape(-1, 3, 3)  # 提取顶点坐标(num_triangle, 3, 3)
    angles = angle_func(vertices)  # 计算内角 (num_triangle, 3)
    data.x_with_angle = [
        {
            'triangle_idx': i,
            'vertices': [tuple(v) for v in vertices[i]],
            'angles': angles[i].tolist()
        }
        for i in range(vertices.shape[0])
    ]

    # 2. 取它们的顶点(去重)
    vertex_info = {}
    for i in room_indices:
        triangle = data.x_with_angle[i]
        face_idx = triangle['triangle_idx']
        vertices = triangle['vertices']
        angles = triangle['angles']
        
        for vertex, angle in zip(vertices, angles):
            vertex_key = tuple(vertex)
            if vertex_key not in vertex_info:
                vertex_info[vertex_key] = {
                    'vertices': vertex,
                    'faces_idx': [],
                    'angles': [],
                    'decoded_x': []
                }
            
            vertex_info[vertex_key]['faces_idx'].append(face_idx)
            vertex_info[vertex_key]['angles'].append(angle)
            vertex_info[vertex_key]['decoded_x'].append(softmax_decoded_x[face_idx])

    # 3. 计算每个顶点的角度总和和融合概率特征
    for vertex_key in vertex_info:
        vertex_info[vertex_key]['angle_merge'] = sum(vertex_info[vertex_key]['angles'])
        vertex_info[vertex_key]['p_merge'] = torch.mean(torch.stack(vertex_info[vertex_key]['decoded_x']), dim=0)

    # 4. angle_loss = p^alpha * min|θ - t|
    t_values = torch.tensor([np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], dtype=torch.float32)
    alpha = self.config.angle_loss_alpha
    total_angle_loss = 0

    for vertex_key, info in vertex_info.items():
        angle_sum = info['angle_merge']
        p = info['p_merge'][2]
        abs_diffs = torch.abs(angle_sum - t_values)
        min_abs_diff = torch.min(abs_diffs)
        angle_loss = p**alpha * min_abs_diff
        total_angle_loss += angle_loss

    return total_angle_loss

for循环太多，考虑能不能优化

2、优化循环

算法优化

# room 顶点
room_vertices = vertices[room_indices].reshape(-1, 3)  

# room angle
room_angles = angles[room_indices].reshape(-1)  # (num_room_triangles * 3)

# 从 softmax_decoded_x 中提取出属于“room”类别的三角形的概率分布，并将每个三角形的概率分布重复 3 次，以对应每个三角形的 3 个顶点
room_decoded_x = softmax_decoded_x[room_indices].repeat_interleave(3, dim=0)  # (num_room_triangles * 3, num_cls)


# 获取去重顶点
unique_vertices
# room_vertices 中相应元素在 unique_vertices 中的索引
inverse_indeces

# angle_merge、p_merge
angle_merge.scatter_add_(0, inverse_indices, room_angles)

优化循环

def angle_loss(self, data, decoded_x):
    """
    doc: 计算角度正则化损失
    return: total_angle_loss: 角度正则化损失
    """

    # 1. 选出所有最大分类概率为room的三角形(及其索引)
    softmax_decoded_x = torch.nn.functional.softmax(decoded_x, dim=1)  # 使用softmax得到不同分类的概率
    max_probs, max_indices = torch.max(softmax_decoded_x, dim=1)
    room_indices = torch.where(max_indices == 2)[0]

    # 获取每个面的三个角 data.x_with_angle
    vertices = data.x[:, :9].reshape(-1, 3, 3)  # 提取顶点坐标(num_triangle, 3, 3)
    angles = angle_func(vertices)  # 计算内角 (num_triangle, 3)


    # 2. 取它们的顶点(去重)
    room_vertices = vertices[room_indices].reshape(-1, 3)  # (num_room_triangles * 3, 3)
    room_angles = angles[room_indices].reshape(-1)  # (num_room_triangles * 3)
    room_decoded_x = softmax_decoded_x[room_indices].repeat_interleave(3, dim=0)  # (num_room_triangles * 3, num_cls)

    # 3. 计算每个顶点的角度总和和融合概率特征
    unique_vertices, inverse_indices = torch.unique(room_vertices, dim=0, return_inverse=True)
    angle_merge = torch.zeros(unique_vertices.shape[0], dtype=torch.float32, device=room_vertices.device)
    p_merge = torch.zeros(unique_vertices.shape[0], softmax_decoded_x.shape[1], dtype=torch.float32, device=room_vertices.device)

    angle_merge.scatter_add_(0, inverse_indices, room_angles)
    p_merge.scatter_add_(0, inverse_indices.unsqueeze(1).expand(-1, softmax_decoded_x.shape[1]), room_decoded_x)
    p_merge = p_merge / torch.bincount(inverse_indices).unsqueeze(1).float()

    # 4. angle_loss = p^alpha * min|θ - t|
    t_values = torch.tensor([np.pi/2, np.pi, 3*np.pi/2, 2*np.pi], dtype=torch.float32, device=angle_merge.device)
    alpha = self.config.angle_loss_alpha
    total_angle_loss = 0

    for angle_sum, p in zip(angle_merge, p_merge):
        abs_diffs = torch.abs(angle_sum - t_values)
        min_abs_diff = torch.min(abs_diffs)
        angle_loss = p[2]**alpha * min_abs_diff
        total_angle_loss += angle_loss

    return total_angle_loss