copy images and labels during forming dataset. add folder for corrupted labels, small refactor

train.py

@@ -1,12 +1,17 @@
-from os import path, replace, remove, listdir, makedirs
+import random
+from os import path, replace, remove, listdir, makedirs, scandir
 from os.path import abspath
 import shutil
 from datetime import datetime
 from pathlib import Path
 from ultralytics import YOLO
-from constants import processed_images_dir, processed_labels_dir, annotation_classes, prefix, date_format, datasets_dir, models_dir
+from constants import (processed_images_dir,
+                       processed_labels_dir,
+                       annotation_classes,
+                       prefix, date_format,
+                       datasets_dir, models_dir,
+                       corrupted_images_dir, corrupted_labels_dir)
 
-latest_model = path.join(models_dir, f'{prefix}latest.pt')
 today_folder = f'{prefix}{datetime.now():{date_format}}'
 today_dataset = path.join(datasets_dir, today_folder)
 train_set = 70
@@ -14,38 +19,61 @@ valid_set = 20
 test_set = 10
 
 
-def form_dataset():
+def form_dataset(set_date: datetime):
     makedirs(today_dataset, exist_ok=True)
-    images = listdir(processed_images_dir)
+    images = []
+    with scandir(processed_images_dir) as imd:
+        for image_file in imd:
+            if not image_file.is_file():
+                continue
+            mod_time = datetime.fromtimestamp(image_file.stat().st_mtime)
+            if set_date is None:
+                images.append(image_file)
+            elif mod_time > set_date:
+                images.append(image_file)
+
+    print('shuffling images')
+    random.shuffle(images)
 
     train_size = int(len(images) * train_set / 100.0)
     valid_size = int(len(images) * valid_set / 100.0)
 
-    move_annotations(images[:train_size], 'train')
-    move_annotations(images[train_size:train_size + valid_size], 'valid')
-    move_annotations(images[train_size + valid_size:], 'test')
+    print(f'copy train dataset, size: {train_size} annotations')
+    copy_annotations(images[:train_size], 'train')
 
+    print(f'copy valid set, size: {valid_size} annotations')
+    copy_annotations(images[train_size:train_size + valid_size], 'valid')
+
+    print(f'copy test set, size: {len(images) - train_size - valid_size} annotations')
+    copy_annotations(images[train_size + valid_size:], 'test')
+
+    print('creating yaml...')
     create_yaml()
 
 
-def move_annotations(images, folder):
+def copy_annotations(images, folder):
     destination_images = path.join(today_dataset, folder, 'images')
     makedirs(destination_images, exist_ok=True)
+
     destination_labels = path.join(today_dataset, folder, 'labels')
     makedirs(destination_labels, exist_ok=True)
-    for image_name in images:
-        image_path = path.join(processed_images_dir, image_name)
-        label_name = f'{Path(image_name).stem}.txt'
+
+    makedirs(corrupted_images_dir, exist_ok=True)
+    makedirs(corrupted_labels_dir, exist_ok=True)
+
+    for image in images:
+        label_name = f'{Path(image.path).stem}.txt'
         label_path = path.join(processed_labels_dir, label_name)
-        if not check_label(label_path):
-            remove(image_path)
+        if check_label(label_path):
+            shutil.copy(image.path, path.join(destination_images, image.name))
+            shutil.copy(label_path, path.join(destination_labels, label_name))
         else:
-            replace(image_path, path.join(destination_images, image_name))
-            replace(label_path, path.join(destination_labels, label_name))
+            shutil.copy(image.path, path.join(corrupted_images_dir, image.name))
+            shutil.copy(label_path, path.join(corrupted_labels_dir, label_name))
+            print(f'Label {label_path} is corrupted! Copy with its image to the corrupted directory ({corrupted_labels_dir})')
 
 
 def check_label(label_path):
-    lines_edited = False
     if not path.exists(label_path):
         return False
     with open(label_path, 'r') as f:
@@ -53,16 +81,7 @@ def check_label(label_path):
         for line in lines:
             for val in line.split(' ')[1:]:
                 if float(val) > 1:
-                    lines.remove(line)
-                    lines_edited = True
-    if len(lines) == 0:
-        return False
-    if not lines_edited:
-        return True
-
-    with open(label_path, 'w') as label_write:
-        label_write.writelines(lines)
-        label_write.close()
+                    return False
     return True
 
 
@@ -97,10 +116,25 @@ def revert_to_processed_data(date):
     shutil.rmtree(date_dataset)
 
 
-if __name__ == '__main__':
-    # form_dataset()
+def get_latest_model():
+    def convert(d: str):
+        dir_date = datetime.strptime(d.replace(prefix, ''), '%Y-%m-%d')
+        dir_model_path = path.join(models_dir, d, 'weights', 'best.pt')
+        return {'date': dir_date, 'path': dir_model_path}
 
-    model_name = latest_model if path.isfile(latest_model) else 'yolov8m.yaml'
+    dates = [convert(d) for d in listdir(models_dir)]
+    sorted_dates = list(sorted(dates, key=lambda x: x['date']))
+    if len(sorted_dates) == 0:
+        return None, None
+    last_model = sorted_dates[-1]
+    return last_model['date'], last_model['path']
+
+
+if __name__ == '__main__':
+    latest_date, latest_model = get_latest_model()
+    # form_dataset(latest_date)
+
+    model_name = latest_model if latest_model is not None and path.isfile(latest_model) else 'yolov8m.yaml'
     print(f'Initial model: {model_name}')
     model = YOLO(model_name)
     
@@ -108,7 +142,7 @@ if __name__ == '__main__':
 
     cur_folder = today_dataset
     yaml = abspath(path.join(cur_folder, 'data.yaml'))
-    results = model.train(data=yaml, epochs=100, batch=55, imgsz=640, save_period=1)
+    results = model.train(data=yaml, epochs=100, batch=60, imgsz=640, save_period=1)
 
     shutil.copy(f'{results.save_dir}/weights/best.pt', latest_model)
     shutil.copytree(results.save_dir, path.join(models_dir, cur_folder))