From 1f7b4a95d269074aa918fdef406c0edc9953832d Mon Sep 17 00:00:00 2001
From: Oleksandr Bezdieniezhnykh <oleksandr.bez@2c2p.com>
Date: Thu, 6 Jun 2024 01:46:48 +0300
Subject: [PATCH] small fix add code for revert dataset add get_recent_model

---
 train.py | 47 ++++++++++++++++++++++++++++-------------------
 1 file changed, 28 insertions(+), 19 deletions(-)

diff --git a/train.py b/train.py
index 787c3c7..c3d3ad4 100644
--- a/train.py
+++ b/train.py
@@ -36,8 +36,11 @@ def move_annotations(images, folder):
         image_path = os.path.join(current_images_dir, image_name)
         label_name = f'{Path(image_name).stem}.txt'
         label_path = os.path.join(current_labels_dir, label_name)
-        os.replace(image_path, os.path.join(destination_images, image_name))
-        os.replace(label_path, os.path.join(destination_labels, label_name))
+        if not os.path.exists(label_path):
+            os.remove(image_path)
+        else:
+            os.replace(image_path, os.path.join(destination_images, image_name))
+            os.replace(label_path, os.path.join(destination_labels, label_name))
 
 
 def create_yaml():
@@ -55,18 +58,17 @@ def create_yaml():
 
 
 def get_recent_model():
-    date_sets = []
-    datasets = [next((file for file in os.listdir(os.path.join('datasets', d)) if file.endswith('pt')), None)
-                for d in os.listdir('datasets')]
-
-    # date_str = d.replace(prefix, '')
-    # if date_str == 'current' or date_str == f'{datetime.now():%Y-%m-%d}':
-    #     continue
-    # if len(date_sets) == 0:
-    #     return None
-
-    recent = max(date_sets)
-    return os.path.join('datasets', f'{prefix}{recent}', f'{prefix}{recent}.pt')
+    cur_model = None
+    cur_date = None
+    for d in os.listdir('datasets'):
+        date_str = d.replace(prefix, '')
+        if date_str == 'current' or date_str == f'{datetime.now():%Y-%m-%d}':
+            continue
+        date = datetime.strptime(date_str, '%Y-%m-%d')
+        for file in os.listdir(os.path.join('datasets', d)):
+            if file.endswith('pt') and cur_date is None or cur_date < date:
+                cur_model = os.path.join('datasets', d, file)
+    return cur_model
 
 
 def retrain():
@@ -75,12 +77,19 @@ def retrain():
 
 
 def revert_to_current(date):
-    def revert_dir(dir):
-        os.listdir(os.path.join(current_images_dir, 'images'))
+    def revert_dir(src_dir, dest_dir):
+        for file in os.listdir(src_dir):
+            s = os.path.join(src_dir, file)
+            d = os.path.join(dest_dir, file)
+            os.replace(s, d)
+    date_dataset = os.path.join('datasets', f'{prefix}{date}')
+    current_dataset = os.path.join('datasets', f'{prefix}current')
+    for subset in ['test', 'train', 'valid']:
+        revert_dir(os.path.join(date_dataset, subset, 'images'), os.path.join(current_dataset, 'images'))
+        revert_dir(os.path.join(date_dataset, subset, 'labels'), os.path.join(current_dataset, 'labels'))
+    shutil.rmtree(date_dataset)
 
-    date_dataset = f'{prefix}{date}'
-    revert_dir(os.path.join(date_dataset, 'test'))
 
 form_dataset()
-create_yaml()
+# revert_to_current('2024-06-06')
 retrain()