import os
from sklearn.model_selection import train_test_split

# Caminho para a pasta que contém as imagens
diretorio = "C:/train/models/ervas/dataset"

# Nomes dos arquivos de saída
arquivo_saida_train = "C:/train/models/ervas/train.txt"
arquivo_saida_test = "C:/train/models/ervas/test.txt"

# Lista para guardar os caminhos das imagens
imagens_para_treinamento = []

# Percorre todos os arquivos no diretório especificado
for arquivo in os.listdir(diretorio):
    # Verifica se o arquivo é um .jpeg e não contém "_labeled" no nome
    if arquivo.endswith(".jpeg") and "_labeled" not in arquivo:
        # Adiciona o caminho completo do arquivo à lista
        caminho_completo = os.path.join(diretorio, arquivo).replace("\\", "/")
        imagens_para_treinamento.append(caminho_completo)

# Divide as imagens em 80% para treinamento e 20% para teste
imagens_train, imagens_test = train_test_split(imagens_para_treinamento, test_size=0.2, random_state=42)

# Escreve os caminhos das imagens no arquivo de treinamento
with open(arquivo_saida_train, "w") as arquivo:
    for caminho in imagens_train:
        arquivo.write(caminho + "\n")

# Escreve os caminhos das imagens no arquivo de teste
with open(arquivo_saida_test, "w") as arquivo:
    for caminho in imagens_test:
        arquivo.write(caminho + "\n")

print(f"Arquivo {arquivo_saida_train} criado com sucesso com {len(imagens_train)} imagens para treinamento.")
print(f"Arquivo {arquivo_saida_test} criado com sucesso com {len(imagens_test)} imagens para teste.")