import cv2
import os

# Função para converter rótulos para o formato YOLO
def convert_labels_to_yolo(input_folder, output_folder, classes):
    for filename in os.listdir(input_folder):
        if filename.endswith('.txt'):
            label_path = os.path.join(input_folder, filename)
            image_path = os.path.join(input_folder, filename.replace('.txt', '.jpeg'))

            image = cv2.imread(image_path)
            height, width, _ = image.shape

            with open(label_path, 'r') as file:
                lines = file.readlines()
                for line in lines:
                    class_id, x_center, y_center, box_width, box_height = map(float, line.strip().split())
                    
                    x_center *= width
                    y_center *= height
                    box_width *= width
                    box_height *= height

                    x_min = int(x_center - (box_width / 2))
                    y_min = int(y_center - (box_height / 2))
                    x_max = int(x_center + (box_width / 2))
                    y_max = int(y_center + (box_height / 2))

                    class_name = classes[int(class_id)]

                    # Escrever no arquivo YOLO
                    yolo_label = f"{class_id} {x_center / width} {y_center / height} {box_width / width} {box_height / height}\n"
                    yolo_label_path = os.path.join(output_folder, filename.replace('.txt', '.txt'))
                    with open(yolo_label_path, 'a') as yolo_file:
                        yolo_file.write(yolo_label)

                    # Desenhar caixa delimitadora na imagem
                    cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
                    cv2.putText(image, class_name, (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            # Salvar imagem com caixas delimitadoras
            output_image_path = os.path.join(output_folder, filename.replace('.txt', '_labeled.jpeg'))
            cv2.imwrite(output_image_path, image)

# Definir classes
classes = {0: 'crop', 1: 'weed'}  # Mapeamento de IDs de classe para nomes

# Chamar função para converter rótulos para formato YOLO
convert_labels_to_yolo('C:\\weed-set\\agri_data\\data', 'C:\\weed-set\\agri_data\\data', classes)