| import os |
| import cv2 |
| import numpy as np |
| from PIL import Image |
| import albumentations as A |
| import scripts.config as config |
|
|
| input_img_dir = config.images |
| out_img_dir = config.images |
|
|
| input_mask_dir = config.masks |
| out_mask_dir = config.masks |
|
|
| print('Starting...') |
|
|
| os.makedirs(out_img_dir, exist_ok=True) |
| os.makedirs(out_mask_dir, exist_ok=True) |
|
|
| transformations = [ |
| ("flip", A.HorizontalFlip(p=1.0)), |
| ("rot15", A.Rotate(limit=15, p=1.0)), |
| ("contrast", A.RandomBrightnessContrast(p=1.0)), |
| ] |
|
|
| for fname in os.listdir(input_img_dir): |
| if not fname.endswith(".jpg"): continue |
| base = fname.replace(".jpg", "") |
| |
| img_path = os.path.join(input_img_dir, fname) |
| mask_path = os.path.join(input_mask_dir, base + "_mask.png") |
|
|
| img = np.array(Image.open(img_path).convert("L")) |
| mask = np.array(Image.open(mask_path).convert("L")) |
| mask = (mask > 127).astype('uint8') |
|
|
| for name, tf in transformations: |
| aug = A.Compose([tf]) |
| augmented = aug(image=img, mask=mask) |
| img_aug = augmented['image'] |
| mask_aug = augmented['mask'] |
|
|
| img_out = os.path.join(out_img_dir, f"{base}_aug_{name}.jpg") |
| mask_out = os.path.join(out_mask_dir, f"{base}_aug_{name}_mask.png") |
|
|
| cv2.imwrite(img_out, img_aug) |
| cv2.imwrite(mask_out, mask_aug * 255) |
| |
| print('Completed...') |
|
|
