Spaces:
Sleeping
Sleeping
| 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}") | |