correct albumentation

try to make augmentation on GPU.
saved llm prompt

train.py

@@ -1,4 +1,4 @@
-import io
+import concurrent.futures
 import os
 import random
 import shutil
@@ -31,7 +31,7 @@ test_set = 10
 old_images_percentage = 75
 
 DEFAULT_CLASS_NUM = 80
-
+total_files_copied = 0
 
 def form_dataset(from_date: datetime):
     makedirs(today_dataset, exist_ok=True)
@@ -67,6 +67,25 @@ def form_dataset(from_date: datetime):
 
 
 def copy_annotations(images, folder):
+    global total_files_copied
+    total_files_copied = 0
+
+    def copy_image(image):
+        global total_files_copied
+        total_files_copied += 1
+        label_name = f'{Path(image.path).stem}.txt'
+        label_path = path.join(processed_labels_dir, label_name)
+        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:
+            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})')
+
+        if total_files_copied % 1000 == 0:
+            print(f'{total_files_copied} copied...')
+
     destination_images = path.join(today_dataset, folder, 'images')
     makedirs(destination_images, exist_ok=True)
 
@@ -78,19 +97,10 @@ def copy_annotations(images, folder):
 
     copied = 0
     print(f'Copying annotations to {destination_images} and {destination_labels} folders:')
-    for image in images:
-        label_name = f'{Path(image.path).stem}.txt'
-        label_path = path.join(processed_labels_dir, label_name)
-        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:
-            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})')
-        copied = copied + 1
-        if copied % 1000 == 0:
-            print(f'{copied} copied...')
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        executor.map(copy_image, images)
+
+
     print(f'Copied all {copied} annotations to {destination_images} and {destination_labels} folders')
 
 
@@ -143,12 +153,15 @@ def revert_to_processed_data(date):
 
 def get_latest_model():
     def convert(d: str):
+        if not d.startswith(prefix):
+            return None
         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}
 
     dates = [convert(d) for d in next(os.walk(models_dir))[1]]
-    sorted_dates = list(sorted(dates, key=lambda x: x['date']))
+    dates = list(filter(lambda x : x is not None, dates))
+    sorted_dates = list(sorted(dates, key=lambda x: x['date'] ))
     if len(sorted_dates) == 0:
         return None, None
     last_model = sorted_dates[-1]
@@ -223,9 +236,8 @@ def validate(model_path):
 
 
 def upload_model(model_path: str):
-
-    # model = YOLO(model_path)
-    # model.export(format="onnx", imgsz=1280, nms=True, batch=4)
+    model = YOLO(model_path)
+    model.export(format="onnx", imgsz=1280, nms=True, batch=4)
     onnx_model = path.dirname(model_path) + Path(model_path).stem + '.onnx'
 
     with open(onnx_model, 'rb') as f_in:
@@ -250,9 +262,8 @@ def upload_model(model_path: str):
     api.upload_file('azaion.onnx.small', onnx_part_small)
 
 if __name__ == '__main__':
-    # model_path = train_dataset('2024-10-26', from_scratch=True)
-    # validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
-    # form_data_sample(500)
-    # convert2rknn()
-    model_path = 'azaion.pt'
-    upload_model(model_path)
+    model_path = train_dataset(from_scratch=True)
+    validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
+    form_data_sample(500)
+    convert2rknn()
+    upload_model('azaion.pt')