100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
import os
|
|
import requests
|
|
from PIL import Image, ImageDraw
|
|
import numpy as np
|
|
import json
|
|
|
|
# Função para enviar requisição para o ChatGPT e obter a resposta
|
|
def get_gpt_response(image_path):
|
|
# Dados do formulário
|
|
form = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Identify the weed in this image, returning just label in yolov7 format",
|
|
}
|
|
]
|
|
}
|
|
formData = json.dumps(form)
|
|
|
|
# Abrir a imagem como arquivo binário
|
|
with open(image_path, 'rb') as f:
|
|
image_data = f.read()
|
|
|
|
# Configuração da requisição
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": "Bearer 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",
|
|
}
|
|
url = "https://api.openai.com/v1/chat/completions"
|
|
|
|
# Fazendo a requisição POST
|
|
response = requests.post(url, headers=headers, data=formData)
|
|
|
|
if response.status_code == 200:
|
|
return response.json()['choices'][0]['message']['content']
|
|
else:
|
|
print(f"Erro ao obter resposta do ChatGPT: {response.status_code}")
|
|
try:
|
|
error_data = response.json()
|
|
print(error_data)
|
|
except ValueError:
|
|
print("Resposta inválida do servidor")
|
|
return None
|
|
|
|
# Função para gerar arquivo de rótulos no formato YOLOv7
|
|
def generate_yolov7_label(label_file, class_name, bbox_coords, image_size):
|
|
x_center = bbox_coords[0] + bbox_coords[2] / 2.0
|
|
y_center = bbox_coords[1] + bbox_coords[3] / 2.0
|
|
width = bbox_coords[2]
|
|
height = bbox_coords[3]
|
|
|
|
x_center /= image_size[0]
|
|
y_center /= image_size[1]
|
|
width /= image_size[0]
|
|
height /= image_size[1]
|
|
|
|
with open(label_file, 'w') as f:
|
|
f.write(f"{class_name} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
|
|
|
|
# Função para desenhar o bbox na imagem
|
|
def draw_bbox(image_path, bbox_coords, output_image_path):
|
|
image = Image.open(image_path)
|
|
draw = ImageDraw.Draw(image)
|
|
draw.rectangle(bbox_coords, outline='red', width=3)
|
|
image.save(output_image_path)
|
|
|
|
# Pasta de entrada e saída
|
|
input_folder = 'imagens/'
|
|
output_folder = 'saida/'
|
|
|
|
# Iterar sobre todas as imagens na pasta de entrada
|
|
for filename in os.listdir(input_folder):
|
|
if filename.endswith('.jpg') or filename.endswith('.png'):
|
|
image_path = os.path.join(input_folder, filename)
|
|
output_label_file = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.txt")
|
|
output_image_path = os.path.join(output_folder, filename)
|
|
|
|
# Obter rótulo do ChatGPT
|
|
label = get_gpt_response(image_path)
|
|
|
|
print(label)
|
|
|
|
# Suponha que o ChatGPT retorna algo como "weed: dandelion, bbox: [x, y, width, height]"
|
|
if label.startswith("weed:"):
|
|
parts = label.split(',')
|
|
class_name = parts[0].split(':')[1].strip()
|
|
bbox_coords = list(map(int, parts[1].split(':')[1].strip()[1:-1].split()))
|
|
|
|
# Gerar arquivo de rótulo YOLOv7
|
|
image_size = Image.open(image_path).size
|
|
generate_yolov7_label(output_label_file, class_name, bbox_coords, image_size)
|
|
|
|
# Desenhar o bbox na imagem
|
|
draw_bbox(image_path, bbox_coords, output_image_path)
|
|
|
|
print(f"Processado: {filename}")
|
|
else:
|
|
print(f"Não foi possível identificar a erva em: {filename}")
|