agrobot_base/Python/trainning/create-label.py

51 lines
2.1 KiB
Python

import cv2
import os
def create_yolo_annotation(image_path, label, contours, output_folder):
# Obter o nome do arquivo sem a extensão
file_name = os.path.splitext(os.path.basename(image_path))[0]
# Obter a largura e altura da imagem
img = cv2.imread(image_path)
height, width, _ = img.shape
# Caminho para a pasta de saída
output_path = os.path.join(output_folder, file_name + ".txt")
# Abrir arquivo para escrever anotações YOLO na pasta de saída
with open(output_path, "w") as f:
for contour in contours:
# Extrair coordenadas do contorno
x, y, w, h = cv2.boundingRect(contour)
# Calcular coordenadas normalizadas do contorno
x_center = (x + w / 2) / width
y_center = (y + h / 2) / height
width_normalized = w / width
height_normalized = h / height
# Escrever no arquivo YOLO
f.write(f"{label} {x_center} {y_center} {width_normalized} {height_normalized}\n")
print(f"Arquivo YOLO criado para {image_path}")
# Pasta contendo as imagens com os contornos
images_folder = "C:\\Zendion Inc\\agrobot_base\\Ruas de cana\\rotulados\\images"
# Pasta de saída para os arquivos YOLO
output_folder = "C:\\Zendion Inc\\agrobot_base\\Ruas de cana\\rotulados\\labels"
# Loop através das imagens no diretório
for image_file in os.listdir(images_folder):
if image_file.endswith(".jpeg"):
image_path = os.path.join(images_folder, image_file)
# Aqui, substitua esta parte pelo código que obtém os contornos específicos da rua
# Exemplo básico: obtenha os contornos utilizando a função cv2.findContours
img = cv2.imread(image_path)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
_, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Chame a função para criar o arquivo YOLO na pasta de saída
create_yolo_annotation(image_path, "0", contours, output_folder) # Rótulo "0" para a classe da rua