File size: 2,014 Bytes
3c8a6d9 e7a1319 3c8a6d9 e7a1319 3c8a6d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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() |