import cv2
import os
import numpy as np

def process_image(image_path, output_image_path, output_label_path):
    # Carregar a imagem
    image = cv2.imread(image_path)

    # Converter para escala de cinza
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Aplicar blur para reduzir o ruído
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # Aplicar a detecção de borda (Canny)
    edges = cv2.Canny(blurred, 50, 150)

    # Encontrar contornos
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Encontrar a maior área de contorno, assumindo que é a erva daninha
    if contours:
        contour = max(contours, key=cv2.contourArea)

        # Obter a caixa delimitadora da erva daninha
        x, y, w, h = cv2.boundingRect(contour)

        # Desenhar a caixa delimitadora na imagem para visualização
        cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)

        # Salvar a imagem com a caixa delimitadora
        cv2.imwrite(output_image_path, image)

        # Normalizar as coordenadas
        height, width = image.shape[:2]
        x_center = (x + x + w) / 2 / width
        y_center = (y + y + h) / 2 / height
        bbox_width = w / width
        bbox_height = h / height

        # Classe (por exemplo, 0 se for a primeira classe)
        class_id = 0

        # Escrever o label no formato YOLO
        label = f"{class_id} {x_center:.6f} {y_center:.6f} {bbox_width:.6f} {bbox_height:.6f}"
        print(f"Label for {image_path}: {label}")

        # Salvar o label em um arquivo .txt
        with open(output_label_path, 'w') as f:
            f.write(label)
    else:
        print(f"No contours found for {image_path}")

def process_images(input_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        if filename.endswith(".jpg"):
            image_path = os.path.join(input_dir, filename)
            output_image_path = os.path.join(output_dir, filename)
            output_label_path = os.path.join(output_dir, filename.replace('.jpg', '.txt'))

            process_image(image_path, output_image_path, output_label_path)

input_dir = 'C:\\Users\\USER\\Downloads\\weed-datasets-master\\corn weed datasets\\bluegrass'  # Substitua pelo caminho da sua pasta de entrada
output_dir = 'C:\\train\\yolov7\\train\\new_ervas\\dataset\\images\\train'  # Substitua pelo caminho da sua pasta de saída

process_images(input_dir, output_dir)