| """Data augmentation pipeline using albumentations. |
| This script augments images and correspondingly transforms YOLO bboxes. |
| """ |
| import os |
| from glob import glob |
| import cv2 |
| import albumentations as A |
| from pathlib import Path |
| import argparse |
|
|
| AUG = A.Compose([ |
| A.HorizontalFlip(p=0.5), |
| A.RandomRotate90(p=0.5), |
| A.RandomBrightnessContrast(p=0.5), |
| A.GaussianBlur(p=0.3), |
| A.GaussNoise(p=0.3), |
| ], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids'])) |
|
|
|
|
| def load_label_txt(txt_path): |
| bboxes = [] |
| labels = [] |
| if not os.path.exists(txt_path): |
| return bboxes, labels |
| with open(txt_path,'r') as f: |
| for line in f: |
| vals = line.strip().split() |
| if not vals: |
| continue |
| cls = int(vals[0]) |
| bbox = list(map(float, vals[1:5])) |
| bboxes.append(bbox) |
| labels.append(cls) |
| return bboxes, labels |
|
|
|
|
| def save_label_txt(txt_path, bboxes, labels): |
| os.makedirs(os.path.dirname(txt_path), exist_ok=True) |
| with open(txt_path,'w') as f: |
| for cls,b in zip(labels,bboxes): |
| f.write(f"{cls} {b[0]:.6f} {b[1]:.6f} {b[2]:.6f} {b[3]:.6f}\n") |
|
|
|
|
| def augment_image(img_path, label_path, out_img_path, out_lbl_path, n=3): |
| img = cv2.imread(img_path) |
| h,w = img.shape[:2] |
| bboxes, labels = load_label_txt(label_path) |
| for i in range(n): |
| try: |
| augmented = AUG(image=img, bboxes=bboxes, category_ids=labels) |
| except Exception: |
| continue |
| aug_img = augmented['image'] |
| aug_bboxes = augmented['bboxes'] |
| save_img_p = out_img_path.replace('{i}',str(i)) |
| save_lbl_p = out_lbl_path.replace('{i}',str(i)) |
| cv2.imwrite(save_img_p, aug_img) |
| save_label_txt(save_lbl_p, aug_bboxes, augmented['category_ids']) |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--src', default='dataset/images/train') |
| parser.add_argument('--labels', default='dataset/labels/train') |
| parser.add_argument('--out', default='dataset_aug') |
| parser.add_argument('--n', type=int, default=3) |
| args = parser.parse_args() |
|
|
| img_files = glob(os.path.join(args.src,'*.jpg')) + glob(os.path.join(args.src,'*.png')) |
| for img_path in img_files: |
| stem = Path(img_path).stem |
| lbl_path = os.path.join(args.labels, stem + '.txt') |
| out_img = os.path.join(args.out, 'images', stem + '_aug_{i}.jpg') |
| out_lbl = os.path.join(args.out, 'labels', stem + '_aug_{i}.txt') |
| augment_image(img_path, lbl_path, out_img, out_lbl, n=args.n) |
|
|
