Quellcode durchsuchen

feat: 添加收集处理结果到CSV文件的功能,支持指定输出路径

zhch158_admin vor 3 Monaten
Ursprung
Commit
e68fa42955
1 geänderte Dateien mit 14 neuen und 2 gelöschten Zeilen
  1. 14 2
      zhch/ppstructurev3_single_process.py

+ 14 - 2
zhch/ppstructurev3_single_process.py

@@ -27,6 +27,7 @@ from utils import (
     get_image_files_from_dir,
     get_image_files_from_dir,
     get_image_files_from_list,
     get_image_files_from_list,
     get_image_files_from_csv,
     get_image_files_from_csv,
+    collect_pid_files
 )
 )
 
 
 def process_images_single_process(image_paths: List[str],
 def process_images_single_process(image_paths: List[str],
@@ -185,7 +186,8 @@ def main():
     parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
     parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
     parser.add_argument("--input_pattern", type=str, default="*", help="Input file pattern")
     parser.add_argument("--input_pattern", type=str, default="*", help="Input file pattern")
     parser.add_argument("--test_mode", action="store_true", help="Test mode (process only 20 images)")
     parser.add_argument("--test_mode", action="store_true", help="Test mode (process only 20 images)")
-    
+    parser.add_argument("--collect_results",type=str, help="收集处理结果到指定CSV文件")
+
     args = parser.parse_args()
     args = parser.parse_args()
     
     
     try:
     try:
@@ -276,7 +278,16 @@ def main():
             json.dump(final_results, f, ensure_ascii=False, indent=2)
             json.dump(final_results, f, ensure_ascii=False, indent=2)
         
         
         print(f"💾 Results saved to: {output_file}")
         print(f"💾 Results saved to: {output_file}")
-        
+
+        if args.collect_results:
+            processed_files = collect_pid_files(output_file)
+            output_file_processed = Path(args.collect_results).resolve()
+            with open(output_file_processed, 'w', encoding='utf-8') as f:
+                f.write("image_path,status\n")
+                for file_path, status in processed_files:
+                    f.write(f"{file_path},{status}\n")
+            print(f"💾 Processed files saved to: {output_file_processed}")
+
         return 0
         return 0
         
         
     except Exception as e:
     except Exception as e:
@@ -300,6 +311,7 @@ if __name__ == "__main__":
             "pipeline": "PP-StructureV3",
             "pipeline": "PP-StructureV3",
             "device": "gpu:0",
             "device": "gpu:0",
             "batch_size": 2,
             "batch_size": 2,
+            "collect_results": "./OmniDocBench_Results_Single/processed_files.csv",
         }
         }
         
         
         # 构造参数
         # 构造参数