智能秤分析

data_preprocessing/data_split.py

@@ -2,7 +2,7 @@ import os
 import shutil
 import random
 from pathlib import Path
-
+import yaml
 def is_image_file(filename):
     """检查文件是否为图像文件"""
     image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff')
@@ -61,10 +61,12 @@ def control_train_number():
     pass
 
 if __name__ == "__main__":
-    # 设置目录路径
-    TRAIN_DIR = "scatter_data/train"
-    VAL_DIR = "scatter_data/val"
+    # # 设置目录路径
+    # TRAIN_DIR = "/home/lc/data_center/electornic/v1/train"
+    # VAL_DIR = "/home/lc/data_center/electornic/v1/val"
 
+    with open('../configs/scatter_data.yml', 'r') as f:
+        conf = yaml.load(f, Loader=yaml.FullLoader)
     print("开始分割数据集...")
-    split_directory(TRAIN_DIR, VAL_DIR)
+    split_directory(conf)
     print("数据集分割完成")