110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
import json, os, cv2
|
|
from PIL import Image
|
|
import albumentations as A
|
|
|
|
# ⚙️ Configurações
|
|
with open("config.json", "r") as f:
|
|
config = json.load(f)
|
|
MODELO = config["camera"]
|
|
|
|
# Pastas
|
|
dataset_path = os.path.join(MODELO, "dataset", "original", "images")
|
|
masks_path = os.path.join(MODELO, "dataset", "original", "masks")
|
|
aug_img_out = os.path.join(MODELO, "dataset", "augmented", "images")
|
|
aug_msk_out = os.path.join(MODELO, "dataset", "augmented", "masks")
|
|
os.makedirs(aug_img_out, exist_ok=True)
|
|
os.makedirs(aug_msk_out, exist_ok=True)
|
|
|
|
# Pipeline de augmentations
|
|
train_tf = A.Compose([
|
|
A.HorizontalFlip(p=0.5),
|
|
|
|
# Geométricas (aplicam em imagem e máscara)
|
|
A.ShiftScaleRotate(
|
|
shift_limit=0.01,
|
|
scale_limit=0.10,
|
|
rotate_limit=5,
|
|
border_mode=cv2.BORDER_REFLECT_101,
|
|
#value=(255,255,255),
|
|
#mask_value=(255,255,255),
|
|
interpolation=cv2.INTER_LINEAR,
|
|
p=0.3
|
|
),
|
|
|
|
# Fotométricas (somente imagem)
|
|
A.OneOf([
|
|
A.RandomBrightnessContrast(0.2, 0.2, p=1),
|
|
A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1),
|
|
A.RandomGamma(gamma_limit=(90,110), p=1),
|
|
], p=0.7),
|
|
|
|
A.OneOf([
|
|
A.MotionBlur(blur_limit=3, p=1),
|
|
A.GaussianBlur(blur_limit=3, p=1),
|
|
], p=0.20),
|
|
|
|
A.OneOf([
|
|
A.GaussNoise(var_limit=(5.0, 15.0), p=1),
|
|
A.ImageCompression(quality_lower=50, quality_upper=85, p=1),
|
|
], p=0.20),
|
|
|
|
A.RandomShadow(p=0.1),
|
|
A.RandomSunFlare(p=0.1),
|
|
A.ChannelShuffle(p=0.05),
|
|
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.1)
|
|
|
|
# Resize final (img=LINEAR, mask=NEAREST)
|
|
#A.Resize(height=H, width=W, interpolation=cv2.INTER_LINEAR, mask_interpolation=cv2.INTER_NEAREST),
|
|
], additional_targets={'mask':'mask'})
|
|
|
|
def load_rgb(path):
|
|
# cv2 lê BGR → converte pra RGB (Albumentations usa RGB por padrão)
|
|
im = cv2.imread(path, cv2.IMREAD_COLOR)
|
|
if im is None:
|
|
raise FileNotFoundError(path)
|
|
return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
|
|
|
def save_rgb(path, arr_rgb):
|
|
# Salva em RGB mantendo cores corretas
|
|
Image.fromarray(arr_rgb).save(path)
|
|
|
|
def augment_images_and_masks(n_copies=6):
|
|
# Faz pareamento por nome base (sem extensão)
|
|
imgs = sorted([f for f in os.listdir(dataset_path) if os.path.isfile(os.path.join(dataset_path,f))])
|
|
msks = sorted([f for f in os.listdir(masks_path) if os.path.isfile(os.path.join(masks_path,f))])
|
|
|
|
# Mapeia máscaras por nome-base
|
|
msk_map = {os.path.splitext(m)[0]: m for m in msks}
|
|
|
|
total = 0
|
|
for img_file in imgs:
|
|
base, ext = os.path.splitext(img_file)
|
|
if base not in msk_map:
|
|
print(f"[WARN] Máscara não encontrada para {img_file}, pulando.")
|
|
continue
|
|
|
|
img_path = os.path.join(dataset_path, img_file)
|
|
msk_path = os.path.join(masks_path, msk_map[base])
|
|
|
|
# Carrega RGB (máscara como RGB também — mantemos as cores exatas)
|
|
img = load_rgb(img_path)
|
|
msk = load_rgb(msk_path)
|
|
|
|
for i in range(n_copies):
|
|
# Aplica aug; máscara recebe só geométricas
|
|
aug = train_tf(image=img, mask=msk)
|
|
img_aug = aug["image"]
|
|
msk_aug = aug["mask"]
|
|
|
|
# Salva
|
|
out_img = os.path.join(aug_img_out, f"{base}_aug_{i:02d}{ext}")
|
|
out_msk = os.path.join(aug_msk_out, f"{base}_aug_{i:02d}{os.path.splitext(msk_map[base])[1]}")
|
|
save_rgb(out_img, img_aug)
|
|
save_rgb(out_msk, msk_aug)
|
|
total += 1
|
|
|
|
print(f"Augmentation completed! {total} pares gerados.")
|
|
|
|
if __name__ == "__main__":
|
|
augment_images_and_masks(n_copies=5)
|