From 8fb24384db5013f4f484c5ef009c01ceee77981f Mon Sep 17 00:00:00 2001
From: zxsanny <zxsanny@gmail.com>
Date: Mon, 19 May 2025 18:32:41 +0300
Subject: [PATCH] add resume training possibility

---
 augmentation.py |  1 -
 train.py        | 11 +++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/augmentation.py b/augmentation.py
index 3c3bef5..d7aa7d4 100644
--- a/augmentation.py
+++ b/augmentation.py
@@ -22,7 +22,6 @@ class Augmentator:
         self.correct_min_bbox_size = 0.01
 
         self.transform = A.Compose([
-            # Flips, rotations and brightness
             A.HorizontalFlip(p=0.6),
             A.RandomBrightnessContrast(p=0.4, brightness_limit=(-0.3, 0.3), contrast_limit=(-0.05, 0.05)),
             A.Affine(p=0.8, scale=(0.8, 1.2), rotate=(-35, 35), shear=(-10, 10)),
diff --git a/train.py b/train.py
index 4ec3b06..dba64cc 100644
--- a/train.py
+++ b/train.py
@@ -172,6 +172,17 @@ def get_latest_model():
     return last_model['date'], last_model['path']
 
 
+def resume_training(last_pt_path):
+    model = YOLO(last_pt_path)
+    results = model.train(data=yaml,
+                          resume=True,
+                          epochs=120,
+                          batch=12,
+                          imgsz=1280,
+                          save_period=1,
+                          workers=24)
+
+
 def train_dataset(existing_date=None, from_scratch=False):
     latest_date, latest_model = get_latest_model() if not from_scratch else None, None