File size: 1,386 Bytes
7b615ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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 #+ "/images_aug"
input_mask_dir = config.masks
out_mask_dir = config.masks #+ "/masks_aug"
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...')
|