70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
|
import os
|
||
|
|
import cv2
|
||
|
|
import albumentations as A
|
||
|
|
from albumentations.pytorch import ToTensorV2
|
||
|
|
import numpy as np
|
||
|
|
import random
|
||
|
|
|
||
|
|
# === CONFIG ===
|
||
|
|
PASTA_ORIGINAL = "dataset/original"
|
||
|
|
PASTA_AUGMENTED = "dataset/augmented"
|
||
|
|
NUM_AUGMENTACOES = 3 # quantas imagens gerar por imagem original
|
||
|
|
|
||
|
|
os.makedirs(os.path.join(PASTA_AUGMENTED, "images"), exist_ok=True)
|
||
|
|
os.makedirs(os.path.join(PASTA_AUGMENTED, "labels"), exist_ok=True)
|
||
|
|
|
||
|
|
transform = A.Compose([
|
||
|
|
A.HorizontalFlip(p=0.5),
|
||
|
|
A.VerticalFlip(p=0.1),
|
||
|
|
A.RandomBrightnessContrast(p=0.3),
|
||
|
|
A.Rotate(limit=10, p=0.4),
|
||
|
|
A.RandomScale(scale_limit=0.1, p=0.3),
|
||
|
|
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
|
||
|
|
|
||
|
|
imagens = [f for f in os.listdir(os.path.join(PASTA_ORIGINAL, "images")) if f.endswith(('.jpg', '.jpeg', '.png'))]
|
||
|
|
|
||
|
|
for nome_img in imagens:
|
||
|
|
caminho_img = os.path.join(PASTA_ORIGINAL, "images", nome_img)
|
||
|
|
caminho_lbl = os.path.join(PASTA_ORIGINAL, "labels", nome_img.replace(".jpg", ".txt").replace(".jpeg", ".txt").replace(".png", ".txt"))
|
||
|
|
|
||
|
|
if not os.path.exists(caminho_lbl):
|
||
|
|
print(f"[!] Label ausente: {nome_img}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
# Carrega imagem e label
|
||
|
|
image = cv2.imread(caminho_img)
|
||
|
|
height, width = image.shape[:2]
|
||
|
|
|
||
|
|
with open(caminho_lbl, 'r') as f:
|
||
|
|
linhas = f.readlines()
|
||
|
|
|
||
|
|
bboxes = []
|
||
|
|
class_labels = []
|
||
|
|
for linha in linhas:
|
||
|
|
parts = linha.strip().split()
|
||
|
|
if len(parts) != 5:
|
||
|
|
continue
|
||
|
|
cls, x, y, w, h = map(float, parts)
|
||
|
|
bboxes.append([x, y, w, h])
|
||
|
|
class_labels.append(int(cls))
|
||
|
|
|
||
|
|
for i in range(NUM_AUGMENTACOES):
|
||
|
|
augmented = transform(image=image, bboxes=bboxes, class_labels=class_labels)
|
||
|
|
img_aug = augmented['image']
|
||
|
|
bboxes_aug = augmented['bboxes']
|
||
|
|
labels_aug = augmented['class_labels']
|
||
|
|
|
||
|
|
nome_base = os.path.splitext(nome_img)[0]
|
||
|
|
nome_img_out = f"{nome_base}_aug{i}.jpg"
|
||
|
|
nome_lbl_out = f"{nome_base}_aug{i}.txt"
|
||
|
|
|
||
|
|
cv2.imwrite(os.path.join(PASTA_AUGMENTED, "images", nome_img_out), img_aug)
|
||
|
|
|
||
|
|
with open(os.path.join(PASTA_AUGMENTED, "labels", nome_lbl_out), 'w') as f:
|
||
|
|
for cls, bbox in zip(labels_aug, bboxes_aug):
|
||
|
|
x, y, w, h = bbox
|
||
|
|
f.write(f"{cls} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n")
|
||
|
|
|
||
|
|
print(f"[+] Augmentado: {nome_img_out}")
|
||
|
|
|
||
|
|
print("\n✅ Augmentation finalizado!")
|