111 lines
4.6 KiB
Python
111 lines
4.6 KiB
Python
|
import os
|
||
|
|
from torchvision import transforms
|
||
|
|
from PIL import ImageOps, Image
|
||
|
|
import torchvision.transforms.functional as TF
|
||
|
|
from torchvision.transforms.functional import to_pil_image
|
||
|
|
import torchvision.transforms as T
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
class ComposeWithSeed(object):
|
||
|
|
def __init__(self, transforms):
|
||
|
|
self.transforms = transforms
|
||
|
|
|
||
|
|
def __call__(self, i, img, mask):
|
||
|
|
transforms.RandomHorizontalFlip(p=0.5)
|
||
|
|
t = self.transforms[i]
|
||
|
|
img = t(img)
|
||
|
|
mask = t(mask)
|
||
|
|
return img, mask
|
||
|
|
|
||
|
|
# Função para aplicar flip horizontal e vertical
|
||
|
|
def flip_image(img):
|
||
|
|
return ImageOps.mirror(ImageOps.flip(img))
|
||
|
|
|
||
|
|
# Função para rotacionar a imagem em ângulos fixos
|
||
|
|
def rotate_image(img, angle):
|
||
|
|
return TF.rotate(img, angle)
|
||
|
|
|
||
|
|
def perspective_image(img, magnitude=0.5):
|
||
|
|
width, height = img.size
|
||
|
|
|
||
|
|
# Pontos de origem
|
||
|
|
points_orig = np.float32([
|
||
|
|
[0, 0],
|
||
|
|
[width, 0],
|
||
|
|
[0, height],
|
||
|
|
[width, height]
|
||
|
|
])
|
||
|
|
|
||
|
|
# Pontos de destino, deslocados com base na magnitude
|
||
|
|
points_dest = np.float32([
|
||
|
|
[int(magnitude * width), int(magnitude * height)],
|
||
|
|
[int((1 - magnitude) * width), 0],
|
||
|
|
[0, int((1 - magnitude) * height)],
|
||
|
|
[width, height]
|
||
|
|
])
|
||
|
|
|
||
|
|
# Calcula a matriz de transformação e aplica a transformação de perspectiva
|
||
|
|
matrix = cv2.getPerspectiveTransform(points_orig, points_dest)
|
||
|
|
img_transformed = cv2.warpPerspective(np.array(img), matrix, (width, height))
|
||
|
|
|
||
|
|
return Image.fromarray(img_transformed)
|
||
|
|
|
||
|
|
# Definindo as transformações
|
||
|
|
transform_list = [
|
||
|
|
T.Lambda(lambda img: flip_image(img)), # Aplica flip nos dois eixos
|
||
|
|
T.Lambda(lambda img: rotate_image(img, 90)), # Rotação de 90 graus
|
||
|
|
T.Lambda(lambda img: rotate_image(img, -90)), # Rotação de -90 graus
|
||
|
|
T.Lambda(lambda img: perspective_image(img, magnitude=0.2)),
|
||
|
|
T.Lambda(lambda img: perspective_image(img, magnitude=0.1)),
|
||
|
|
# Aqui, você pode adicionar outras transformações fixas conforme necessário
|
||
|
|
#T.Resize((256, 256)), # Redimensionamento para o tamanho desejado
|
||
|
|
T.ToTensor(), # Converte as imagens PIL para tensores PyTorch
|
||
|
|
# Normalização pode ser adicionada aqui, se necessário
|
||
|
|
]
|
||
|
|
|
||
|
|
# Agora, definimos a transformação composta com a classe personalizada
|
||
|
|
transform = ComposeWithSeed(transform_list)
|
||
|
|
|
||
|
|
# Caminhos para os diretórios onde suas imagens e máscaras originais estão armazenadas
|
||
|
|
dataset_path = 'dataset/images'
|
||
|
|
masks_path = 'dataset/masks'
|
||
|
|
|
||
|
|
# Caminhos para os diretórios onde as imagens e máscaras aumentadas serão salvas
|
||
|
|
augmented_images_path = 'dataset/augmented_images'
|
||
|
|
augmented_masks_path = 'dataset/augmented_masks'
|
||
|
|
|
||
|
|
# Verifica se os diretórios de destino existem, caso contrário, cria os diretórios
|
||
|
|
os.makedirs(augmented_images_path, exist_ok=True)
|
||
|
|
os.makedirs(augmented_masks_path, exist_ok=True)
|
||
|
|
|
||
|
|
# Função para aplicar a transformação e salvar as imagens e máscaras transformadas
|
||
|
|
def augment_images_and_masks(dataset_path, masks_path, augmented_images_path, augmented_masks_path, transform, num_copies=5):
|
||
|
|
# Lista todos os arquivos nos diretórios do dataset de imagens e máscaras
|
||
|
|
image_files = [f for f in os.listdir(dataset_path) if os.path.isfile(os.path.join(dataset_path, f))]
|
||
|
|
mask_files = [f for f in os.listdir(masks_path) if os.path.isfile(os.path.join(masks_path, f))]
|
||
|
|
|
||
|
|
for image_file, mask_file in zip(image_files, mask_files):
|
||
|
|
image_path = os.path.join(dataset_path, image_file)
|
||
|
|
mask_path = os.path.join(masks_path, mask_file)
|
||
|
|
image = Image.open(image_path).convert('RGB')
|
||
|
|
mask = Image.open(mask_path).convert('RGB')
|
||
|
|
|
||
|
|
for i in range(num_copies):
|
||
|
|
# Aplica a transformação de maneira consistente em ambos, imagem e máscara
|
||
|
|
transformed_image, transformed_mask = transform(i, image, mask)
|
||
|
|
|
||
|
|
# Salva a imagem e a máscara transformadas
|
||
|
|
image_save_path = os.path.join(augmented_images_path, f"{os.path.splitext(image_file)[0]}_aug_{i}{os.path.splitext(image_file)[1]}")
|
||
|
|
mask_save_path = os.path.join(augmented_masks_path, f"{os.path.splitext(mask_file)[0]}_aug_{i}{os.path.splitext(mask_file)[1]}")
|
||
|
|
|
||
|
|
#transformed_image_pil = to_pil_image(transformed_image)
|
||
|
|
transformed_image.save(image_save_path)
|
||
|
|
#transformed_mask_pil = to_pil_image(transformed_mask)
|
||
|
|
transformed_mask.save(mask_save_path)
|
||
|
|
|
||
|
|
# Chama a função para iniciar o processo de aumento de dados
|
||
|
|
augment_images_and_masks(dataset_path, masks_path, augmented_images_path, augmented_masks_path, transform, num_copies=5)
|
||
|
|
|
||
|
|
print("Augmentation completed!")
|