优化目前版本中的问题

app/service/project_train_service.py

@@ -36,17 +36,18 @@ def add_project(info: ProjectInfoIn, session: Session, user_id: int):
     return project_info.id
 
 
-def check_image_name(project_id: int, files: List[UploadFile], session: Session):
+def check_image_name(project_id: int, img_type: str, files: List[UploadFile], session: Session):
     for file in files:
-        if not pimc.check_img_name(project_id, file.filename, session):
+        if not pimc.check_img_name(project_id, img_type, file.filename, session):
             return False, file.filename
     return True, None
 
 
-def upload_project_image(project_info: ProjectInfoOut, files: List[UploadFile], session: Session):
+def upload_project_image(project_info: ProjectInfoOut, img_type: str, files: List[UploadFile], session: Session):
     """
     上传项目的图片
     :param files: 上传的图片
+    :param img_type: 上传的图片类别
     :param project_info: 项目信息
     :param session:
     :return:
@@ -56,6 +57,7 @@ def upload_project_image(project_info: ProjectInfoOut, files: List[UploadFile],
         image = ProjectImage()
         image.project_id = project_info.id
         image.file_name = file.filename
+        image.img_type = img_type
         # 保存原图
         path = os.save_images(images_url, project_info.project_no, file=file)
         image.image_url = path
@@ -124,11 +126,9 @@ def run_train_yolo(project_info: ProjectInfoOut, session: Session):
     :param session: 数据库session
     :return:
     """
-    # 先获取项目的所有图片
-    project_images = pimc.get_images(project_info.id, session)
-
-    # 将图片根据,根据3:1的比例将图片分成train：val的两个数组
-    project_images_train, project_images_val = split_array(project_images)
+    # 先查询两个图片列表
+    project_images_train = pimc.get_images(project_info.id, 'train', session)
+    project_images_val = pimc.get_images(project_info.id, 'val', session)
 
     # 得到训练版本
     version_path = 'v' + str(project_info.train_version + 1)