50 lines
2.2 KiB
Python
50 lines
2.2 KiB
Python
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)
|