correct albumentation

try to make augmentation on GPU.
saved llm prompt

preprocessing.py

@@ -5,11 +5,29 @@ from pathlib import Path
 import albumentations as A
 import cv2
 import numpy as np
+import concurrent.futures
 
 from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir,
                        annotation_classes, checkpoint_file, checkpoint_date_format)
 from dto.imageLabel import ImageLabel
 
+total_files_processed = 0
+transform = A.Compose([
+    # Flips, rotations and brightness
+    A.HorizontalFlip(),
+    A.RandomBrightnessContrast(brightness_limit=(-0.05, 0.05), contrast_limit=(-0.05, 0.05)),
+    A.Affine(p=0.7, scale=(0.8, 1.2), rotate=25, translate_percent=0.1),
+
+    # Weather
+    A.RandomFog(p=0.2, fog_coef_range=(0, 0.3)),
+    A.RandomShadow(p=0.2),
+
+    # Image Quality/Noise
+    A.MotionBlur(p=0.2, blur_limit=(3, 5)),
+
+    # Color Variations
+    A.HueSaturationValue(p=0.3, hue_shift_limit=8, sat_shift_limit=8, val_shift_limit=8)
+], bbox_params=A.BboxParams(format='yolo'))
 
 def correct_bboxes(labels):
     margin = 0.0005
@@ -37,31 +55,18 @@ def correct_bboxes(labels):
 
 
 def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
-    transforms = [
-        A.Compose([A.HorizontalFlip(always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo', )),
-        A.Compose([A.RandomBrightnessContrast(always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.SafeRotate(limit=90, always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.SafeRotate(limit=90, always_apply=True),
-                   A.RandomBrightnessContrast(always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True),
-                   A.VerticalFlip(always_apply=True), ],
-                  bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.SafeRotate(limit=90, always_apply=True),
-                   A.RandomBrightnessContrast(always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo'))
-    ]
-
     results = []
     labels = correct_bboxes(img_ann.labels)
     if len(labels) == 0 and len(img_ann.labels) != 0:
         print('no labels but was!!!')
-    for i, transform in enumerate(transforms):
+    results.append(ImageLabel(
+            image=img_ann.image,
+            labels=img_ann.labels,
+            image_path=os.path.join(processed_images_dir, Path(img_ann.image_path).name),
+            labels_path=os.path.join(processed_labels_dir, Path(img_ann.labels_path).name)
+        )
+    )
+    for i in range(7):
         try:
             res = transform(image=img_ann.image, bboxes=labels)
             path = Path(img_ann.image_path)
@@ -87,7 +92,8 @@ def write_result(img_ann: ImageLabel):
                  img_ann.labels]
         f.writelines(lines)
         f.close()
-    print(f'{img_ann.labels_path} written')
+    global total_files_processed
+    print(f'{total_files_processed}. {img_ann.labels_path} written')
 
 
 def read_labels(labels_path) -> [[]]:
@@ -104,19 +110,10 @@ def read_labels(labels_path) -> [[]]:
         return arr
 
 
-def process_image(img_ann):
-    results = image_processing(img_ann)
-    for res_ann in results:
-        write_result(res_ann)
-    write_result(ImageLabel(
-        image=img_ann.image,
-        labels=img_ann.labels,
-        image_path=os.path.join(processed_images_dir, Path(img_ann.image_path).name),
-        labels_path=os.path.join(processed_labels_dir, Path(img_ann.labels_path).name)
-    ))
-
-
 def preprocess_annotations():
+    global total_files_processed  # Indicate that we're using the global counter
+    total_files_processed = 0
+
     os.makedirs(processed_images_dir, exist_ok=True)
     os.makedirs(processed_labels_dir, exist_ok=True)
 
@@ -126,20 +123,31 @@ def preprocess_annotations():
         for image_file in imd:
             if image_file.is_file() and image_file.name not in processed_images:
                 images.append(image_file)
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        executor.map(process_image_file, images)
 
-    for image_file in images:
+def process_image_file(image_file): # this function will be executed in thread
+    try:
+        image_path = os.path.join(data_images_dir, image_file.name)
+        labels_path = os.path.join(data_labels_dir, f'{Path(image_path).stem}.txt')
+        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
+
+        img_ann = ImageLabel(
+            image_path=image_path,
+            image=image,
+            labels_path=labels_path,
+            labels=read_labels(labels_path)
+        )
         try:
-            image_path = os.path.join(data_images_dir, image_file.name)
-            labels_path = os.path.join(data_labels_dir, f'{Path(image_path).stem}.txt')
-            image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
-            process_image(ImageLabel(
-                image_path=image_path,
-                image=image,
-                labels_path=labels_path,
-                labels=read_labels(labels_path)
-            ))
+            results = image_processing(img_ann)
+            for res_ann in results:
+                write_result(res_ann)
         except Exception as e:
-            print(f'Error appeared {e}')
+            print(e)
+        global total_files_processed
+        total_files_processed += 1
+    except Exception as e:
+        print(f'Error appeared in thread for {image_file.name}: {e}')
 
 
 def main():