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