增加对训练和推理中图片情况的判断

2025-03-12 09:39:50 +08:00
parent f49a6caf10
commit 76f4a5ecd9
5 changed files with 51 additions and 1 deletions
--- a/app/api/business/project_train_api.py
+++ b/app/api/business/project_train_api.py
@ -243,6 +243,12 @@ async def run_train(train_in: ProjectTrainIn, session: Session = Depends(get_db)
        return rc.response_error("请先上传验证图片")
    if val_img_count < 5:
        return rc.response_error("验证图片少于5张，请继续上传验证图片")
+    train_label_count = pimc.check_image_label(project_id, 'train', session)
+    if train_label_count > 0:
+        return rc.response_error("训练图片中存在未标注的图片")
+    val_label_count = pimc.check_image_label(project_id, 'val', session)
+    if val_label_count > 0:
+        return rc.response_error("验证图片中存在未标注的图片")
    data, project, name = ps.run_train_yolo(project_info, train_in, session)
    thread_train = threading.Thread(target=run_event_loop, args=(data, project, name, train_in,
                                    project_id, session,))