GPU augmentation. try3. Failed.

preprocessing-cuda.py

@@ -1,170 +1,268 @@
 import os
 import time
-import cv2
+from datetime import datetime, timedelta
 from pathlib import Path
-import nvidia.dali as dali
+import cv2
-import nvidia.dali.fn as fn
+import numpy as np
+import torch
+import kornia.augmentation as K
+import kornia.utils as KU
+from torch.utils.data import Dataset, DataLoader
+import concurrent.futures
 
-from constants import (
+from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir)
-    data_images_dir,
+from dto.imageLabel import ImageLabel
-    data_labels_dir,
-    processed_images_dir,
-    processed_labels_dir
-)
 
-NUM_AUGMENTATIONS = 7
+# Configurable parameters
+num_augmented_images = 7
+augmentation_probability = 0.5 # general probability for augmentations, can be adjusted per augmentation
+RESIZE_SIZE = (1080, 1920) # Resize images to Full HD 1920x1080 (height, width)
+processed_images_dir = processed_images_dir + '_cuda'
+processed_labels_dir = processed_labels_dir + '_cuda'
 
-
+# Ensure directories exist
-class DataLoader:
-    def __init__(self, batch_size=32):
-        self.batch_size = batch_size
 os.makedirs(processed_images_dir, exist_ok=True)
 os.makedirs(processed_labels_dir, exist_ok=True)
 
+# Custom Augmentations
+class RandomFog(K.AugmentationBase2D):
+    def __init__(self, fog_coef_range=(0, 0.3), p=augmentation_probability, same_on_batch=True):
+        super().__init__(p=p, same_on_batch=same_on_batch)
+        self.fog_coef_range = fog_coef_range
+
+    def compute_transformation(self, input_shape: torch.Size, params: dict) -> dict:
+        return {"fog_factor": torch.rand(input_shape[0], device=self.device) * (self.fog_coef_range[1] - self.fog_coef_range[0]) + self.fog_coef_range[0]}
+
+    def apply_transform(self, input: torch.Tensor, params: dict, transform: dict) -> torch.Tensor:
+        fog_factor = transform['fog_factor'].view(-1, 1, 1, 1)
+        return input * (1.0 - fog_factor) + fog_factor
+
+class RandomShadow(K.AugmentationBase2D):
+    def __init__(self, shadow_factor_range=(0.2, 0.8), p=augmentation_probability, same_on_batch=True):
+        super().__init__(p=p, same_on_batch=same_on_batch)
+        self.shadow_factor_range = shadow_factor_range
+
+    def compute_transformation(self, input_shape: torch.Size, params: dict) -> dict:
+        batch_size, _, height, width = input_shape
+        x1 = torch.randint(0, width, (batch_size,), device=self.device)
+        y1 = torch.randint(0, height, (batch_size,), device=self.device)
+        x2 = torch.randint(x1, width, (batch_size,), device=self.device)
+        y2 = torch.randint(y1, height, (batch_size,), device=self.device)
+        shadow_factor = torch.rand(batch_size, device=self.device) * (self.shadow_factor_range[1] - self.shadow_factor_range[0]) + self.shadow_factor_range[0]
+        return {"x1": x1, "y1": y1, "x2": x2, "y2": y2, "shadow_factor": shadow_factor}
+
+    def apply_transform(self, input: torch.Tensor, params: dict, transform: dict) -> torch.Tensor:
+        batch_size, _, height, width = input.size()
+        mask = torch.zeros_like(input, device=self.device)
+        for b in range(batch_size):
+            mask[b, :, transform['y1'][b]:transform['y2'][b], transform['x1'][b]:transform['x2'][b]] = 1
+        shadow_factor = transform['shadow_factor'].view(-1, 1, 1, 1)
+        return input * (1.0 - mask) + input * mask * shadow_factor
+
+
+class ImageDataset(Dataset):
+    def __init__(self, images_dir, labels_dir):
+        self.images_dir = images_dir
+        self.labels_dir = labels_dir
+        self.image_filenames = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]
+        self.resize = K.Resize(RESIZE_SIZE) # Add resize transform here
+
+    def __len__(self):
+        return len(self.image_filenames)
+
+    def __getitem__(self, idx):
+        image_filename = self.image_filenames[idx]
+        image_path = os.path.join(self.images_dir, image_filename)
+        label_path = os.path.join(self.labels_dir, Path(image_filename).stem + '.txt')
+
+        image_np = cv2.imread(image_path)
+        if image_np is None:
+            raise FileNotFoundError(f"Error reading image: {image_path}")
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) # Convert to RGB for Kornia
+        image = KU.image_to_tensor(image_np, keepdim=False).float()  # HWC -> CHW, and to tensor, convert to float here!
+        image = self.resize(image) # Resize image here to fixed size
+
+        print(f"Image shape after resize (index {idx}, filename {image_filename}): {image.shape}") # DEBUG PRINT
+
+        labels = []
+        if os.path.exists(label_path):
+            labels = self._read_labels(label_path)
+        return image, labels, image_filename
+
     def _read_labels(self, labels_path):
+        labels = []
         with open(labels_path, 'r') as f:
-            rows = f.readlines()
+            for row in f.readlines():
-            arr = []
+                str_coordinates = row.strip().split(' ')
-            for row in rows:
+                class_num = int(str_coordinates[0])
-                str_coordinates = row.split(' ')
+                coordinates = [float(n) for n in str_coordinates[1:]] # x_center, y_center, width, height (normalized YOLO)
-                class_num = str_coordinates.pop(0)
+                labels.append([*coordinates, class_num])
-                coordinates = [float(n.replace(',', '.')) for n in str_coordinates]
+        return labels
-                coordinates.append(class_num)
-                arr.append(coordinates)
-            return arr
 
-    def _get_image_label_pairs(self):
+def yolo_to_xyxy(bboxes_yolo, image_width, image_height):
-        processed_images = set(f.name for f in os.scandir(processed_images_dir))
+    bboxes_xyxy = []
-        pairs = []
+    for bbox in bboxes_yolo:
+        x_center, y_center, w, h, class_id = bbox
+        x_min = int((x_center - w / 2) * image_width)
+        y_min = int((y_center - h / 2) * image_height)
+        x_max = int((x_center + w / 2) * image_width)
+        y_max = int((y_center + h / 2) * image_height)
+        bboxes_xyxy.append([x_min, y_min, x_max, y_max, class_id])
+    return torch.tensor(bboxes_xyxy) if bboxes_xyxy else torch.empty((0, 5))
 
-        for image_file in os.scandir(data_images_dir):
+def xyxy_to_yolo(bboxes_xyxy, image_width, image_height):
-            if image_file.is_file() and image_file.name not in processed_images:
+    bboxes_yolo = []
-                image_path = os.path.join(data_images_dir, image_file.name)
+    for bbox in bboxes_xyxy:
-                labels_path = os.path.join(data_labels_dir, f'{Path(image_path).stem}.txt')
+        x_min, y_min, x_max, y_max, class_id = bbox
+        x_center = ((x_min + x_max) / 2) / image_width
+        y_center = ((y_min + y_max) / 2) / image_height
+        w = (x_max - x_min) / image_width
+        h = (y_max - y_min) / image_height
+        bboxes_yolo.append([x_center, y_center, w, h, int(class_id)])
+    return bboxes_yolo
 
-                if os.path.exists(labels_path):
+def correct_bboxes(labels):
-                    pairs.append((image_path, labels_path))
+    margin = 0.0005
+    min_size = 0.01
+    res = []
+    for bboxes in labels:
+        x = bboxes[0]
+        y = bboxes[1]
+        half_width = 0.5*bboxes[2]
+        half_height = 0.5*bboxes[3]
 
-        return pairs
+        w_diff = min( (1 - margin) - (x + half_width), (x - half_width) - margin, 0 )
+        w = bboxes[2] + 2*w_diff
+        if w < min_size:
+            continue
+        h_diff = min( (1 - margin) - (y + half_height), ((y - half_height) - margin), 0)
+        h = bboxes[3] + 2 * h_diff
+        if h < min_size:
+            continue
+        res.append([x, y, w, h, bboxes[4]])
+    return res
 
-    def create_dali_pipeline(self, file_paths):
+def process_image_and_labels(image, labels_yolo, image_filename, geometric_pipeline, intensity_pipeline, device):
-        @dali.pipeline_def(batch_size=self.batch_size, num_threads=32, device_id=0)
+    image = image.float() / 255.0
-        def augmentation_pipeline():
+    original_height, original_width = RESIZE_SIZE[0], RESIZE_SIZE[1] # Use fixed resize size (Height, Width)
-            # Read images
-            jpegs, _ = fn.file_reader(file_root=data_images_dir, file_list=file_paths, random_shuffle=False)
 
-            # Decode images
+    processed_image_labels = []
-            images = fn.decoders.image(jpegs, device='mixed')
 
-            # Random augmentations with GPU acceleration
+    # 1. Original image and labels
-            augmented_images = []
+    current_labels_yolo_corrected = correct_bboxes(labels_yolo)
-            for _ in range(NUM_AUGMENTATIONS):
+    processed_image_labels.append(ImageLabel(
-                aug_image = fn.random_resized_crop(
+        image=KU.tensor_to_image(image.byte()), # Convert back to numpy uint8 for saving
-                    images,
+        labels=current_labels_yolo_corrected,
-                    device='gpu',
+        image_path=os.path.join(processed_images_dir, image_filename),
-                    size=(images.shape[1], images.shape[2]),
+        labels_path=os.path.join(processed_labels_dir, Path(image_filename).stem + '.txt')
-                    random_area=(0.8, 1.0)
+    ))
-                )
 
-                # Apply multiple random augmentations
+    # 2-8. Augmented images
-                aug_image = fn.flip(aug_image, horizontal=fn.random.coin_flip())
+    for i in range(num_augmented_images):
-                aug_image = fn.brightness_contrast(
+        img_batch = image.unsqueeze(0).to(device) # BCHW
-                    aug_image,
+        bboxes_xyxy = yolo_to_xyxy(labels_yolo, original_width, original_height).unsqueeze(0).to(device) # B N 5
-                    brightness=fn.random.uniform(range=(-0.05, 0.05)),
-                    contrast=fn.random.uniform(range=(-0.05, 0.05))
-                )
-                aug_image = fn.rotate(
-                    aug_image,
-                    angle=fn.random.uniform(range=(-25, 25)),
-                    fill_value=0
-                )
 
-                # Add noise and color jittering
+        augmented_batch = geometric_pipeline(img_batch, params={"bbox": bboxes_xyxy})
-                aug_image = fn.noise.gaussian(aug_image, mean=0, stddev=fn.random.uniform(range=(0, 0.1)))
+        geo_augmented_image = augmented_batch["input"]
-                aug_image = fn.hsv(
+        geo_augmented_bboxes_xyxy = augmented_batch["bbox"]
-                    aug_image,
-                    hue=fn.random.uniform(range=(-8, 8)),
-                    saturation=fn.random.uniform(range=(-8, 8)),
-                    value=fn.random.uniform(range=(-8, 8))
-                )
 
-                augmented_images.append(aug_image)
+        intensity_augmented_image = intensity_pipeline(geo_augmented_image)
 
-            # Also include original image
+        # Convert back to CPU and numpy
-            augmented_images.append(images)
+        augmented_image_np = KU.tensor_to_image((intensity_augmented_image.squeeze(0).cpu() * 255.0).byte())
+        augmented_bboxes_xyxy_cpu = geo_augmented_bboxes_xyxy.squeeze(0).cpu()
 
-            return tuple(augmented_images)
+        augmented_bboxes_yolo = xyxy_to_yolo(augmented_bboxes_xyxy_cpu, original_width, original_height)
+        augmented_bboxes_yolo_corrected = correct_bboxes(augmented_bboxes_yolo)
 
-        return augmentation_pipeline()
 
-    def process_batch(self):
+        processed_image_labels.append(ImageLabel(
-        image_label_pairs = self._get_image_label_pairs()
+            image=augmented_image_np,
+            labels=augmented_bboxes_yolo_corrected,
+            image_path=os.path.join(processed_images_dir, f'{Path(image_filename).stem}_{i + 1}{Path(image_filename).suffix}'),
+            labels_path=os.path.join(processed_labels_dir, f'{Path(image_filename).stem}_{i + 1}.txt')
+        ))
+    return processed_image_labels
 
-        # Create file list for DALI
+def write_result(img_ann: ImageLabel):
-        file_list_path = os.path.join(processed_images_dir, 'file_list.txt')
+    cv2.imwrite(img_ann.image_path, cv2.cvtColor(img_ann.image, cv2.COLOR_RGB2BGR)) # Save as BGR
-        with open(file_list_path, 'w') as f:
+    print(f'{img_ann.image_path} written')
-            for img_path, _ in image_label_pairs:
-                f.write(f'{img_path}\n')
 
-        # Create DALI pipeline
+    with open(img_ann.labels_path, 'w') as f:
-        pipeline = self.create_dali_pipeline(file_list_path)
+        lines = [f'{ann[4]} {round(ann[0], 5)} {round(ann[1], 5)} {round(ann[2], 5)} {round(ann[3], 5)}\n' for ann in
-        pipeline.build()
+                 img_ann.labels]
-
-        # Process images
-        for batch_idx in range(0, len(image_label_pairs), self.batch_size):
-            batch_pairs = image_label_pairs[batch_idx:batch_idx + self.batch_size]
-
-            pipeline.run()
-
-            # Get augmented images
-            for img_idx, (orig_img_path, orig_labels_path) in enumerate(batch_pairs):
-                # Read original labels
-                orig_labels = self._read_labels(orig_labels_path)
-
-                # Write original image and labels
-                self._write_image_and_labels(
-                    pipeline.output[NUM_AUGMENTATIONS][img_idx],
-                    orig_img_path,
-                    orig_labels,
-                    is_original=True
-                )
-
-                # Write augmented images
-                for aug_idx in range(NUM_AUGMENTATIONS):
-                    self._write_image_and_labels(
-                        pipeline.output[aug_idx][img_idx],
-                        orig_img_path,
-                        orig_labels,
-                        aug_idx=aug_idx
-                    )
-
-    def _write_image_and_labels(self, image, orig_img_path, labels, is_original=False, aug_idx=None):
-        path = Path(orig_img_path)
-
-        if is_original:
-            img_name = path.name
-            label_name = f'{path.stem}.txt'
-        else:
-            img_name = f'{path.stem}_{aug_idx + 1}{path.suffix}'
-            label_name = f'{path.stem}_{aug_idx + 1}.txt'
-
-        # Write image
-        img_path = os.path.join(processed_images_dir, img_name)
-        cv2.imencode('.jpg', image.asnumpy())[1].tofile(img_path)
-
-        # Write labels
-        label_path = os.path.join(processed_labels_dir, label_name)
-        with open(label_path, 'w') as f:
-            lines = [f'{ann[4]} {round(ann[0], 5)} {round(ann[1], 5)} {round(ann[2], 5)} {round(ann[3], 5)}\n' for ann in labels]
         f.writelines(lines)
+        f.close()
+    print(f'{img_ann.labels_path} written')
+
+
+def process_batch_wrapper(batch_data, geometric_pipeline, intensity_pipeline, device):
+    processed_batch_image_labels = []
+    for image, labels_yolo, image_filename in batch_data:
+        results = process_image_and_labels(image, labels_yolo, image_filename, geometric_pipeline, intensity_pipeline, device)
+        processed_batch_image_labels.extend(results)
+    return processed_batch_image_labels
+
+def save_batch_results(batch_image_labels):
+    global total_files_processed
+    for img_ann in batch_image_labels:
+        write_result(img_ann)
+        total_files_processed += 1
+        print(f"Total processed images: {total_files_processed}")
 
 
 def main():
-    while True:
+    global total_files_processed
-        loader = DataLoader()
+    total_files_processed = 0
-        loader.process_batch()
+
-        print('All processed, waiting for 2 minutes...')
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        time.sleep(120)
+    print(f"Using device: {device}")
+
+    geometric_pipeline = K.AugmentationSequential(
+        K.RandomHorizontalFlip(p=0.5),
+        K.RandomAffine(degrees=25, translate=(0.1, 0.1), scale=(0.8, 1.2), p=0.5),
+        data_keys=["input", "bbox"],
+        same_on_batch=False
+    ).to(device)
+
+    intensity_pipeline = K.AugmentationSequential(
+        K.ColorJitter(brightness=0.1, contrast=0.07, saturation=0.1, hue=0.1, p=0.5),
+        RandomFog(p=0.2),
+        RandomShadow(p=0.3),
+        K.RandomMotionBlur(kernel_size=3, angle=35., direction=0.5, p=0.3),
+        data_keys=["input"],
+        same_on_batch=False
+    ).to(device)
+
+
+    dataset = ImageDataset(data_images_dir, data_labels_dir)
+    dataloader = DataLoader(dataset, batch_size=32, shuffle=False, num_workers=os.cpu_count()) # Adjust batch_size as needed
+
+    processed_images_set = set(os.listdir(processed_images_dir))
+    images_to_process_indices = [i for i, filename in enumerate(dataset.image_filenames) if filename not in processed_images_set]
+    dataloader_filtered = torch.utils.data.Subset(dataset, images_to_process_indices)
+    filtered_dataloader = DataLoader(dataloader_filtered, batch_size=32, shuffle=False, num_workers=os.cpu_count())
+
+
+    start_time = time.time()
+
+    try:
+        for batch_data in filtered_dataloader:
+            batch_image = batch_data[0]
+            batch_labels = batch_data[1]
+            batch_filenames = batch_data[2]
+
+            batch_processed_image_labels = process_batch_wrapper(list(zip(batch_image, batch_labels, batch_filenames)), geometric_pipeline, intensity_pipeline, device)
+            save_batch_results(batch_processed_image_labels)
+    except Exception as e:
+        print(e)
+
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    print(f"Total processing time: {elapsed_time:.2f} seconds")
+    images_per_hour = (total_files_processed / elapsed_time) * 3600
+    print(f"Processed images per hour: {images_per_hour:.2f}")
+    print("Augmentation process completed.")
 
 
 if __name__ == '__main__':

preprocessing.py

@@ -1,14 +1,13 @@
+import concurrent.futures
 import os.path
 import time
-from datetime import datetime, timedelta
 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,
+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
@@ -154,7 +153,7 @@ def main():
     while True:
         preprocess_annotations()
         print('All processed, waiting for 2 minutes...')
-        time.sleep(120)
+        time.sleep(300)
 
 
 if __name__ == '__main__':

prompts/DALI_initial_prompt.txt

@@ -1,172 +0,0 @@
-I have a code for augmenting photos for dataset, I'm using albumentations. The problem is - I have 38k photos and more in a future, and albumentations works on CPU. It's working very slow, around 1800/ hour. I want to use GPU approach for augmentation task, DALI, since it's an original Nvidia implementation.
-Note, it should create 7 augmented images  + original one.
-
-Here is a code I'm using now:
-
-import os.path
-import time
-from datetime import datetime, timedelta
-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
-    min_size = 0.01
-    res = []
-    for bboxes in labels:
-        x = bboxes[0]
-        y = bboxes[1]
-        half_width = 0.5*bboxes[2]
-        half_height = 0.5*bboxes[3]
-
-        # calc how much bboxes are outside borders ( +small margin ).
-        # value should be negative. If it's positive, then put 0, as no correction
-        w_diff = min( (1 - margin) - (x + half_width), (x - half_width) - margin, 0 )
-        w = bboxes[2] + 2*w_diff
-        if w < min_size:
-            continue
-        h_diff = min( (1 - margin) - (y + half_height), ((y - half_height) - margin), 0)
-        h = bboxes[3] + 2 * h_diff
-        if h < min_size:
-            continue
-        res.append([x, y, w, h, bboxes[4]])
-    return res
-    pass
-
-
-def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
-    results = []
-    labels = correct_bboxes(img_ann.labels)
-    if len(labels) == 0 and len(img_ann.labels) != 0:
-        print('no labels but was!!!')
-    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)
-            name = f'{path.stem}_{i + 1}'
-            img = ImageLabel(
-                image=res['image'],
-                labels=res['bboxes'],
-                image_path=os.path.join(processed_images_dir, f'{name}{path.suffix}'),
-                labels_path=os.path.join(processed_labels_dir, f'{name}.txt')
-            )
-            results.append(img)
-        except Exception as e:
-            print(f'Error during transformation: {e}')
-    return results
-
-
-def write_result(img_ann: ImageLabel):
-    cv2.imencode('.jpg', img_ann.image)[1].tofile(img_ann.image_path)
-    print(f'{img_ann.image_path} written')
-
-    with open(img_ann.labels_path, 'w') as f:
-        lines = [f'{ann[4]} {round(ann[0], 5)} {round(ann[1], 5)} {round(ann[2], 5)} {round(ann[3], 5)}\n' for ann in
-                 img_ann.labels]
-        f.writelines(lines)
-        f.close()
-    global total_files_processed
-    print(f'{total_files_processed}. {img_ann.labels_path} written')
-
-
-def read_labels(labels_path) -> [[]]:
-    with open(labels_path, 'r') as f:
-        rows = f.readlines()
-        arr = []
-        for row in rows:
-            str_coordinates = row.split(' ')
-            class_num = str_coordinates.pop(0)
-            coordinates = [float(n.replace(',', '.')) for n in str_coordinates]
-            # noinspection PyTypeChecker
-            coordinates.append(class_num)
-            arr.append(coordinates)
-        return arr
-
-
-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)
-
-    processed_images = set(f.name for f in os.scandir(processed_images_dir))
-    images = []
-    with os.scandir(data_images_dir) as imd:
-        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)
-
-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:
-            results = image_processing(img_ann)
-            for res_ann in results:
-                write_result(res_ann)
-        except Exception as 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():
-    while True:
-        preprocess_annotations()
-        print('All processed, waiting for 2 minutes...')
-        time.sleep(120)
-
-
-if __name__ == '__main__':
-    main()
-
-
-please rewrite the whole code to DALI, utilizing GPU.
-I do have 128Gb of RAM, RTX4090 and 32CPU Cores 
-Also note, that for each image I'm making 7 augmented version (each of this versions should be different, cause of random factor of apply one or another augmentation + different random parameters in each augmentation mechanism)
-Make this number 7 configurable in the beginning of the file. Also utilize GPU as match as possible, use batching

requirements.txt

@@ -13,4 +13,3 @@ numpy
 requests
 pyyaml
 boto3
-nvidia-dali-cuda120

train.py

@@ -7,7 +7,6 @@ from datetime import datetime
 from os import path, replace, listdir, makedirs, scandir
 from os.path import abspath
 from pathlib import Path
-from utils import Dotdict
 
 import yaml
 from ultralytics import YOLO
@@ -15,13 +14,14 @@ from ultralytics import YOLO
 import constants
 from azaion_api import ApiCredentials, Api
 from cdn_manager import CDNCredentials, CDNManager
-from security import Security
 from constants import (processed_images_dir,
                        processed_labels_dir,
                        annotation_classes,
                        prefix, date_format,
                        datasets_dir, models_dir,
                        corrupted_images_dir, corrupted_labels_dir, sample_dir)
+from security import Security
+from utils import Dotdict
 
 today_folder = f'{prefix}{datetime.now():{date_format}}'
 today_dataset = path.join(datasets_dir, today_folder)