mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 09:16:36 +00:00
GPU augmentation. try3. Failed.
This commit is contained in:
zxsanny
+233
-135
@@ -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 nvidia.dali.fn as fn
|
||||
import cv2
|
||||
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 (
|
||||
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)
|
||||
from dto.imageLabel import ImageLabel
|
||||
|
||||
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
|
||||
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 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)
|
||||
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()
|
||||
arr = []
|
||||
for row in rows:
|
||||
str_coordinates = row.split(' ')
|
||||
class_num = str_coordinates.pop(0)
|
||||
coordinates = [float(n.replace(',', '.')) for n in str_coordinates]
|
||||
coordinates.append(class_num)
|
||||
arr.append(coordinates)
|
||||
return arr
|
||||
for row in f.readlines():
|
||||
str_coordinates = row.strip().split(' ')
|
||||
class_num = int(str_coordinates[0])
|
||||
coordinates = [float(n) for n in str_coordinates[1:]] # x_center, y_center, width, height (normalized YOLO)
|
||||
labels.append([*coordinates, class_num])
|
||||
return labels
|
||||
|
||||
def _get_image_label_pairs(self):
|
||||
processed_images = set(f.name for f in os.scandir(processed_images_dir))
|
||||
pairs = []
|
||||
def yolo_to_xyxy(bboxes_yolo, image_width, image_height):
|
||||
bboxes_xyxy = []
|
||||
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):
|
||||
if image_file.is_file() and image_file.name not in processed_images:
|
||||
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')
|
||||
def xyxy_to_yolo(bboxes_xyxy, image_width, image_height):
|
||||
bboxes_yolo = []
|
||||
for bbox in bboxes_xyxy:
|
||||
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):
|
||||
pairs.append((image_path, labels_path))
|
||||
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]
|
||||
|
||||
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):
|
||||
@dali.pipeline_def(batch_size=self.batch_size, num_threads=32, device_id=0)
|
||||
def augmentation_pipeline():
|
||||
# Read images
|
||||
jpegs, _ = fn.file_reader(file_root=data_images_dir, file_list=file_paths, random_shuffle=False)
|
||||
def process_image_and_labels(image, labels_yolo, image_filename, geometric_pipeline, intensity_pipeline, device):
|
||||
image = image.float() / 255.0
|
||||
original_height, original_width = RESIZE_SIZE[0], RESIZE_SIZE[1] # Use fixed resize size (Height, Width)
|
||||
|
||||
# Decode images
|
||||
images = fn.decoders.image(jpegs, device='mixed')
|
||||
processed_image_labels = []
|
||||
|
||||
# Random augmentations with GPU acceleration
|
||||
augmented_images = []
|
||||
for _ in range(NUM_AUGMENTATIONS):
|
||||
aug_image = fn.random_resized_crop(
|
||||
images,
|
||||
device='gpu',
|
||||
size=(images.shape[1], images.shape[2]),
|
||||
random_area=(0.8, 1.0)
|
||||
)
|
||||
# 1. Original image and labels
|
||||
current_labels_yolo_corrected = correct_bboxes(labels_yolo)
|
||||
processed_image_labels.append(ImageLabel(
|
||||
image=KU.tensor_to_image(image.byte()), # Convert back to numpy uint8 for saving
|
||||
labels=current_labels_yolo_corrected,
|
||||
image_path=os.path.join(processed_images_dir, image_filename),
|
||||
labels_path=os.path.join(processed_labels_dir, Path(image_filename).stem + '.txt')
|
||||
))
|
||||
|
||||
# Apply multiple random augmentations
|
||||
aug_image = fn.flip(aug_image, horizontal=fn.random.coin_flip())
|
||||
aug_image = fn.brightness_contrast(
|
||||
aug_image,
|
||||
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
|
||||
)
|
||||
# 2-8. Augmented images
|
||||
for i in range(num_augmented_images):
|
||||
img_batch = image.unsqueeze(0).to(device) # BCHW
|
||||
bboxes_xyxy = yolo_to_xyxy(labels_yolo, original_width, original_height).unsqueeze(0).to(device) # B N 5
|
||||
|
||||
# Add noise and color jittering
|
||||
aug_image = fn.noise.gaussian(aug_image, mean=0, stddev=fn.random.uniform(range=(0, 0.1)))
|
||||
aug_image = fn.hsv(
|
||||
aug_image,
|
||||
hue=fn.random.uniform(range=(-8, 8)),
|
||||
saturation=fn.random.uniform(range=(-8, 8)),
|
||||
value=fn.random.uniform(range=(-8, 8))
|
||||
)
|
||||
augmented_batch = geometric_pipeline(img_batch, params={"bbox": bboxes_xyxy})
|
||||
geo_augmented_image = augmented_batch["input"]
|
||||
geo_augmented_bboxes_xyxy = augmented_batch["bbox"]
|
||||
|
||||
augmented_images.append(aug_image)
|
||||
intensity_augmented_image = intensity_pipeline(geo_augmented_image)
|
||||
|
||||
# Also include original image
|
||||
augmented_images.append(images)
|
||||
# Convert back to CPU and numpy
|
||||
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):
|
||||
image_label_pairs = self._get_image_label_pairs()
|
||||
processed_image_labels.append(ImageLabel(
|
||||
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
|
||||
file_list_path = os.path.join(processed_images_dir, 'file_list.txt')
|
||||
with open(file_list_path, 'w') as f:
|
||||
for img_path, _ in image_label_pairs:
|
||||
f.write(f'{img_path}\n')
|
||||
def write_result(img_ann: ImageLabel):
|
||||
cv2.imwrite(img_ann.image_path, cv2.cvtColor(img_ann.image, cv2.COLOR_RGB2BGR)) # Save as BGR
|
||||
print(f'{img_ann.image_path} written')
|
||||
|
||||
# Create DALI pipeline
|
||||
pipeline = self.create_dali_pipeline(file_list_path)
|
||||
pipeline.build()
|
||||
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()
|
||||
print(f'{img_ann.labels_path} written')
|
||||
|
||||
# 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()
|
||||
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
|
||||
|
||||
# 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)
|
||||
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:
|
||||
loader = DataLoader()
|
||||
loader.process_batch()
|
||||
print('All processed, waiting for 2 minutes...')
|
||||
time.sleep(120)
|
||||
global total_files_processed
|
||||
total_files_processed = 0
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
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__':
|
||||
|
||||
+4
-5
@@ -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,
|
||||
annotation_classes, checkpoint_file, checkpoint_date_format)
|
||||
from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir)
|
||||
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__':
|
||||
|
||||
@@ -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
|
||||
+1
-2
@@ -12,5 +12,4 @@ cryptography
|
||||
numpy
|
||||
requests
|
||||
pyyaml
|
||||
boto3
|
||||
nvidia-dali-cuda120
|
||||
boto3
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user