Implementing 2 stage training

2025-12-11 12:04:08 +02:00
parent 221c80aa8c
commit c0684a9c14
5 changed files with 315 additions and 5 deletions
--- a/src/utils/config_manager.py
+++ b/src/utils/config_manager.py
@@ -58,6 +58,10 @@ class ConfigManager:
            "models": {
                "default_base_model": "yolov8s-seg.pt",
                "models_directory": "data/models",
+                "base_model_choices": [
+                    "yolov8s-seg.pt",
+                    "yolov11s-seg.pt",
+                ],
            },
            "training": {
                "default_epochs": 100,
@@ -65,6 +69,20 @@ class ConfigManager:
                "default_imgsz": 640,
                "default_patience": 50,
                "default_lr0": 0.01,
+                "two_stage": {
+                    "enabled": False,
+                    "stage1": {
+                        "epochs": 20,
+                        "lr0": 0.0005,
+                        "patience": 10,
+                        "freeze": 10,
+                    },
+                    "stage2": {
+                        "epochs": 150,
+                        "lr0": 0.0003,
+                        "patience": 30,
+                    },
+                },
            },
            "detection": {
                "default_confidence": 0.25,