agrobot_base/Python/yolov7-label-weed/detect.py

70 lines
2.5 KiB
Python

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)