import torch import argparse import os import logging import subprocess import numpy as np import cv2 import time from collections import defaultdict import tqdm from multiprocessing import Pool, cpu_count import multiprocessing from typing import List, Tuple, Dict # 如果还未导入 def read_vertices(obj_path): vertices = [] with open(obj_path, 'r') as file: lines = file.readlines() for line in lines: if line.startswith('v '): # 顶点坐标 vertices.append(list(map(float, line.split()[1:4]))) vertices = torch.tensor(vertices) return vertices def read_uvs(obj_path): uv_coordinates = [] with open(obj_path, 'r') as file: lines = file.readlines() for line in lines: if line.startswith('vt '): # UV 坐标 uv_coordinates.append(list(map(float, line.split()[1:3]))) uv_coordinates = torch.tensor(uv_coordinates) return uv_coordinates def read_faces(obj_path): vertex_indices = [] uv_indices = [] with open(obj_path, 'r') as file: lines = file.readlines() for line in lines: if line.startswith('f '): # 面 parts = line.split()[1:] v_indices = [] uv_indices_temp = [] for face in parts: v, vt = map(int, face.split('/')[:2]) v_indices.append(v - 1) uv_indices_temp.append(vt - 1) vertex_indices.append(v_indices) uv_indices.append(uv_indices_temp) vertex_indices = torch.tensor(vertex_indices) uv_indices = torch.tensor(uv_indices) return vertex_indices, uv_indices def read_missing_faces(missing_faces_path): with open(missing_faces_path, 'r') as file: lines = file.readlines() missing_color_faces = torch.tensor( [int(line.strip()) for line in lines] ) return missing_color_faces def read_uv_map(input_texture_path): uv_map = cv2.imread(input_texture_path) uv_map = cv2.cvtColor(uv_map, cv2.COLOR_BGR2RGB) uv_map = torch.from_numpy(uv_map) return uv_map def parse_obj_file_and_uv_map(obj_path, missing_faces_path, input_texture_path, device): print(f"Reading OBJ file: {obj_path}") # vertices = [] # uv_coordinates = [] # vertex_indices = [] # uv_indices = [] # multiprocessing.set_start_method('spawn', force=True) # multiprocessing.freeze_support() start_time = time.time() p = Pool(5) uv_map_result = p.apply_async(read_uv_map, (input_texture_path,)) vertices_result = p.apply_async(read_vertices, (obj_path,)) uv_coordinates_result = p.apply_async(read_uvs, (obj_path,)) faces_result = p.apply_async(read_faces, (obj_path,)) missing_faces_result = p.apply_async(read_missing_faces, (missing_faces_path,)) p.close() p.join() vertices = vertices_result.get() uv_coordinates = uv_coordinates_result.get() vertex_indices, uv_indices = faces_result.get() missing_color_faces = missing_faces_result.get() uv_map = uv_map_result.get() vertices = vertices.to(device) uv_coordinates = uv_coordinates.to(device) vertex_indices = vertex_indices.to(device) uv_indices = uv_indices.to(device) missing_color_faces = missing_color_faces.to(device) uv_map = uv_map.to(device) end_time = time.time() print(f"using: {end_time - start_time} seconds") # exit() print("Converting to tensors...") return vertices, uv_coordinates, vertex_indices, uv_indices, missing_color_faces, uv_map def write_obj_with_uv_coordinates(filename, vertices, uvs, vertex_indices, uv_indices): """ 高性能OBJ文件写入函数 Parameters: filename (str): 输出OBJ文件路径 vertices (np.ndarray): 顶点数组 uvs (np.ndarray): UV坐标数组 vertex_indices (np.ndarray): 面的顶点索引 uv_indices (np.ndarray): 面的UV索引 """ # 估算数据大小（以字节为单位） estimated_size = ( len(vertices) * 40 + # 每个顶点约40字节 (v x.xxxxxx y.xxxxxx z.xxxxxx\n) len(uvs) * 30 + # 每个UV坐标约30字节 (vt x.xxxxxx y.xxxxxx\n) len(vertex_indices) * 40 # 每个面约40字节 (f v1/vt1 v2/vt2 v3/vt3\n) ) # 设置缓冲区大小为估算大小的1.2倍，最小256MB，最大1GB buffer_size = min(max(int(estimated_size * 1.2), 256 * 1024 * 1024), 1024 * 1024 * 1024) # 使用格式化字符串和列表推导式优化字符串生成 vertex_lines = ['v %.6f %.6f %.6f' % (v[0], v[1], v[2]) for v in vertices] uv_lines = ['vt %.6f %.6f' % (uv[0], uv[1]) for uv in uvs] # 优化face数据处理 face_lines = [] face_format = 'f %d/%d %d/%d %d/%d' for v_idx, uv_idx in zip(vertex_indices, uv_indices): face_lines.append(face_format % ( v_idx[0] + 1, uv_idx[0] + 1, v_idx[1] + 1, uv_idx[1] + 1, v_idx[2] + 1, uv_idx[2] + 1 )) # 使用join一次性构建完整内容 content = ['mtllib mesh.mtl'] + vertex_lines + [''] + uv_lines + [''] + ['usemtl material_0'] + face_lines # 一次性写入所有数据 with open(filename, 'w', buffering=buffer_size) as f: f.write('\n'.join(content)) def load_regions(filename): regions = [] with open(filename, 'r') as file: for line in file: parts = line.split(";") if len(parts) != 2: continue # Skip any lines that don't have exactly two parts first_set = set(int(x) for x in parts[0].strip().split()) second_set = set(int(x) for x in parts[1].strip().split()) regions.append((first_set, second_set)) return regions def build_face_adjacency(vertices, faces): """ 构建面的邻接关系，基于共享边 Args: vertices: 顶点数组 faces: 面片索引数组 (N x 3) Returns: dict: 面片邻接关系字典，key为面片索引，value为邻接面片索引列表 """ # 将faces转换为numpy数组以加快处理速度 faces = np.asarray(faces) num_faces = len(faces) # 为每个面创建所有边 (Nx3x2) edges = np.stack([ np.column_stack((faces[:, i], faces[:, (i + 1) % 3])) for i in range(3) ], axis=1) # 确保边的方向一致 (较小的顶点索引在前) edges.sort(axis=2) # 将边展平为 (Nx3, 2) 的形状 edges = edges.reshape(-1, 2) # 创建边到面的映射 edge_faces = np.repeat(np.arange(num_faces), 3) # 使用复合键对边进行排序 edge_keys = edges[:, 0] * vertices.shape[0] + edges[:, 1] sort_idx = np.argsort(edge_keys) edges = edges[sort_idx] edge_faces = edge_faces[sort_idx] # 找到重复的边（共享边） same_edges = (edge_keys[sort_idx][1:] == edge_keys[sort_idx][:-1]) edge_start_idx = np.where(same_edges)[0] # 构建邻接字典 face_adjacency = defaultdict(list) for idx in edge_start_idx: face1, face2 = edge_faces[idx], edge_faces[idx + 1] face_adjacency[face1].append(face2) face_adjacency[face2].append(face1) return dict(face_adjacency) def find_groups_and_subgroups(face_adjacency, missing_faces): """ 找到相连的面组和它们的邻接面 返回: regions: 列表，每个元素是一个元组 (missing_faces_set, adjacent_faces_set)， 与 load_regions() 函数返回格式保持一致 """ missing_faces_set = set(missing_faces.cpu().numpy()) unused_faces = set(missing_faces.cpu().numpy()) regions = [] total_faces = len(unused_faces) with tqdm.tqdm(total=total_faces, desc="Processing faces") as pbar: while unused_faces: start_face = unused_faces.pop() current_group = {start_face} current_subgroup = set() stack = [start_face] while stack: face_idx = stack.pop() for neighbor in face_adjacency.get(face_idx, []): if neighbor in unused_faces: current_group.add(neighbor) unused_faces.remove(neighbor) stack.append(neighbor) elif neighbor not in missing_faces_set: current_subgroup.add(neighbor) regions.append((current_group, current_subgroup)) pbar.update(total_faces - len(unused_faces) - pbar.n) # 输出统计信息 print(f"\nTotal regions: {len(regions)}") print(f"Average missing faces group size: {sum(len(g[0]) for g in regions)/len(regions):.2f}") print(f"Largest missing faces group size: {max(len(g[0]) for g in regions)}") print(f"Smallest missing faces group size: {min(len(g[0]) for g in regions)}") # 检查每个组是否都有邻接面 for i, (group, subgroup) in enumerate(regions): if not subgroup: print(f"Warning: Region {i} with {len(group)} missing faces has no adjacent faces!") return regions def compute_regions_face_colors( regions: List[Tuple[set, set]], uv_map: torch.Tensor, uvs: torch.Tensor, face_uv_indices: torch.Tensor, device: str ) -> Dict[int, torch.Tensor]: """ 根据每个区域的边缘面UV坐标计算加权平均的颜色， 当无有效采样时更新对应face_uv_indices。 参数: regions (List[Tuple[set, set]]): 每个区域为 (缺失面集合, 邻接面集合) uv_map (torch.Tensor): 原始纹理贴图，RGB格式 uvs (torch.Tensor): 原始UV坐标 face_uv_indices (torch.Tensor): 每个面对应的UV索引 device (str): 使用的设备("cuda"或"cpu") 返回: Dict[int, torch.Tensor]: 键为区域索引，值为该区域加权平均计算得到的颜色（uint8） """ regions_face_color: Dict[int, torch.Tensor] = {} for r_index, region in enumerate(tqdm.tqdm(regions, desc="Processing regions")): region_faces_indexes = torch.tensor(list(region[0]), device=device) region_edge_faces_indexes = torch.tensor(list(region[1]), device=device) if len(region_edge_faces_indexes) == 0: continue # 获取边缘面的UV索引 edge_face_uv_indices = face_uv_indices[region_edge_faces_indexes] # 使用三角形的质心UV坐标来采样颜色 triangle_uvs = uvs[edge_face_uv_indices] # shape: [num_faces, 3, 2] centroid_uvs = triangle_uvs.mean(dim=1) # shape: [num_faces, 2] # 将UV坐标转换为像素坐标 scale_tensor = torch.tensor([uv_map.shape[1] - 1, uv_map.shape[0] - 1], device=device) pixel_coords = torch.round(centroid_uvs * scale_tensor) pixel_coords[:, 1] = uv_map.shape[0] - 1 - pixel_coords[:, 1] pixel_coords = pixel_coords.long().clamp(0, uv_map.shape[0] - 1) # 直接采样质心位置的颜色 colors = uv_map[pixel_coords[:, 1], pixel_coords[:, 0]] # shape: [num_faces, 3] # 使用面积加权平均来计算最终颜色 areas = torch.abs( (triangle_uvs[:, 1, 0] - triangle_uvs[:, 0, 0]) * (triangle_uvs[:, 2, 1] - triangle_uvs[:, 0, 1]) - (triangle_uvs[:, 2, 0] - triangle_uvs[:, 0, 0]) * (triangle_uvs[:, 1, 1] - triangle_uvs[:, 0, 1]) ) * 0.5 if len(colors) > 0: weighted_color = (colors.float() * areas.unsqueeze(1)).sum(dim=0) / areas.sum() regions_face_color[r_index] = weighted_color.round().clamp(0, 255).to(torch.uint8) else: # 如果没有有效的采样点，使用第一个相邻面的UV坐标更新face_uv_indices face_uv_indices[region_faces_indexes] = face_uv_indices[region_edge_faces_indexes[0]].unsqueeze(dim=0).clone() return regions_face_color 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。 参数: uv_map (torch.Tensor): 原始纹理贴图，RGB格式 uvs (torch.Tensor): 原始UV坐标 face_uv_indices (torch.Tensor): 原始面的UV索引 regions (List[Tuple[set, set]]): 每个区域为 (缺失面集合, 邻接面集合) regions_face_color (Dict[int, torch.Tensor]): 每个区域计算得到的颜色 device (str): 使用的设备("cuda"或"cpu") 返回: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: new_uv_map: 更新后的UV贴图 uvs_updated: 更新后的UV坐标（拼接上新计算的UV） face_uv_indices: 更新后的face UV索引 """ total_regions = len(regions_face_color) grid_size = uv_map.shape[1] // 3 all_c = torch.div(torch.arange(total_regions, device=device), grid_size, rounding_mode='floor') all_r = torch.remainder(torch.arange(total_regions, device=device), grid_size) # 创建新的颜色UV贴图 color_uv_map = torch.full((int(uv_map.shape[0] / 2), uv_map.shape[1], 3), 255, dtype=torch.uint8, device=device) # 调整原始uvs的纵坐标 uvs[:, 1] = uvs[:, 1] * (2 / 3) + 1 / 3 # 批量创建所有颜色块的坐标 c_indices = all_c.unsqueeze(1).repeat(1, 9) * 3 + torch.tensor([0, 1, 2, 0, 1, 2, 0, 1, 2], device=device).unsqueeze(0) r_indices = all_r.unsqueeze(1).repeat(1, 9) * 3 + torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2], device=device).unsqueeze(0) # 批量设置颜色 colors = torch.stack([color for _, color in sorted(regions_face_color.items(), key=lambda x: x[0])]) colors_repeated = colors.unsqueeze(1).repeat(1, 9, 1) color_uv_map[c_indices.flatten(), r_indices.flatten()] = colors_repeated.reshape(-1, 3) # 批量计算新的UV坐标 pixels = torch.stack([ all_r * 3 + 1, uv_map.shape[0] + all_c * 3 + 1 ], dim=1).to(device) u_new = pixels[:, 0].float() / (uv_map.shape[1] - 1) new_height = int(uv_map.shape[0] + uv_map.shape[0] / 2) v_new = (new_height - 1 - pixels[:, 1].float()) / (new_height - 1) new_uvs = torch.stack([u_new, v_new], dim=1) # 更新UV坐标：拼接新计算的UV uvs_updated = torch.cat([uvs, new_uvs], dim=0) uv_coordinates_start = uvs_updated.shape[0] - total_regions # 批量更新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) # 合并原始UV贴图和新的颜色UV贴图 new_uv_map = torch.cat((uv_map, color_uv_map), dim=0) return new_uv_map, uvs_updated, face_uv_indices def group_regions_by_y_axis( regions: List[Tuple[set, set]], vertices: torch.Tensor, triangle_vertex_indices: torch.Tensor, device: str, interval_size: float = 0.1 ) -> Dict[int, List[int]]: """ 将区域按照y轴高度分组 Args: regions: 区域列表，每个区域为(缺失面集合, 邻接面集合)的元组 vertices: 顶点坐标张量 triangle_vertex_indices: 三角形顶点索引张量 device: 计算设备 ('cuda' 或 'cpu') interval_size: y轴分组的间隔大小，默认为0.1 Returns: Dict[int, List[int]]: 以y轴区间为键，区域索引列表为值的字典 """ y_intervals = defaultdict(list) for r_index, region in enumerate(regions): region_faces_indexes = torch.tensor(list(region[0]), device=device) # 计算面组的平均y轴位置 face_vertices = vertices[triangle_vertex_indices[region_faces_indexes]] avg_y = face_vertices[:, :, 1].mean(dim=(0, 1)) # 根据y轴位置分配到对应区间 interval_key = int(avg_y // interval_size) y_intervals[interval_key].append(r_index) return dict(y_intervals) def align_regions_colors( regions_face_color: Dict[int, torch.Tensor], y_intervals: Dict[int, List[int]], regions: List[Tuple[set, set]] ) -> Dict[int, torch.Tensor]: """ 对齐区间内的颜色 Args: regions_face_color: 每个区域的颜色 y_intervals: 每个y轴区间的区域索引列表 Returns: Dict[int, torch.Tensor]: 以y轴区间为键，颜色为值的字典 """ # aligned_regions_face_color = {} large_group_threshold_min = 5000 large_group_threshold_max = 100000 for interval_key, region_indices in y_intervals.items(): large_groups = [] # normal_groups = [] for r_index in region_indices: region = regions[r_index] if len(region[0]) >= large_group_threshold_min and len(region[0]) <= large_group_threshold_max: large_groups.append((r_index, len(region[0]), regions_face_color[r_index])) # 查找 large_groups 中 len(region[0]) 最大的组，并获取其颜色 if large_groups: largest_group = max(large_groups, key=lambda x: x[1]) color: torch.Tensor = largest_group[2] for large_group in large_groups: regions_face_color[large_group[0]] = color return regions_face_color def process(input_obj_path, input_texture_path, missing_faces_path, output_obj_path, output_texture_path): start_time = time.time() device = 'cuda' if torch.cuda.is_available() else 'cpu' vertices, uvs, triangle_vertex_indices, face_uv_indices, missing_color_faces, uv_map = parse_obj_file_and_uv_map( input_obj_path, missing_faces_path, input_texture_path, device=device) # 构建面的邻接关系和找到区域 start_face_adjacency_time = time.time() face_adjacency = build_face_adjacency(vertices.cpu().numpy(), triangle_vertex_indices.cpu().numpy()) end_face_adjacency_time = time.time() print(f"face_adjacency using: {end_face_adjacency_time - start_face_adjacency_time} seconds") start_find_groups_time = time.time() regions = find_groups_and_subgroups(face_adjacency, missing_color_faces) end_find_groups_time = time.time() print(f"find_groups_and_subgroups using: {end_find_groups_time - start_find_groups_time} seconds") start_texture_map_time = time.time() # 使用新封装的函数计算每个区域的加权平均颜色 regions_face_color = compute_regions_face_colors(regions, uv_map, uvs, face_uv_indices, device) end_texture_map_time = time.time() print(f"texture_mapping_to_triangle using: {end_texture_map_time - start_texture_map_time} seconds") # 按y轴区间分组 y_intervals = group_regions_by_y_axis( regions, vertices, triangle_vertex_indices, device ) # 对齐区间内的颜色 regions_face_color = align_regions_colors(regions_face_color, y_intervals, regions) # 更新UV贴图和面索引 start_color_map_time = time.time() new_uv_map, uvs, face_uv_indices = update_uv_map_and_indices(uv_map, uvs, face_uv_indices, regions, regions_face_color, device) end_color_map_time = time.time() print(f"color_mapping_to_triangle using: {end_color_map_time - start_color_map_time} seconds") end_time = time.time() print(f"using: {end_time - start_time} seconds") # 写入OBJ和纹理贴图 start_write_time = time.time() vertices_cpu = vertices.cpu().numpy() uvs_cpu = uvs.cpu().numpy() triangle_vertex_indices_cpu = triangle_vertex_indices.cpu().numpy() face_uv_indices_cpu = face_uv_indices.cpu().numpy() new_uv_map_cpu = new_uv_map.cpu().numpy() new_uv_map_bgr = cv2.cvtColor(new_uv_map_cpu, cv2.COLOR_RGB2BGR) with Pool(2) as p: # 异步执行OBJ和纹理图写入操作 obj_future = p.apply_async(write_obj_with_uv_coordinates, (output_obj_path, vertices_cpu, uvs_cpu, triangle_vertex_indices_cpu, face_uv_indices_cpu)) img_future = p.apply_async(cv2.imwrite, (output_texture_path, new_uv_map_bgr, [cv2.IMWRITE_PNG_COMPRESSION, 3])) obj_future.get() img_future.get() end_write_time = time.time() end_time = time.time() print(f"Total file writing time: {end_write_time - start_write_time:.2f} seconds") print(f"using: {end_time - start_time} seconds") def main(): parser = argparse.ArgumentParser(description='Process OBJ files to fix missing color faces.') parser.add_argument('--input_obj', type=str, required = True, help='Path to the input OBJ file') parser.add_argument('--input_texture', type=str, required = True, help='Path to the texture file') parser.add_argument('--missing_faces', type=str, required = True, help='Path to the file with indices of missing color faces') parser.add_argument('--output_obj', type=str, required = True, help='Path to the output OBJ file') parser.add_argument('--output_texture', type=str, required = True, help='Path to the texture file') args = parser.parse_args() process(args.input_obj, args.input_texture, args.missing_faces, args.output_obj, args.output_texture) if __name__ == '__main__': main()