Refactor: split_data() 예외처리 추가

2024-09-27 19:06:23 +09:00 · 2024-09-27 19:06:23 +09:00 · 7e75f41c64
commit 7e75f41c64
parent 7dd09182b8
1 changed files with 3 additions and 0 deletions
--- a/ai/app/api/yolo/detection.py
+++ b/ai/app/api/yolo/detection.py
@ -161,6 +161,9 @@ def split_data(data:list[TrainDataInfo], ratio:float):
        random.shuffle(data)
        train_data = data[:train_size]
        val_data = data[train_size:]
+        
+        if not train_data or not val_data:
+            raise Exception("data size is too small")
        return train_data, val_data
    except Exception as e:
        raise HTTPException(status_code=500, detail="exception in split_data(): " + str(e))