pinthoz's picture
First commit
097f176 verified
import os
import cv2
import albumentations as A
# Função para ajustar as anotações após transformação
def adjust_annotations(annotations, original_size, resized_size, padded_size):
orig_h, orig_w = original_size
resized_h, resized_w = resized_size
padded_h, padded_w = padded_size
scale_x = resized_w / orig_w # Escala para largura
scale_y = resized_h / orig_h # Escala para altura
pad_x = (padded_w - resized_w) / 2 # Padding horizontal
pad_y = (padded_h - resized_h) / 2 # Padding vertical
adjusted_annotations = []
for ann in annotations:
cls, x_c, y_c, w, h = ann
# Ajustar coordenadas normalizadas para o redimensionamento
x_c = x_c * orig_w * scale_x
y_c = y_c * orig_h * scale_y
w = w * orig_w * scale_x
h = h * orig_h * scale_y
# Ajustar para o padding e normalizar para o novo tamanho
x_c = (x_c + pad_x) / padded_w
y_c = (y_c + pad_y) / padded_h
w = w / padded_w
h = h / padded_h
# Adicionar anotação ajustada
adjusted_annotations.append((cls, x_c, y_c, w, h))
return adjusted_annotations
# Função para ler as anotações de um arquivo
def read_annotations(file_path):
annotations = []
with open(file_path, "r") as f:
for line in f:
parts = line.strip().split()
cls = int(parts[0]) # Classe
x_c, y_c, w, h = map(float, parts[1:])
annotations.append((cls, x_c, y_c, w, h))
return annotations
# Configuração de transformações
resize_height, resize_width = 128, 128
padded_height, padded_width = 256, 256
augmentation = A.Compose([
A.Resize(height=resize_height, width=resize_width),
A.PadIfNeeded(min_height=padded_height, min_width=padded_width, border_mode=cv2.BORDER_CONSTANT, value=(0, 0, 0)),
])
# Caminhos das pastas
images_folder = r"C:\Users\anoca\Documents\GitHub\analogic-watch-detector\dataset\images\train"
labels_folder = r"C:\Users\anoca\Documents\GitHub\analogic-watch-detector\dataset\labels\train"
output_images_folder = r"C:\Users\anoca\Documents\GitHub\analogic-watch-detector\dataset\images\train"
output_labels_folder = r"C:\Users\anoca\Documents\GitHub\analogic-watch-detector\dataset\labels\train"
# Certificar-se de que as pastas de saída existem
os.makedirs(output_images_folder, exist_ok=True)
os.makedirs(output_labels_folder, exist_ok=True)
# Processar todos os arquivos na pasta
for image_filename in os.listdir(images_folder):
# Pular arquivos com "rotated" no nome
if "rotated" in image_filename:
print(f"Pulado: {image_filename} (contém 'rotated').")
continue
if image_filename.endswith((".jpg", ".png", ".jpeg")): # Verifica formatos de imagem
image_path = os.path.join(images_folder, image_filename)
label_path = os.path.join(labels_folder, image_filename.replace(".jpg", ".txt").replace(".png", ".txt").replace(".jpeg", ".txt"))
# Verificar se o arquivo de anotações correspondente existe
if not os.path.exists(label_path):
print(f"Anotação não encontrada para {image_filename}, pulando.")
continue
# Carregar a imagem e as anotações
image = cv2.imread(image_path)
original_size = image.shape[:2] # (altura, largura)
annotations = read_annotations(label_path)
# Aplicar transformações
augmented = augmentation(image=image)
augmented_image = augmented["image"]
# Ajustar as anotações
new_annotations = adjust_annotations(
annotations,
original_size=(original_size[0], original_size[1]),
resized_size=(resize_height, resize_width),
padded_size=(padded_height, padded_width)
)
# Criar novo nome para a imagem e as anotações
base_name, ext = os.path.splitext(image_filename)
new_image_name = f"{base_name}_zoom_out{ext}"
new_label_name = f"{base_name}_zoom_out.txt"
# Salvar imagem transformada
output_image_path = os.path.join(output_images_folder, new_image_name)
cv2.imwrite(output_image_path, augmented_image)
# Salvar anotações transformadas
output_label_path = os.path.join(output_labels_folder, new_label_name)
with open(output_label_path, "w") as f:
for ann in new_annotations:
f.write(f"{ann[0]} {ann[1]:.6f} {ann[2]:.6f} {ann[3]:.6f} {ann[4]:.6f}\n")
print(f"Processado e salvo como: {new_image_name}")