forked from datong-new/CPIAD
-
Notifications
You must be signed in to change notification settings - Fork 0
/
yolov4_helper.py
111 lines (96 loc) · 3.49 KB
/
yolov4_helper.py
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import torch
from torch.nn import functional as F
from torch import nn
import cv2
from torchvision import transforms
from PIL import ImageFile
from PIL import Image
from tool.darknet2pytorch import *
import numpy
from constant import *
def resize(img:torch.tensor):
# resize
img_pil = img.round().byte().detach().cpu().numpy()
img_pil = Image.fromarray(img_pil.astype('uint8'), 'RGB')
resize_small = transforms.Compose([
transforms.Resize((608, 608)),
])
img_pil = resize_small(img_pil)
img_pil = numpy.array(img_pil)
img_pil = torch.from_numpy(img_pil).float().to(device)
img = img.permute(2,0,1)
img = nn.functional.interpolate(img.unsqueeze(0), size=(608, 608),mode="bilinear", align_corners=False)
img = img.squeeze(0).permute(1,2,0)
img = img + (img_pil-img.detach())
return img
def _input_transform(img:torch.tensor):
## bgr2rgb
new_img = torch.zeros(img.shape).to(device)
new_img[:,:,0] = new_img[:,:,0]+img[:,:,2]
new_img[:,:,1] = new_img[:,:,1]+img[:,:,1]
new_img[:,:,2] = new_img[:,:,2]+img[:,:,0]
img = new_img
# reisze
img = resize(img)
img = img.permute(2,0,1).contiguous()
#img = img.float().div(255.0)
img = img.float()/255.0
img = img.unsqueeze(0)
return img
def transforms_test(img_path="../images/991.png"):
img1 = Image.open(img_path).convert('RGB')
resize_small = transforms.Compose([
transforms.Resize((608, 608)),
])
img_pil = numpy.array(img1)
img1 = resize_small(img1)
width = img1.width
height = img1.height
img1 = torch.ByteTensor(torch.ByteStorage.from_buffer(img1.tobytes()))
img1 = img1.view(height, width, 3).transpose(0, 1).transpose(0, 2).contiguous()
img1 = img1.view(1, 3, height, width)
img1 = img1.float().div(255.0).to(device)
img2 = cv2.imread(img_path)
img2 = torch.from_numpy(img2).float()
img2 = _input_transform(img2).to(device)
print((img1-img2).sum())
assert (img1-img2).sum()==0
class Helper():
def __init__(self):
cfgfile = "models/yolov4.cfg"
weightfile = "models/yolov4.weights"
self.darknet_model = Darknet(cfgfile)
self.darknet_model.load_weights(weightfile)
#self.darknet_model = self.darknet_model.train().to(device)
self.darknet_model = self.darknet_model.eval().to(device)
def input_transforms(self, img:torch.tensor):
img = _input_transform(img)
def get_cls_scores(self, img:torch.tensor):
img = _input_transform(img).to(device)
output = self.darknet_model(img)
self.features = self.darknet_model.features
scores = []
for item in output:
h, w = item.shape[-2], item.shape[-1]
item = item.reshape(-1, 5+80, h*w).permute(1,0,2).reshape(5+80, -1)
scores += [item[4, :].sigmoid()]
return scores
async def attack_loss(self, img, t=0.45):
img = img.to(device)
scores = self.get_cls_scores(img)
thresh_loss = 0
objects_num = 0
for score in scores:
objects_num += (score>0.5).sum().item()
mask = score>0.45
score = score *mask
if mask.sum()!=0: thresh_loss += (score.sum() / mask.sum())
return thresh_loss, objects_num
if __name__ == "__main__":
transforms_test()
img_path = "../images/991.png"
img = cv2.imread(img_path)
img = torch.from_numpy(img).float()
helper = Helper()
attack_loss = helper.attack_loss(img)
print(attack_loss)