GPU augmentation. try1

This commit is contained in:
zxsanny
2025-03-05 10:47:13 +02:00
parent b5e5f0b297
commit 5bd519b968
+7 -7
View File
@@ -1,13 +1,9 @@
import os
import time
import numpy as np
import cv2
from pathlib import Path
import concurrent.futures
import nvidia.dali as dali
import nvidia.dali.fn as fn
import nvidia.dali.types as types
from constants import (
data_images_dir,
@@ -16,7 +12,6 @@ from constants import (
processed_labels_dir
)
# Configurable number of augmentations per image
NUM_AUGMENTATIONS = 7
@@ -64,7 +59,12 @@ class DataLoader:
# Random augmentations with GPU acceleration
augmented_images = []
for _ in range(NUM_AUGMENTATIONS):
aug_image = fn.random_resized_crop(images, random_area=(0.8, 1.0))
aug_image = fn.random_resized_crop(
images,
device='gpu',
min_scale=0.8,
max_scale=1.0
)
# Apply multiple random augmentations
aug_image = fn.flip(aug_image, horizontal=fn.random.coin_flip())
@@ -93,7 +93,7 @@ class DataLoader:
# Also include original image
augmented_images.append(images)
return augmented_images
return tuple(augmented_images)
return augmentation_pipeline()