37 lines
1.5 KiB
Python
37 lines
1.5 KiB
Python
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.")
|