C++和python剥离

libs/MVS/mask_face_occlusion.py

@@ -14,6 +14,9 @@ import torch
 import torch.nn.functional as F
 from torch.utils.dlpack import to_dlpack, from_dlpack
 
+import os
+from typing import Dict, List, Set
+
 class ModelProcessor:
     def __init__(self):
 
@@ -455,6 +458,8 @@ class ModelProcessor:
         print(f"所有图像处理完成，总耗时: {total_time:.2f}秒")
         print(f"平均每张图像耗时: {total_time/len(images):.2f}秒")
         
+        self.save_occlusion_data(visible_faces_dict, edge_faces_dict, delete_edge_faces_dict, self.asset_dir)
+        
         return {
             "result1": visible_faces_dict, 
             "result2": edge_faces_dict, 
@@ -1183,8 +1188,82 @@ class ModelProcessor:
         print(f"所有图像处理完成，总耗时: {total_time:.2f}秒")
         print(f"平均每张图像耗时: {total_time/len(images):.2f}秒")
 
+        self.save_occlusion_data(visible_faces_dict, edge_faces_dict, delete_edge_faces_dict, self.asset_dir)
+
         return {"result1": visible_faces_dict, "result2": edge_faces_dict, "result3": delete_edge_faces_dict}
 
+    def save_occlusion_data(self, result1: Dict[str, List[int]], 
+                        result2: Dict[str, List[int]], 
+                        result3: Dict[str, List[int]], 
+                        base_path: str) -> None:
+        """
+        保存遮挡数据到文件
+        
+        Args:
+            result1: 可见面字典，包含图像名称和对应的可见面列表
+            result2: 边面字典，包含图像名称和对应的边面列表  
+            result3: 删除边面字典，包含图像名称和对应的删除边面列表
+            base_path: 基础文件路径
+        """
+
+        print(f"save_occlusion_data {base_path}, {len(result1)}, {len(result2)}, {len(result3)}")
+        
+        # 处理返回的可见面字典 - 转换为图像名到面编号集合的映射
+        visible_faces_map: Dict[str, Set[int]] = {}
+        for image_name, face_list in result1.items():
+            visible_faces_map[image_name] = set(face_list)
+        
+        # 计算所有可见面的并集
+        face_visible_relative: Set[int] = set()
+        for face_set in visible_faces_map.values():
+            face_visible_relative.update(face_set)
+        
+        # 处理边面字典
+        edge_faces_map: Dict[str, Set[int]] = {}
+        for image_name, face_list in result2.items():
+            edge_faces_map[image_name] = set(face_list)
+        
+        # 处理删除边面字典
+        delete_edge_faces_map: Dict[str, Set[int]] = {}
+        for image_name, face_list in result3.items():
+            delete_edge_faces_map[image_name] = set(face_list)
+        
+        # 保存 visible_faces_map
+        try:
+            with open(base_path + "_visible_faces_map.txt", "w", encoding='utf-8') as map_file:
+                for image_name, face_set in visible_faces_map.items():
+                    # 写入图像名称和所有面ID，用空格分隔
+                    line = image_name + " " + " ".join(str(face) for face in face_set) + "\n"
+                    map_file.write(line)
+        except IOError as e:
+            print(f"Error writing visible_faces_map file: {e}")
+        
+        # 保存 face_visible_relative
+        try:
+            with open(base_path + "_face_visible_relative.txt", "w", encoding='utf-8') as relative_file:
+                for face in face_visible_relative:
+                    relative_file.write(str(face) + "\n")
+        except IOError as e:
+            print(f"Error writing face_visible_relative file: {e}")
+        
+        # 保存 edge_faces_map
+        try:
+            with open(base_path + "_edge_faces_map.txt", "w", encoding='utf-8') as map_file2:
+                for image_name, face_set in edge_faces_map.items():
+                    line = image_name + " " + " ".join(str(face) for face in face_set) + "\n"
+                    map_file2.write(line)
+        except IOError as e:
+            print(f"Error writing edge_faces_map file: {e}")
+        
+        # 保存 delete_edge_faces_map
+        try:
+            with open(base_path + "_delete_edge_faces_map.txt", "w", encoding='utf-8') as map_file3:
+                for image_name, face_set in delete_edge_faces_map.items():
+                    line = image_name + " " + " ".join(str(face) for face in face_set) + "\n"
+                    map_file3.write(line)
+        except IOError as e:
+            print(f"Error writing delete_edge_faces_map file: {e}")
+
     def process(self):
 
         print("process")