import torch
from yolov5 import YOLOv5
from art.attacks.evasion import FastGradientMethod
from art.estimators.classification import PyTorchClassifier  # Alteração para nova versão do ART
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

# Carregar o modelo YOLOv5
model = YOLOv5('yolov5s.pt', device='cuda')  # Use yolov5s.pt ou outro modelo

# Carregar uma imagem de exemplo
image_path = 'caminho/para/imagem.jpg'  # Substitua pelo caminho correto
img = Image.open(image_path).convert('RGB')
img = np.array(img)  # Converter para array numpy

# Preparar a imagem para o modelo
img_tensor = torch.from_numpy(img).float() / 255.0
img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).cuda()  # Adicionar batch e mover para GPU

# Fazer a detecção normal
results = model(img)
results.show()

# Criar o classificador PyTorch para ART
def predict_fn(x):
    # Função de predição para ART
    with torch.no_grad():
        return model(x).xywh[0][:, 4].cpu().numpy()  # Retorna as probabilidades de deteção de objetos

# Criando o classificador ART
classifier = PyTorchClassifier(
    model=model.model,  # Passando o modelo YOLO para ART
    loss_fn=torch.nn.CrossEntropyLoss(),
    input_shape=(3, 640, 640),  # YOLOv5 trabalha com imagens de 640x640
    nb_classes=80,  # 80 classes do COCO dataset
    clip_values=(0, 1),  # Normalização das imagens
    preprocessing=(0, 1)
)

# Criar o ataque adversarial com o ART (exemplo com FGSM)
attack = FastGradientMethod(estimator=classifier, eps=0.1)  # eps controla a intensidade do ataque
adv_img = attack.generate(x=img_tensor.cpu().numpy())

# Converter a imagem adversarial para um formato visualizável
adv_img = np.moveaxis(adv_img[0], 0, -1)  # Move as dimensões para (H, W, C)
adv_img = np.clip(adv_img * 255, 0, 255).astype(np.uint8)

# Mostrar a imagem adversarial
plt.imshow(adv_img)
plt.title('Imagem Adversarial')
plt.show()

# Aplicar o modelo YOLO na imagem adversarial
adv_results = model(adv_img)
adv_results.show()