去小黄点bug修复

fill_all_empty_faces_v1.2.py

@@ -323,7 +323,7 @@ def compute_regions_face_colors(
     return regions_face_color
 
 
-def update_uv_map_and_indices(
+def update_uv_map_and_indices2(
     uv_map: torch.Tensor,
     uvs: torch.Tensor,
     face_uv_indices: torch.Tensor,
@@ -394,6 +394,75 @@ def update_uv_map_and_indices(
 
     return new_uv_map, uvs_updated, face_uv_indices
 
+def update_uv_map_and_indices(
+    uv_map: torch.Tensor,
+    uvs: torch.Tensor,
+    face_uv_indices: torch.Tensor,
+    regions: List[Tuple[set, set]],
+    regions_face_color: Dict[int, torch.Tensor],
+    device: str
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    根据计算得到的区域颜色，更新UV贴图及对应的UV坐标，并批量更新face_uv_indices。
+    """
+    total_regions = len(regions_face_color)
+    
+    # 1. 在CPU上创建新的UV贴图
+    new_height = int(uv_map.shape[0] + uv_map.shape[0] / 2)
+    new_width = uv_map.shape[1]
+    
+    # 将原始UV贴图移到CPU
+    uv_map_cpu = uv_map.cpu().numpy()
+    
+    # 创建新的UV贴图（在CPU上）
+    new_uv_map_cpu = np.zeros((new_height, new_width, 3), dtype=np.uint8)
+    new_uv_map_cpu[:uv_map.shape[0], :, :] = uv_map_cpu
+    
+    # 2. 在CPU上创建颜色块
+    grid_size = new_width // 3
+    for i, (region_index, color) in enumerate(regions_face_color.items()):
+        all_c = i // grid_size
+        all_r = i % grid_size
+        
+        # 计算3x3颜色块的位置
+        for dr in range(3):
+            for dc in range(3):
+                r = all_r * 3 + dr
+                c = all_c * 3 + dc + uv_map.shape[0]  # 偏移到下半部分
+                
+                if r < new_width and c < new_height:
+                    new_uv_map_cpu[c, r, :] = color.cpu().numpy()
+    
+    # 3. 更新UV坐标
+    uvs[:, 1] = uvs[:, 1] * (2 / 3) + 1 / 3
+    
+    # 计算新的UV坐标
+    new_uvs = []
+    for i, (region_index, _) in enumerate(sorted(regions_face_color.items(), key=lambda x: x[0])):
+        r = (i % grid_size) * 3 + 1
+        c = uv_map.shape[0] + (i // grid_size) * 3 + 1
+        
+        u_new = r / (new_width - 1)
+        v_new = (new_height - 1 - c) / (new_height - 1)
+        new_uvs.append([u_new, v_new])
+    
+    new_uvs_tensor = torch.tensor(new_uvs, device=device)
+    uvs_updated = torch.cat([uvs, new_uvs_tensor], dim=0)
+    uv_coordinates_start = uvs_updated.shape[0] - total_regions
+    
+    # 4. 批量更新face_uv_indices
+    for i, (region_index, _) in enumerate(sorted(regions_face_color.items(), key=lambda x: x[0])):
+        region_faces_indexes = torch.tensor(list(regions[region_index][0]), device=device)
+        face_uv_indices[region_faces_indexes] = torch.full(
+            (1, 3), uv_coordinates_start + i, 
+            device=device, dtype=face_uv_indices.dtype
+        )
+    
+    # 将结果转回GPU
+    new_uv_map = torch.from_numpy(new_uv_map_cpu).to(device)
+    
+    return new_uv_map, uvs_updated, face_uv_indices
+
 def group_regions_by_y_axis(
     regions: List[Tuple[set, set]],
     vertices: torch.Tensor,