yolo pytorch

Halcom

yolov3 pytorch

1. # -*- coding: utf-8 -*-
   
2. """
   
3. Created on Sat Jun 17 23:33:13 2023
   
4. @author: ysw
   
5. """
   
6. #-------------------------------------#
   
7. #       对数据集进行训练
   
8. #-------------------------------------#
   
9. import datetime
   
10. import os
    
11. import math
    
12. import cv2
    
13. from PIL import Image
    
14. import numpy as np
    
15. from collections import OrderedDict
    
16. import torch
    
17. import torch.backends.cudnn as cudnn
    
18. import torch.distributed as dist
    
19. import torch.nn as nn
    
20. import torch.optim as optim
    
21. from torch import nn
    
22. from torch.utils.data import DataLoader
    
23. from torch.utils.data.dataset import Dataset
    
24. # from nets.yolo import YoloBody
    
25. # from nets.yolo_training import YOLOLoss
    
26. # from utils.dataloader import YoloDataset, yolo_dataset_collate
    
27. # from utils.utils import get_anchors, get_classes
    
28. 
    
29. def conv2d(filter_in, filter_out, kernel_size):
    
30.     pad = (kernel_size - 1) // 2 if kernel_size else 0
    
31.     return nn.Sequential(OrderedDict([
    
32.         ("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=1, padding=pad, bias=False)),
    
33.         ("bn", nn.BatchNorm2d(filter_out)),
    
34.         ("relu", nn.LeakyReLU(0.1)),
    
35.     ]))
    
36. #------------------------------------------------------------------------#
    
37. #   make_last_layers里面一共有七个卷积，前五个用于提取特征。
    
38. #   后两个用于获得yolo网络的预测结果
    
39. #------------------------------------------------------------------------#
    
40. def make_last_layers(filters_list, in_filters, out_filter):
    
41.     m = nn.Sequential(
    
42.         conv2d(in_filters, filters_list[0], 1),
    
43.         conv2d(filters_list[0], filters_list[1], 3),
    
44.         conv2d(filters_list[1], filters_list[0], 1),
    
45.         conv2d(filters_list[0], filters_list[1], 3),
    
46.         conv2d(filters_list[1], filters_list[0], 1),
    
47.         conv2d(filters_list[0], filters_list[1], 3),
    
48.         nn.Conv2d(filters_list[1], out_filter, kernel_size=1, stride=1, padding=0, bias=True)
    
49.     )
    
50.     return m
    
51. 
    
52. class YoloBody(nn.Module):
    
53.     def __init__(self, anchors_mask, num_classes, phi=0, load_weights = False):
    
54.         super(YoloBody, self).__init__()
    
55.         #---------------------------------------------------#   
    
56.         #   生成efficientnet的主干模型，以efficientnetB0为例
    
57.         #   获得三个有效特征层，他们的shape分别是：
    
58.         #   52, 52, 40
    
59.         #   26, 26, 112
    
60.         #   13, 13, 320
    
61.         #---------------------------------------------------#
    
62.         self.backbone = EfficientNet(phi, load_weights = load_weights)
    
63. 
    
64.         out_filters = {
    
65.             0: [40, 112, 320],
    
66.             1: [40, 112, 320],
    
67.             2: [48, 120, 352],
    
68.             3: [48, 136, 384],
    
69.             4: [56, 160, 448],
    
70.             5: [64, 176, 512],
    
71.             6: [72, 200, 576],
    
72.             7: [80, 224, 640],
    
73.         }[phi]
    
74. 
    
75.         self.last_layer0            = make_last_layers([out_filters[-1], int(out_filters[-1]*2)], out_filters[-1], len(anchors_mask[0]) * (num_classes + 5))
    
76. 
    
77.         self.last_layer1_conv       = conv2d(out_filters[-1], out_filters[-2], 1)
    
78.         self.last_layer1_upsample   = nn.Upsample(scale_factor=2, mode='nearest')
    
79.         self.last_layer1            = make_last_layers([out_filters[-2], int(out_filters[-2]*2)], out_filters[-2] + out_filters[-2], len(anchors_mask[1]) * (num_classes + 5))
    
80. 
    
81.         self.last_layer2_conv       = conv2d(out_filters[-2], out_filters[-3], 1)
    
82.         self.last_layer2_upsample   = nn.Upsample(scale_factor=2, mode='nearest')
    
83.         self.last_layer2            = make_last_layers([out_filters[-3], int(out_filters[-3]*2)], out_filters[-3] + out_filters[-3], len(anchors_mask[2]) * (num_classes + 5))
    
84. 
    
85.     def forward(self, x):
    
86.         #---------------------------------------------------#   
    
87.         #   获得三个有效特征层，他们的shape分别是：
    
88.         #   52, 52, 40
    
89.         #   26, 26, 112
    
90.         #   13, 13, 320
    
91.         #---------------------------------------------------#
    
92.         x2, x1, x0 = self.backbone(x)
    
93. 
    
94.         #---------------------------------------------------#
    
95.         #   第一个特征层
    
96.         #   out0 = (batch_size,255,13,13)
    
97.         #---------------------------------------------------#
    
98.         out0_branch = self.last_layer0[:5](x0)
    
99.         out0        = self.last_layer0[5:](out0_branch)
    
100. 
     
101.         x1_in = self.last_layer1_conv(out0_branch)
     
102.         x1_in = self.last_layer1_upsample(x1_in)
     
103. 
     
104.         x1_in = torch.cat([x1_in, x1], 1)
     
105.         #---------------------------------------------------#
     
106.         #   第二个特征层
     
107.         #   out1 = (batch_size,255,26,26)
     
108.         #---------------------------------------------------#
     
109.         out1_branch = self.last_layer1[:5](x1_in)
     
110.         out1        = self.last_layer1[5:](out1_branch)
     
111. 
     
112.         x2_in = self.last_layer2_conv(out1_branch)
     
113.         x2_in = self.last_layer2_upsample(x2_in)
     
114. 
     
115.         x2_in = torch.cat([x2_in, x2], 1)
     
116.         #---------------------------------------------------#
     
117.         #   第一个特征层
     
118.         #   out3 = (batch_size,255,52,52)
     
119.         #---------------------------------------------------#
     
120.         out2 = self.last_layer2(x2_in)
     
121.         return out0, out1, out2
     
122. 
     
123. class YOLOLoss(nn.Module):
     
124.     def __init__(self, anchors, num_classes, input_shape, cuda, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]):
     
125.         super(YOLOLoss, self).__init__()
     
126.         #-----------------------------------------------------------#
     
127.         #   13x13的特征层对应的anchor是[116,90],[156,198],[373,326]
     
128.         #   26x26的特征层对应的anchor是[30,61],[62,45],[59,119]
     
129.         #   52x52的特征层对应的anchor是[10,13],[16,30],[33,23]
     
130.         #-----------------------------------------------------------#
     
131.         self.anchors        = anchors
     
132.         self.num_classes    = num_classes
     
133.         self.bbox_attrs     = 5 + num_classes
     
134.         self.input_shape    = input_shape
     
135.         self.anchors_mask   = anchors_mask
     
136. 
     
137.         self.giou           = True
     
138.         self.balance        = [0.4, 1.0, 4]
     
139.         self.box_ratio      = 0.05
     
140.         self.obj_ratio      = 5 * (input_shape[0] * input_shape[1]) / (416 ** 2)
     
141.         self.cls_ratio      = 1 * (num_classes / 80)
     
142. 
     
143.         self.ignore_threshold = 0.5
     
144.         self.cuda           = cuda
     
145. 
     
146.     def clip_by_tensor(self, t, t_min, t_max):
     
147.         t = t.float()
     
148.         result = (t >= t_min).float() * t + (t < t_min).float() * t_min
     
149.         result = (result <= t_max).float() * result + (result > t_max).float() * t_max
     
150.         return result
     
151. 
     
152.     def MSELoss(self, pred, target):
     
153.         return torch.pow(pred - target, 2)
     
154. 
     
155.     def BCELoss(self, pred, target):
     
156.         epsilon = 1e-7
     
157.         pred    = self.clip_by_tensor(pred, epsilon, 1.0 - epsilon)
     
158.         output  = - target * torch.log(pred) - (1.0 - target) * torch.log(1.0 - pred)
     
159.         return output
     
160. 
     
161.     def box_giou(self, b1, b2):
     
162.         """
     
163.         输入为：
     
164.         ----------
     
165.         b1: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh
     
166.         b2: tensor, shape=(batch, feat_w, feat_h, anchor_num, 4), xywh
     
167. 
     
168.         返回为：
     
169.         -------
     
170.         giou: tensor, shape=(batch, feat_w, feat_h, anchor_num, 1)
     
171.         """
     
172.         #----------------------------------------------------#
     
173.         #   求出预测框左上角右下角
     
174.         #----------------------------------------------------#
     
175.         b1_xy       = b1[..., :2]
     
176.         b1_wh       = b1[..., 2:4]
     
177.         b1_wh_half  = b1_wh/2.
     
178.         b1_mins     = b1_xy - b1_wh_half
     
179.         b1_maxes    = b1_xy + b1_wh_half
     
180.         #----------------------------------------------------#
     
181.         #   求出真实框左上角右下角
     
182.         #----------------------------------------------------#
     
183.         b2_xy       = b2[..., :2]
     
184.         b2_wh       = b2[..., 2:4]
     
185.         b2_wh_half  = b2_wh/2.
     
186.         b2_mins     = b2_xy - b2_wh_half
     
187.         b2_maxes    = b2_xy + b2_wh_half
     
188. 
     
189.         #----------------------------------------------------#
     
190.         #   求真实框和预测框所有的iou
     
191.         #----------------------------------------------------#
     
192.         intersect_mins  = torch.max(b1_mins, b2_mins)
     
193.         intersect_maxes = torch.min(b1_maxes, b2_maxes)
     
194.         intersect_wh    = torch.max(intersect_maxes - intersect_mins, torch.zeros_like(intersect_maxes))
     
195.         intersect_area  = intersect_wh[..., 0] * intersect_wh[..., 1]
     
196.         b1_area         = b1_wh[..., 0] * b1_wh[..., 1]
     
197.         b2_area         = b2_wh[..., 0] * b2_wh[..., 1]
     
198.         union_area      = b1_area + b2_area - intersect_area
     
199.         iou             = intersect_area / union_area
     
200. 
     
201.         #----------------------------------------------------#
     
202.         #   找到包裹两个框的最小框的左上角和右下角
     
203.         #----------------------------------------------------#
     
204.         enclose_mins    = torch.min(b1_mins, b2_mins)
     
205.         enclose_maxes   = torch.max(b1_maxes, b2_maxes)
     
206.         enclose_wh      = torch.max(enclose_maxes - enclose_mins, torch.zeros_like(intersect_maxes))
     
207.         #----------------------------------------------------#
     
208.         #   计算对角线距离
     
209.         #----------------------------------------------------#
     
210.         enclose_area    = enclose_wh[..., 0] * enclose_wh[..., 1]
     
211.         giou            = iou - (enclose_area - union_area) / enclose_area
     
212.         
     
213.         return giou
     
214.         
     
215.     def forward(self, l, input, targets=None):
     
216.         #----------------------------------------------------#
     
217.         #   l代表的是，当前输入进来的有效特征层，是第几个有效特征层
     
218.         #   input的shape为  bs, 3*(5+num_classes), 13, 13
     
219.         #                   bs, 3*(5+num_classes), 26, 26
     
220.         #                   bs, 3*(5+num_classes), 52, 52
     
221.         #   targets代表的是真实框。
     
222.         #----------------------------------------------------#
     
223.         #--------------------------------#
     
224.         #   获得图片数量，特征层的高和宽
     
225.         #   13和13
     
226.         #--------------------------------#
     
227.         bs      = input.size(0)
     
228.         in_h    = input.size(2)
     
229.         in_w    = input.size(3)
     
230.         #-----------------------------------------------------------------------#
     
231.         #   计算步长
     
232.         #   每一个特征点对应原来的图片上多少个像素点
     
233.         #   如果特征层为13x13的话，一个特征点就对应原来的图片上的32个像素点
     
234.         #   如果特征层为26x26的话，一个特征点就对应原来的图片上的16个像素点
     
235.         #   如果特征层为52x52的话，一个特征点就对应原来的图片上的8个像素点
     
236.         #   stride_h = stride_w = 32、16、8
     
237.         #   stride_h和stride_w都是32。
     
238.         #-----------------------------------------------------------------------#
     
239.         stride_h = self.input_shape[0] / in_h
     
240.         stride_w = self.input_shape[1] / in_w
     
241.         #-------------------------------------------------#
     
242.         #   此时获得的scaled_anchors大小是相对于特征层的
     
243.         #-------------------------------------------------#
     
244.         scaled_anchors  = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors]
     
245.         #-----------------------------------------------#
     
246.         #   输入的input一共有三个，他们的shape分别是
     
247.         #   bs, 3*(5+num_classes), 13, 13 => batch_size, 3, 13, 13, 5 + num_classes
     
248.         #   batch_size, 3, 26, 26, 5 + num_classes
     
249.         #   batch_size, 3, 52, 52, 5 + num_classes
     
250.         #-----------------------------------------------#
     
251.         prediction = input.view(bs, len(self.anchors_mask[l]), self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
     
252.         
     
253.         #-----------------------------------------------#
     
254.         #   先验框的中心位置的调整参数
     
255.         #-----------------------------------------------#
     
256.         x = torch.sigmoid(prediction[..., 0])
     
257.         y = torch.sigmoid(prediction[..., 1])
     
258.         #-----------------------------------------------#
     
259.         #   先验框的宽高调整参数
     
260.         #-----------------------------------------------#
     
261.         w = prediction[..., 2]
     
262.         h = prediction[..., 3]
     
263.         #-----------------------------------------------#
     
264.         #   获得置信度，是否有物体
     
265.         #-----------------------------------------------#
     
266.         conf = torch.sigmoid(prediction[..., 4])
     
267.         #-----------------------------------------------#
     
268.         #   种类置信度
     
269.         #-----------------------------------------------#
     
270.         pred_cls = torch.sigmoid(prediction[..., 5:])
     
271. 
     
272.         #-----------------------------------------------#
     
273.         #   获得网络应该有的预测结果
     
274.         #-----------------------------------------------#
     
275.         y_true, noobj_mask, box_loss_scale = self.get_target(l, targets, scaled_anchors, in_h, in_w)
     
276. 
     
277.         #---------------------------------------------------------------#
     
278.         #   将预测结果进行解码，判断预测结果和真实值的重合程度
     
279.         #   如果重合程度过大则忽略，因为这些特征点属于预测比较准确的特征点
     
280.         #   作为负样本不合适
     
281.         #----------------------------------------------------------------#
     
282.         noobj_mask, pred_boxes = self.get_ignore(l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask)
     
283. 
     
284.         if self.cuda:
     
285.             y_true          = y_true.type_as(x)
     
286.             noobj_mask      = noobj_mask.type_as(x)
     
287.             box_loss_scale  = box_loss_scale.type_as(x)
     
288.         #--------------------------------------------------------------------------#
     
289.         #   box_loss_scale是真实框宽高的乘积，宽高均在0-1之间，因此乘积也在0-1之间。
     
290.         #   2-宽高的乘积代表真实框越大，比重越小，小框的比重更大。
     
291.         #--------------------------------------------------------------------------#
     
292.         box_loss_scale = 2 - box_loss_scale
     
293.             
     
294.         loss        = 0
     
295.         obj_mask    = y_true[..., 4] == 1
     
296.         n           = torch.sum(obj_mask)
     
297.         if n != 0:
     
298.             if self.giou:
     
299.                 #---------------------------------------------------------------#
     
300.                 #   计算预测结果和真实结果的giou
     
301.                 #----------------------------------------------------------------#
     
302.                 giou        = self.box_giou(pred_boxes, y_true[..., :4]).type_as(x)
     
303.                 loss_loc    = torch.mean((1 - giou)[obj_mask])
     
304.             else:
     
305.                 #-----------------------------------------------------------#
     
306.                 #   计算中心偏移情况的loss，使用BCELoss效果好一些
     
307.                 #-----------------------------------------------------------#
     
308.                 loss_x      = torch.mean(self.BCELoss(x[obj_mask], y_true[..., 0][obj_mask]) * box_loss_scale[obj_mask])
     
309.                 loss_y      = torch.mean(self.BCELoss(y[obj_mask], y_true[..., 1][obj_mask]) * box_loss_scale[obj_mask])
     
310.                 #-----------------------------------------------------------#
     
311.                 #   计算宽高调整值的loss
     
312.                 #-----------------------------------------------------------#
     
313.                 loss_w      = torch.mean(self.MSELoss(w[obj_mask], y_true[..., 2][obj_mask]) * box_loss_scale[obj_mask])
     
314.                 loss_h      = torch.mean(self.MSELoss(h[obj_mask], y_true[..., 3][obj_mask]) * box_loss_scale[obj_mask])
     
315.                 loss_loc    = (loss_x + loss_y + loss_h + loss_w) * 0.1
     
316. 
     
317.             loss_cls    = torch.mean(self.BCELoss(pred_cls[obj_mask], y_true[..., 5:][obj_mask]))
     
318.             loss        += loss_loc * self.box_ratio + loss_cls * self.cls_ratio
     
319. 
     
320.         loss_conf   = torch.mean(self.BCELoss(conf, obj_mask.type_as(conf))[noobj_mask.bool() | obj_mask])
     
321.         loss        += loss_conf * self.balance[l] * self.obj_ratio
     
322.         # if n != 0:
     
323.         #     print(loss_loc * self.box_ratio, loss_cls * self.cls_ratio, loss_conf * self.balance[l] * self.obj_ratio)
     
324.         return loss
     
325. 
     
326.     def calculate_iou(self, _box_a, _box_b):
     
327.         #-----------------------------------------------------------#
     
328.         #   计算真实框的左上角和右下角
     
329.         #-----------------------------------------------------------#
     
330.         b1_x1, b1_x2 = _box_a[:, 0] - _box_a[:, 2] / 2, _box_a[:, 0] + _box_a[:, 2] / 2
     
331.         b1_y1, b1_y2 = _box_a[:, 1] - _box_a[:, 3] / 2, _box_a[:, 1] + _box_a[:, 3] / 2
     
332.         #-----------------------------------------------------------#
     
333.         #   计算先验框获得的预测框的左上角和右下角
     
334.         #-----------------------------------------------------------#
     
335.         b2_x1, b2_x2 = _box_b[:, 0] - _box_b[:, 2] / 2, _box_b[:, 0] + _box_b[:, 2] / 2
     
336.         b2_y1, b2_y2 = _box_b[:, 1] - _box_b[:, 3] / 2, _box_b[:, 1] + _box_b[:, 3] / 2
     
337. 
     
338.         #-----------------------------------------------------------#
     
339.         #   将真实框和预测框都转化成左上角右下角的形式
     
340.         #-----------------------------------------------------------#
     
341.         box_a = torch.zeros_like(_box_a)
     
342.         box_b = torch.zeros_like(_box_b)
     
343.         box_a[:, 0], box_a[:, 1], box_a[:, 2], box_a[:, 3] = b1_x1, b1_y1, b1_x2, b1_y2
     
344.         box_b[:, 0], box_b[:, 1], box_b[:, 2], box_b[:, 3] = b2_x1, b2_y1, b2_x2, b2_y2
     
345. 
     
346.         #-----------------------------------------------------------#
     
347.         #   A为真实框的数量，B为先验框的数量
     
348.         #-----------------------------------------------------------#
     
349.         A = box_a.size(0)
     
350.         B = box_b.size(0)
     
351. 
     
352.         #-----------------------------------------------------------#
     
353.         #   计算交的面积
     
354.         #-----------------------------------------------------------#
     
355.         max_xy  = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
     
356.         min_xy  = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))
     
357.         inter   = torch.clamp((max_xy - min_xy), min=0)
     
358.         inter   = inter[:, :, 0] * inter[:, :, 1]
     
359.         #-----------------------------------------------------------#
     
360.         #   计算预测框和真实框各自的面积
     
361.         #-----------------------------------------------------------#
     
362.         area_a = ((box_a[:, 2]-box_a[:, 0]) * (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
     
363.         area_b = ((box_b[:, 2]-box_b[:, 0]) * (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
     
364.         #-----------------------------------------------------------#
     
365.         #   求IOU
     
366.         #-----------------------------------------------------------#
     
367.         union = area_a + area_b - inter
     
368.         return inter / union  # [A,B]
     
369.    
     
370.     def get_target(self, l, targets, anchors, in_h, in_w):
     
371.         #-----------------------------------------------------#
     
372.         #   计算一共有多少张图片
     
373.         #-----------------------------------------------------#
     
374.         bs              = len(targets)
     
375.         #-----------------------------------------------------#
     
376.         #   用于选取哪些先验框不包含物体
     
377.         #-----------------------------------------------------#
     
378.         noobj_mask      = torch.ones(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False)
     
379.         #-----------------------------------------------------#
     
380.         #   让网络更加去关注小目标
     
381.         #-----------------------------------------------------#
     
382.         box_loss_scale  = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, requires_grad = False)
     
383.         #-----------------------------------------------------#
     
384.         #   batch_size, 3, 13, 13, 5 + num_classes
     
385.         #-----------------------------------------------------#
     
386.         y_true          = torch.zeros(bs, len(self.anchors_mask[l]), in_h, in_w, self.bbox_attrs, requires_grad = False)
     
387.         for b in range(bs):            
     
388.             if len(targets[b])==0:
     
389.                 continue
     
390.             batch_target = torch.zeros_like(targets[b])
     
391.             #-------------------------------------------------------#
     
392.             #   计算出正样本在特征层上的中心点
     
393.             #-------------------------------------------------------#
     
394.             batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w
     
395.             batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
     
396.             batch_target[:, 4] = targets[b][:, 4]
     
397.             batch_target = batch_target.cpu()
     
398.             
     
399.             #-------------------------------------------------------#
     
400.             #   将真实框转换一个形式
     
401.             #   num_true_box, 4
     
402.             #-------------------------------------------------------#
     
403.             gt_box          = torch.FloatTensor(torch.cat((torch.zeros((batch_target.size(0), 2)), batch_target[:, 2:4]), 1))
     
404.             #-------------------------------------------------------#
     
405.             #   将先验框转换一个形式
     
406.             #   9, 4
     
407.             #-------------------------------------------------------#
     
408.             anchor_shapes   = torch.FloatTensor(torch.cat((torch.zeros((len(anchors), 2)), torch.FloatTensor(anchors)), 1))
     
409.             #-------------------------------------------------------#
     
410.             #   计算交并比
     
411.             #   self.calculate_iou(gt_box, anchor_shapes) = [num_true_box, 9]每一个真实框和9个先验框的重合情况
     
412.             #   best_ns:
     
413.             #   [每个真实框最大的重合度max_iou, 每一个真实框最重合的先验框的序号]
     
414.             #-------------------------------------------------------#
     
415.             best_ns = torch.argmax(self.calculate_iou(gt_box, anchor_shapes), dim=-1)
     
416. 
     
417.             for t, best_n in enumerate(best_ns):
     
418.                 if best_n not in self.anchors_mask[l]:
     
419.                     continue
     
420.                 #----------------------------------------#
     
421.                 #   判断这个先验框是当前特征点的哪一个先验框
     
422.                 #----------------------------------------#
     
423.                 k = self.anchors_mask[l].index(best_n)
     
424.                 #----------------------------------------#
     
425.                 #   获得真实框属于哪个网格点
     
426.                 #----------------------------------------#
     
427.                 i = torch.floor(batch_target[t, 0]).long()
     
428.                 j = torch.floor(batch_target[t, 1]).long()
     
429.                 #----------------------------------------#
     
430.                 #   取出真实框的种类
     
431.                 #----------------------------------------#
     
432.                 c = batch_target[t, 4].long()
     
433. 
     
434.                 #----------------------------------------#
     
435.                 #   noobj_mask代表无目标的特征点
     
436.                 #----------------------------------------#
     
437.                 noobj_mask[b, k, j, i] = 0
     
438.                 #----------------------------------------#
     
439.                 #   tx、ty代表中心调整参数的真实值
     
440.                 #----------------------------------------#
     
441.                 if not self.giou:
     
442.                     #----------------------------------------#
     
443.                     #   tx、ty代表中心调整参数的真实值
     
444.                     #----------------------------------------#
     
445.                     y_true[b, k, j, i, 0] = batch_target[t, 0] - i.float()
     
446.                     y_true[b, k, j, i, 1] = batch_target[t, 1] - j.float()
     
447.                     y_true[b, k, j, i, 2] = math.log(batch_target[t, 2] / anchors[best_n][0])
     
448.                     y_true[b, k, j, i, 3] = math.log(batch_target[t, 3] / anchors[best_n][1])
     
449.                     y_true[b, k, j, i, 4] = 1
     
450.                     y_true[b, k, j, i, c + 5] = 1
     
451.                 else:
     
452.                     #----------------------------------------#
     
453.                     #   tx、ty代表中心调整参数的真实值
     
454.                     #----------------------------------------#
     
455.                     y_true[b, k, j, i, 0] = batch_target[t, 0]
     
456.                     y_true[b, k, j, i, 1] = batch_target[t, 1]
     
457.                     y_true[b, k, j, i, 2] = batch_target[t, 2]
     
458.                     y_true[b, k, j, i, 3] = batch_target[t, 3]
     
459.                     y_true[b, k, j, i, 4] = 1
     
460.                     y_true[b, k, j, i, c + 5] = 1
     
461.                 #----------------------------------------#
     
462.                 #   用于获得xywh的比例
     
463.                 #   大目标loss权重小，小目标loss权重大
     
464.                 #----------------------------------------#
     
465.                 box_loss_scale[b, k, j, i] = batch_target[t, 2] * batch_target[t, 3] / in_w / in_h
     
466.         return y_true, noobj_mask, box_loss_scale
     
467. 
     
468.     def get_ignore(self, l, x, y, h, w, targets, scaled_anchors, in_h, in_w, noobj_mask):
     
469.         #-----------------------------------------------------#
     
470.         #   计算一共有多少张图片
     
471.         #-----------------------------------------------------#
     
472.         bs = len(targets)
     
473. 
     
474.         #-----------------------------------------------------#
     
475.         #   生成网格，先验框中心，网格左上角
     
476.         #-----------------------------------------------------#
     
477.         grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_h, 1).repeat(
     
478.             int(bs * len(self.anchors_mask[l])), 1, 1).view(x.shape).type_as(x)
     
479.         grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_w, 1).t().repeat(
     
480.             int(bs * len(self.anchors_mask[l])), 1, 1).view(y.shape).type_as(x)
     
481. 
     
482.         # 生成先验框的宽高
     
483.         scaled_anchors_l = np.array(scaled_anchors)[self.anchors_mask[l]]
     
484.         anchor_w = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([0])).type_as(x)
     
485.         anchor_h = torch.Tensor(scaled_anchors_l).index_select(1, torch.LongTensor([1])).type_as(x)
     
486.         
     
487.         anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
     
488.         anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
     
489.         #-------------------------------------------------------#
     
490.         #   计算调整后的先验框中心与宽高
     
491.         #-------------------------------------------------------#
     
492.         pred_boxes_x    = torch.unsqueeze(x + grid_x, -1)
     
493.         pred_boxes_y    = torch.unsqueeze(y + grid_y, -1)
     
494.         pred_boxes_w    = torch.unsqueeze(torch.exp(w) * anchor_w, -1)
     
495.         pred_boxes_h    = torch.unsqueeze(torch.exp(h) * anchor_h, -1)
     
496.         pred_boxes      = torch.cat([pred_boxes_x, pred_boxes_y, pred_boxes_w, pred_boxes_h], dim = -1)
     
497.         
     
498.         for b in range(bs):           
     
499.             #-------------------------------------------------------#
     
500.             #   将预测结果转换一个形式
     
501.             #   pred_boxes_for_ignore      num_anchors, 4
     
502.             #-------------------------------------------------------#
     
503.             pred_boxes_for_ignore = pred_boxes[b].view(-1, 4)
     
504.             #-------------------------------------------------------#
     
505.             #   计算真实框，并把真实框转换成相对于特征层的大小
     
506.             #   gt_box      num_true_box, 4
     
507.             #-------------------------------------------------------#
     
508.             if len(targets[b]) > 0:
     
509.                 batch_target = torch.zeros_like(targets[b])
     
510.                 #-------------------------------------------------------#
     
511.                 #   计算出正样本在特征层上的中心点
     
512.                 #-------------------------------------------------------#
     
513.                 batch_target[:, [0,2]] = targets[b][:, [0,2]] * in_w
     
514.                 batch_target[:, [1,3]] = targets[b][:, [1,3]] * in_h
     
515.                 batch_target = batch_target[:, :4].type_as(x)
     
516.                 #-------------------------------------------------------#
     
517.                 #   计算交并比
     
518.                 #   anch_ious       num_true_box, num_anchors
     
519.                 #-------------------------------------------------------#
     
520.                 anch_ious = self.calculate_iou(batch_target, pred_boxes_for_ignore)
     
521.                 #-------------------------------------------------------#
     
522.                 #   每个先验框对应真实框的最大重合度
     
523.                 #   anch_ious_max   num_anchors
     
524.                 #-------------------------------------------------------#
     
525.                 anch_ious_max, _    = torch.max(anch_ious, dim = 0)
     
526.                 anch_ious_max       = anch_ious_max.view(pred_boxes[b].size()[:3])
     
527.                 noobj_mask[b][anch_ious_max > self.ignore_threshold] = 0
     
528.         return noobj_mask, pred_boxes
     
529.    
     
530. #---------------------------------------------------------#
     
531. #   将图像转换成RGB图像，防止灰度图在预测时报错。
     
532. #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
     
533. #---------------------------------------------------------#
     
534. def cvtColor(image):
     
535.     if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
     
536.         return image
     
537.     else:
     
538.         image = image.convert('RGB')
     
539.         return image
     
540.    
     
541. #---------------------------------------------------#
     
542. #   对输入图像进行resize
     
543. #---------------------------------------------------#
     
544. def resize_image(image, size, letterbox_image):
     
545.     iw, ih  = image.size
     
546.     w, h    = size
     
547.     if letterbox_image:
     
548.         scale   = min(w/iw, h/ih)
     
549.         nw      = int(iw*scale)
     
550.         nh      = int(ih*scale)
     
551. 
     
552.         image   = image.resize((nw,nh), Image.BICUBIC)
     
553.         new_image = Image.new('RGB', size, (128,128,128))
     
554.         new_image.paste(image, ((w-nw)//2, (h-nh)//2))
     
555.     else:
     
556.         new_image = image.resize((w, h), Image.BICUBIC)
     
557.     return new_image
     
558. 
     
559. #---------------------------------------------------#
     
560. #   获得类
     
561. #---------------------------------------------------#
     
562. def get_classes(classes_path):
     
563.     with open(classes_path, encoding='utf-8') as f:
     
564.         class_names = f.readlines()
     
565.     class_names = [c.strip() for c in class_names]
     
566.     return class_names, len(class_names)
     
567. 
     
568. #---------------------------------------------------#
     
569. #   获得先验框
     
570. #---------------------------------------------------#
     
571. def get_anchors(anchors_path):
     
572.     '''loads the anchors from a file'''
     
573.     with open(anchors_path, encoding='utf-8') as f:
     
574.         anchors = f.readline()
     
575.     anchors = [float(x) for x in anchors.split(',')]
     
576.     anchors = np.array(anchors).reshape(-1, 2)
     
577.     return anchors, len(anchors)
     
578. 
     
579. #---------------------------------------------------#
     
580. #   获得学习率
     
581. #---------------------------------------------------#
     
582. def get_lr(optimizer):
     
583.     for param_group in optimizer.param_groups:
     
584.         return param_group['lr']
     
585. 
     
586. def preprocess_input(image):
     
587.     image /= 255.0
     
588.     return image
     
589. 
     
590. class YoloDataset(Dataset):
     
591.     def __init__(self, annotation_lines, input_shape, num_classes, train):
     
592.         super(YoloDataset, self).__init__()
     
593.         self.annotation_lines   = annotation_lines
     
594.         self.input_shape        = input_shape
     
595.         self.num_classes        = num_classes
     
596.         self.length             = len(self.annotation_lines)
     
597.         self.train              = train
     
598. 
     
599.     def __len__(self):
     
600.         return self.length
     
601. 
     
602.     def __getitem__(self, index):
     
603.         index       = index % self.length
     
604.         #---------------------------------------------------#
     
605.         #   训练时进行数据的随机增强
     
606.         #   验证时不进行数据的随机增强
     
607.         #---------------------------------------------------#
     
608.         image, box  = self.get_random_data(self.annotation_lines[index], self.input_shape[0:2], random = self.train)
     
609.         image       = np.transpose(preprocess_input(np.array(image, dtype=np.float32)), (2, 0, 1))
     
610.         box         = np.array(box, dtype=np.float32)
     
611.         if len(box) != 0:
     
612.             box[:, [0, 2]] = box[:, [0, 2]] / self.input_shape[1]
     
613.             box[:, [1, 3]] = box[:, [1, 3]] / self.input_shape[0]
     
614. 
     
615.             box[:, 2:4] = box[:, 2:4] - box[:, 0:2]
     
616.             box[:, 0:2] = box[:, 0:2] + box[:, 2:4] / 2
     
617.         return image, box
     
618. 
     
619.     def rand(self, a=0, b=1):
     
620.         return np.random.rand()*(b-a) + a
     
621. 
     
622.     def get_random_data(self, annotation_line, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.4, random=True):
     
623.         line    = annotation_line.split()
     
624.         #------------------------------#
     
625.         #   读取图像并转换成RGB图像
     
626.         #------------------------------#
     
627.         image   = Image.open(line[0])
     
628.         image   = cvtColor(image)
     
629.         #------------------------------#
     
630.         #   获得图像的高宽与目标高宽
     
631.         #------------------------------#
     
632.         iw, ih  = image.size
     
633.         h, w    = input_shape
     
634.         #------------------------------#
     
635.         #   获得预测框
     
636.         #------------------------------#
     
637.         box     = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
     
638. 
     
639.         if not random:
     
640.             scale = min(w/iw, h/ih)
     
641.             nw = int(iw*scale)
     
642.             nh = int(ih*scale)
     
643.             dx = (w-nw)//2
     
644.             dy = (h-nh)//2
     
645. 
     
646.             #---------------------------------#
     
647.             #   将图像多余的部分加上灰条
     
648.             #---------------------------------#
     
649.             image       = image.resize((nw,nh), Image.BICUBIC)
     
650.             new_image   = Image.new('RGB', (w,h), (128,128,128))
     
651.             new_image.paste(image, (dx, dy))
     
652.             image_data  = np.array(new_image, np.float32)
     
653. 
     
654.             #---------------------------------#
     
655.             #   对真实框进行调整
     
656.             #---------------------------------#
     
657.             if len(box)>0:
     
658.                 np.random.shuffle(box)
     
659.                 box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
     
660.                 box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
     
661.                 box[:, 0:2][box[:, 0:2]<0] = 0
     
662.                 box[:, 2][box[:, 2]>w] = w
     
663.                 box[:, 3][box[:, 3]>h] = h
     
664.                 box_w = box[:, 2] - box[:, 0]
     
665.                 box_h = box[:, 3] - box[:, 1]
     
666.                 box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid box
     
667. 
     
668.             return image_data, box
     
669.                
     
670.         #------------------------------------------#
     
671.         #   对图像进行缩放并且进行长和宽的扭曲
     
672.         #------------------------------------------#
     
673.         new_ar = iw/ih * self.rand(1-jitter,1+jitter) / self.rand(1-jitter,1+jitter)
     
674.         scale = self.rand(.25, 2)
     
675.         if new_ar < 1:
     
676.             nh = int(scale*h)
     
677.             nw = int(nh*new_ar)
     
678.         else:
     
679.             nw = int(scale*w)
     
680.             nh = int(nw/new_ar)
     
681.         image = image.resize((nw,nh), Image.BICUBIC)
     
682. 
     
683.         #------------------------------------------#
     
684.         #   将图像多余的部分加上灰条
     
685.         #------------------------------------------#
     
686.         dx = int(self.rand(0, w-nw))
     
687.         dy = int(self.rand(0, h-nh))
     
688.         new_image = Image.new('RGB', (w,h), (128,128,128))
     
689.         new_image.paste(image, (dx, dy))
     
690.         image = new_image
     
691. 
     
692.         #------------------------------------------#
     
693.         #   翻转图像
     
694.         #------------------------------------------#
     
695.         flip = self.rand()<.5
     
696.         if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)
     
697. 
     
698.         image_data      = np.array(image, np.uint8)
     
699.         #---------------------------------#
     
700.         #   对图像进行色域变换
     
701.         #   计算色域变换的参数
     
702.         #---------------------------------#
     
703.         r               = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
     
704.         #---------------------------------#
     
705.         #   将图像转到HSV上
     
706.         #---------------------------------#
     
707.         hue, sat, val   = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))
     
708.         dtype           = image_data.dtype
     
709.         #---------------------------------#
     
710.         #   应用变换
     
711.         #---------------------------------#
     
712.         x       = np.arange(0, 256, dtype=r.dtype)
     
713.         lut_hue = ((x * r[0]) % 180).astype(dtype)
     
714.         lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
     
715.         lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
     
716. 
     
717.         image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
     
718.         image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)
     
719. 
     
720.         #---------------------------------#
     
721.         #   对真实框进行调整
     
722.         #---------------------------------#
     
723.         if len(box)>0:
     
724.             np.random.shuffle(box)
     
725.             box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
     
726.             box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
     
727.             if flip: box[:, [0,2]] = w - box[:, [2,0]]
     
728.             box[:, 0:2][box[:, 0:2]<0] = 0
     
729.             box[:, 2][box[:, 2]>w] = w
     
730.             box[:, 3][box[:, 3]>h] = h
     
731.             box_w = box[:, 2] - box[:, 0]
     
732.             box_h = box[:, 3] - box[:, 1]
     
733.             box = box[np.logical_and(box_w>1, box_h>1)]
     
734.         
     
735.         return image_data, box
     
736.    
     
737. # DataLoader中collate_fn使用
     
738. def yolo_dataset_collate(batch):
     
739.     images = []
     
740.     bboxes = []
     
741.     for img, box in batch:
     
742.         images.append(img)
     
743.         bboxes.append(box)
     
744.     images = torch.from_numpy(np.array(images)).type(torch.FloatTensor)
     
745.     bboxes = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in bboxes]
     
746.     return images, bboxes
     
747. 
     
748. if __name__ == "__main__":
     
749.     #---------------------------------#
     
750.     #   input_shape     输入的shape大小，一定要是32的倍数
     
751.     #------------------------------------------------------#
     
752.     input_shape     = [416, 416]
     
753.     img_channel = 1
     
754.     img_size = 416
     
755.     num_epochs      = 300
     
756.     lr_init = 0.001
     
757.     num_workers         = 0
     
758.     #----------------------------------------------------#
     
759.     #   获得图片路径和标签
     
760.     #----------------------------------------------------#
     
761.     train_annotation_path   = '2007_train.txt'
     
762.     val_annotation_path     = '2007_val.txt'
     
763.    
     
764.     device          = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
765.     #----------------------------------------------------#
     
766.     #   获取classes和anchor
     
767.     #----------------------------------------------------#
     
768.     # class_names, num_classes = get_classes(classes_path)
     
769.     num_classes = 1
     
770.     class_names = ["Pos"]*num_classes
     
771.     # anchors, num_anchors     = get_anchors(anchors_path)
     
772.     anchors = np.array([[ 10.,  13.],
     
773.            [ 16.,  30.],
     
774.            [ 33.,  23.],
     
775.            [ 30.,  61.],
     
776.            [ 62.,  45.],
     
777.            [ 59., 119.],
     
778.            [116.,  90.],
     
779.            [156., 198.],
     
780.            [373., 326.]])
     
781.     num_anchors = len(anchors)
     
782.     anchors_mask    = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
     
783.    
     
784.     #------------------------------------------------------#
     
785.     #   创建yolo模型
     
786.     #------------------------------------------------------#
     
787.     model = YoloBody(anchors_mask, num_classes, phi=0, load_weights=False)
     
788.     # if not pretrained:
     
789.     #     weights_init(model)
     
790.     yolo_loss    = YOLOLoss(anchors, num_classes, input_shape, True, anchors_mask)
     
791.         
     
792.     #------------------------------------------------------------------#
     
793.     #   torch 1.2不支持amp，建议使用torch 1.7.1及以上正确使用fp16
     
794.     #   因此torch1.2这里显示"could not be resolve"
     
795.     #------------------------------------------------------------------#
     
796.    
     
797.     GPU_index = 0
     
798.     # model_train = torch.nn.DataParallel(model)
     
799.     cudnn.benchmark = True
     
800.    
     
801.     #---------------------------#
     
802.     #   读取数据集对应的txt
     
803.     #---------------------------#
     
804.     with open(train_annotation_path) as f:
     
805.         train_lines = f.readlines()
     
806.     with open(val_annotation_path) as f:
     
807.         val_lines   = f.readlines()
     
808.     num_train   = len(train_lines)
     
809.     num_val     = len(val_lines)
     
810.    
     
811.     #------------------------------------------------------#
     
812.     #-------------------------------------------------------------------#
     
813.     #   如果不冻结训练的话，直接设置batch_size为Unfreeze_batch_size
     
814.     #-------------------------------------------------------------------#
     
815.     batch_size = 8
     
816.    
     
817.     optimizer = torch.optim.Adam([dict(params=model.parameters(), lr = lr_init)])
     
818.    
     
819.     #---------------------------------------#
     
820.     #   构建数据集加载器。
     
821.     #---------------------------------------#
     
822.     train_dataset   = YoloDataset(train_lines, input_shape, num_classes, train = True)
     
823.     val_dataset     = YoloDataset(val_lines, input_shape, num_classes, train = False)
     
824.    
     
825.     train_dataloaders             = DataLoader(train_dataset, shuffle = True, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
     
826.                                 drop_last=True, collate_fn=yolo_dataset_collate, sampler=None)
     
827.     val_dataloaders         = DataLoader(val_dataset  , shuffle = True, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
     
828.                                 drop_last=True, collate_fn=yolo_dataset_collate, sampler=None)
     
829.    
     
830.    
     
831.     model.cuda(GPU_index)
     
832.     model.train()
     
833.     train_sum_loss = np.inf
     
834.     val_sum_loss = np.inf
     
835.     for epoch in range(num_epochs):
     
836.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
     
837.         # dt_size = len(train_dataloaders.dataset)
     
838.         step = 0
     
839.         epoch_loss = 0
     
840.         model.train()
     
841.         for x, y in train_dataloaders:
     
842.             step += 1
     
843.             inputs = x.to(device)
     
844.             labels = y.to(device)
     
845.             labels = [ann.cuda(0) for ann in labels]
     
846.             
     
847.             optimizer.zero_grad()
     
848.             outputs = model(inputs)
     
849.             
     
850.             loss_value_all  = 0
     
851.             for l in range(len(outputs)):
     
852.                 loss_item = yolo_loss(l, outputs[l], outputs)
     
853.                 loss_value_all  += loss_item
     
854.             loss = loss_value_all
     
855.             
     
856.             loss.backward()
     
857.             optimizer.step()
     
858.             epoch_loss += loss.item()
     
859.    
     
860.             print("%d/%d,train_loss:%0.3f,train_metric:%0.3f" % (step, (dt_size - 1) // train_dataloaders.batch_size + 1, loss.item(), epoch_metric))
     
861.             del x, y, inputs, labels, outputs
     
862.             torch.cuda.empty_cache()
     
863.         # print("epoch %d loss:%0.3f" % (epoch, epoch_loss))
     
864.         if(epoch_loss < train_sum_loss):
     
865.             train_sum_loss = epoch_loss
     
866.             model.eval()
     
867.             example = torch.rand(1, img_channel, img_size, img_size).type(torch.FloatTensor).cuda(GPU_index)
     
868.             pt_model = torch.jit.trace(model, example)
     
869.             pt_model.save( os.path.join(train_project_path, "Model", str(epoch) +'-Loss.pt') )

复制代码

参考：
【1】https://github.com/bubbliiiing/efficientnet-yolo3-pytorch

训练所需的各类权重可以在百度云下载。
链接: https://pan.baidu.com/s/1BQ9rIyxWiMIlUgwyb8ltVw
提取码: dp7r

VOC数据集下载地址如下，里面已经包括了训练集、测试集、验证集（与测试集一样），无需再次划分：
链接: https://pan.baidu.com/s/19Mw2u_df_nBzsC2lg20fQA
提取码: j5ge

【2】COCO-2014下载
https://blog.csdn.net/boy854456187/article/details/119277637

【3】人脸检测+口罩：(3条消息) 使用YOLOv5实现人脸口罩佩戴检测（详细）_Stick_2的博客-CSDN博客

【4】人脸检测：(3条消息) 用yolov5做人脸检测_yolov5人脸识别_noob_qing的博客-CSDN博客

Halcom

【0基础上手yolov3目标定位-哔哩哔哩】https://b23.tv/DRN1k4h

1. '''
   
2.     0,0 ------> x (width)
   
3.      |
   
4.      |  (Left,Top)
   
5.      |      *_________
   
6.      |      |         |
   
7.             |         |
   
8.      y      |_________|
   
9.   (height)            *
   
10.                 (Right,Bottom)
    
11. '''

复制代码

路径；Col1,Row1,Col2,Row2,Class【Class从0开始】

1. C:/DetecrionProjects/Image/1196770001008.jpg;396,121,504,219,0
   
2. C:/DetecrionProjects/Image/1196770001009.jpg;344,88,503,215,0
   
3. C:/DetecrionProjects/Image/1196770001010.jpg;359,105,510,224,0
   
4. C:/DetecrionProjects/Image/1196770001011.jpg;376,86,511,225,0
   
5. C:/DetecrionProjects/Image/1196770001012.jpg;268,75,435,191,0
   
6. C:/DetecrionProjects/Image/1196770001014.jpg;365,120,508,224,0
   
7. C:/DetecrionProjects/Image/1196770001015.jpg;314,88,501,216,0
   
8. C:/DetecrionProjects/Image/1196770001016.jpg;340,102,511,222,0

复制代码

Halcom

1. #-----------------------------------------------------------------------#
   
2. #   predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能
   
3. #   整合到了一个py文件中，通过指定mode进行模式的修改。
   
4. #-----------------------------------------------------------------------#
   
5. import time
   
6. import colorsys
   
7. import os
   
8. import cv2
   
9. import numpy as np
   
10. from PIL import Image, ImageDraw, ImageFont
    
11. import torch
    
12. import torch.nn as nn
    
13. from torchvision.ops import nms
    
14. 
    
15. from nets.yolo import YoloBody
    
16. # from utils.utils import (cvtColor, get_anchors, get_classes, preprocess_input, resize_image)
    
17. # from utils.utils_bbox import DecodeBox
    
18. 
    
19. #---------------------------------------------------------#
    
20. #   将图像转换成RGB图像，防止灰度图在预测时报错。
    
21. #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
    
22. #---------------------------------------------------------#
    
23. def cvtColor(image):
    
24.     if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
    
25.         return image
    
26.     else:
    
27.         image = image.convert('RGB')
    
28.         return image
    
29.    
    
30. #---------------------------------------------------#
    
31. #   对输入图像进行resize
    
32. #---------------------------------------------------#
    
33. def resize_image(image, size):
    
34.     iw, ih  = image.size
    
35.     w, h    = size
    
36.     new_image = image.resize((w, h), Image.BICUBIC)
    
37.     return new_image
    
38. 
    
39. #---------------------------------------------------#
    
40. #   获得类
    
41. #---------------------------------------------------#
    
42. def get_classes(classes_path):
    
43.     with open(classes_path, encoding='utf-8') as f:
    
44.         class_names = f.readlines()
    
45.     class_names = [c.strip() for c in class_names]
    
46.     return class_names, len(class_names)
    
47. 
    
48. #---------------------------------------------------#
    
49. #   获得先验框
    
50. #---------------------------------------------------#
    
51. def get_anchors(anchors_path):
    
52.     '''loads the anchors from a file'''
    
53.     with open(anchors_path, encoding='utf-8') as f:
    
54.         anchors = f.readline()
    
55.     anchors = [float(x) for x in anchors.split(',')]
    
56.     anchors = np.array(anchors).reshape(-1, 2)
    
57.     return anchors, len(anchors)
    
58. 
    
59. #---------------------------------------------------#
    
60. #   获得学习率
    
61. #---------------------------------------------------#
    
62. def get_lr(optimizer):
    
63.     for param_group in optimizer.param_groups:
    
64.         return param_group['lr']
    
65. 
    
66. def preprocess_input(image):
    
67.     image /= 255.0
    
68.     return image
    
69. 
    
70. class DecodeBox():
    
71.     def __init__(self, anchors, num_classes, input_shape, anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]):
    
72.         super(DecodeBox, self).__init__()
    
73.         self.anchors        = anchors
    
74.         self.num_classes    = num_classes
    
75.         self.bbox_attrs     = 5 + num_classes
    
76.         self.input_shape    = input_shape
    
77.         #-----------------------------------------------------------#
    
78.         #   13x13的特征层对应的anchor是[116,90],[156,198],[373,326]
    
79.         #   26x26的特征层对应的anchor是[30,61],[62,45],[59,119]
    
80.         #   52x52的特征层对应的anchor是[10,13],[16,30],[33,23]
    
81.         #-----------------------------------------------------------#
    
82.         self.anchors_mask   = anchors_mask
    
83. 
    
84.     def decode_box(self, inputs):
    
85.         outputs = []
    
86.         for i, input in enumerate(inputs):
    
87.             #-----------------------------------------------#
    
88.             #   输入的input一共有三个，他们的shape分别是
    
89.             #   batch_size, 255, 13, 13
    
90.             #   batch_size, 255, 26, 26
    
91.             #   batch_size, 255, 52, 52
    
92.             #-----------------------------------------------#
    
93.             batch_size      = input.size(0)
    
94.             input_height    = input.size(2)
    
95.             input_width     = input.size(3)
    
96. 
    
97.             #-----------------------------------------------#
    
98.             #   输入为416x416时
    
99.             #   stride_h = stride_w = 32、16、8
    
100.             #-----------------------------------------------#
     
101.             stride_h = self.input_shape[0] / input_height
     
102.             stride_w = self.input_shape[1] / input_width
     
103.             #-------------------------------------------------#
     
104.             #   此时获得的scaled_anchors大小是相对于特征层的
     
105.             #-------------------------------------------------#
     
106.             scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]
     
107. 
     
108.             #-----------------------------------------------#
     
109.             #   输入的input一共有三个，他们的shape分别是
     
110.             #   batch_size, 3, 13, 13, 85
     
111.             #   batch_size, 3, 26, 26, 85
     
112.             #   batch_size, 3, 52, 52, 85
     
113.             #-----------------------------------------------#
     
114.             prediction = input.view(batch_size, len(self.anchors_mask[i]),
     
115.                                     self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()
     
116. 
     
117.             #-----------------------------------------------#
     
118.             #   先验框的中心位置的调整参数
     
119.             #-----------------------------------------------#
     
120.             x = torch.sigmoid(prediction[..., 0])  
     
121.             y = torch.sigmoid(prediction[..., 1])
     
122.             #-----------------------------------------------#
     
123.             #   先验框的宽高调整参数
     
124.             #-----------------------------------------------#
     
125.             w = prediction[..., 2]
     
126.             h = prediction[..., 3]
     
127.             #-----------------------------------------------#
     
128.             #   获得置信度，是否有物体
     
129.             #-----------------------------------------------#
     
130.             conf        = torch.sigmoid(prediction[..., 4])
     
131.             #-----------------------------------------------#
     
132.             #   种类置信度
     
133.             #-----------------------------------------------#
     
134.             pred_cls    = torch.sigmoid(prediction[..., 5:])
     
135. 
     
136.             FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
     
137.             LongTensor  = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
     
138. 
     
139.             #----------------------------------------------------------#
     
140.             #   生成网格，先验框中心，网格左上角
     
141.             #   batch_size,3,13,13
     
142.             #----------------------------------------------------------#
     
143.             grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(
     
144.                 batch_size * len(self.anchors_mask[i]), 1, 1).view(x.shape).type(FloatTensor)
     
145.             grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(
     
146.                 batch_size * len(self.anchors_mask[i]), 1, 1).view(y.shape).type(FloatTensor)
     
147. 
     
148.             #----------------------------------------------------------#
     
149.             #   按照网格格式生成先验框的宽高
     
150.             #   batch_size,3,13,13
     
151.             #----------------------------------------------------------#
     
152.             anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
     
153.             anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
     
154.             anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
     
155.             anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)
     
156. 
     
157.             #----------------------------------------------------------#
     
158.             #   利用预测结果对先验框进行调整
     
159.             #   首先调整先验框的中心，从先验框中心向右下角偏移
     
160.             #   再调整先验框的宽高。
     
161.             #----------------------------------------------------------#
     
162.             pred_boxes          = FloatTensor(prediction[..., :4].shape)
     
163.             pred_boxes[..., 0]  = x.data + grid_x
     
164.             pred_boxes[..., 1]  = y.data + grid_y
     
165.             pred_boxes[..., 2]  = torch.exp(w.data) * anchor_w
     
166.             pred_boxes[..., 3]  = torch.exp(h.data) * anchor_h
     
167. 
     
168.             #----------------------------------------------------------#
     
169.             #   将输出结果归一化成小数的形式
     
170.             #----------------------------------------------------------#
     
171.             _scale = torch.Tensor([input_width, input_height, input_width, input_height]).type(FloatTensor)
     
172.             output = torch.cat((pred_boxes.view(batch_size, -1, 4) / _scale,
     
173.                                 conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
     
174.             outputs.append(output.data)
     
175.         return outputs
     
176. 
     
177.     def yolo_correct_boxes(self, box_xy, box_wh, input_shape, image_shape):
     
178.         #-----------------------------------------------------------------#
     
179.         #   把y轴放前面是因为方便预测框和图像的宽高进行相乘
     
180.         #-----------------------------------------------------------------#
     
181.         box_yx = box_xy[..., ::-1]
     
182.         box_hw = box_wh[..., ::-1]
     
183.         input_shape = np.array(input_shape)
     
184.         image_shape = np.array(image_shape)
     
185. 
     
186.         box_mins    = box_yx - (box_hw / 2.)
     
187.         box_maxes   = box_yx + (box_hw / 2.)
     
188.         boxes  = np.concatenate([box_mins[..., 0:1], box_mins[..., 1:2], box_maxes[..., 0:1], box_maxes[..., 1:2]], axis=-1)
     
189.         boxes *= np.concatenate([image_shape, image_shape], axis=-1)
     
190.         return boxes
     
191. 
     
192.     def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, conf_thres=0.5, nms_thres=0.4):
     
193.         #----------------------------------------------------------#
     
194.         #   将预测结果的格式转换成左上角右下角的格式。
     
195.         #   prediction  [batch_size, num_anchors, 85]
     
196.         #----------------------------------------------------------#
     
197.         box_corner          = prediction.new(prediction.shape)
     
198.         box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
     
199.         box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
     
200.         box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
     
201.         box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
     
202.         prediction[:, :, :4] = box_corner[:, :, :4]
     
203. 
     
204.         output = [None for _ in range(len(prediction))]
     
205.         for i, image_pred in enumerate(prediction):
     
206.             #----------------------------------------------------------#
     
207.             #   对种类预测部分取max。
     
208.             #   class_conf  [num_anchors, 1]    种类置信度
     
209.             #   class_pred  [num_anchors, 1]    种类
     
210.             #----------------------------------------------------------#
     
211.             class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True)
     
212. 
     
213.             #----------------------------------------------------------#
     
214.             #   利用置信度进行第一轮筛选
     
215.             #----------------------------------------------------------#
     
216.             conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()
     
217. 
     
218.             #----------------------------------------------------------#
     
219.             #   根据置信度进行预测结果的筛选
     
220.             #----------------------------------------------------------#
     
221.             image_pred = image_pred[conf_mask]
     
222.             class_conf = class_conf[conf_mask]
     
223.             class_pred = class_pred[conf_mask]
     
224.             if not image_pred.size(0):
     
225.                 continue
     
226.             #-------------------------------------------------------------------------#
     
227.             #   detections  [num_anchors, 7]
     
228.             #   7的内容为：x1, y1, x2, y2, obj_conf, class_conf, class_pred
     
229.             #-------------------------------------------------------------------------#
     
230.             detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)
     
231. 
     
232.             #------------------------------------------#
     
233.             #   获得预测结果中包含的所有种类
     
234.             #------------------------------------------#
     
235.             unique_labels = detections[:, -1].cpu().unique()
     
236. 
     
237.             if prediction.is_cuda:
     
238.                 unique_labels = unique_labels.cuda()
     
239.                 detections = detections.cuda()
     
240. 
     
241.             for c in unique_labels:
     
242.                 #------------------------------------------#
     
243.                 #   获得某一类得分筛选后全部的预测结果
     
244.                 #------------------------------------------#
     
245.                 detections_class = detections[detections[:, -1] == c]
     
246. 
     
247.                 #------------------------------------------#
     
248.                 #   使用官方自带的非极大抑制会速度更快一些！
     
249.                 #------------------------------------------#
     
250.                 keep = nms(
     
251.                     detections_class[:, :4],
     
252.                     detections_class[:, 4] * detections_class[:, 5],
     
253.                     nms_thres
     
254.                 )
     
255.                 max_detections = detections_class[keep]
     
256. 
     
257.                 # Add max detections to outputs
     
258.                 output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))
     
259.             
     
260.             if output[i] is not None:
     
261.                 output[i]           = output[i].cpu().numpy()
     
262.                 box_xy, box_wh      = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]
     
263.                 output[i][:, :4]    = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape)
     
264.         return output
     
265. 
     
266. '''
     
267. 训练自己的数据集必看注释！
     
268. '''
     
269. class YOLO(object):
     
270.     _defaults = {
     
271.         #--------------------------------------------------------------------------#
     
272.         #   使用自己训练好的模型进行预测一定要修改model_path和classes_path！
     
273.         #   model_path指向logs文件夹下的权值文件，classes_path指向model_data下的txt
     
274.         #
     
275.         #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
     
276.         #   验证集损失较低不代表mAP较高，仅代表该权值在验证集上泛化性能较好。
     
277.         #   如果出现shape不匹配，同时要注意训练时的model_path和classes_path参数的修改
     
278.         #--------------------------------------------------------------------------#
     
279.         "model_path"        : 'model_data/Valnet.pt',
     
280.         "classes_path"      : 'model_data/user.txt',
     
281.         #---------------------------------------------------------------------#
     
282.         #   anchors_path代表先验框对应的txt文件，一般不修改。
     
283.         #   anchors_mask用于帮助代码找到对应的先验框，一般不修改。
     
284.         #---------------------------------------------------------------------#
     
285.         "anchors_path"      : 'model_data/yolo_anchors.txt',
     
286.         "anchors_mask"      : [[6, 7, 8], [3, 4, 5], [0, 1, 2]],
     
287.         #---------------------------------------------------------------------#
     
288.         #   输入图片的大小，必须为32的倍数。
     
289.         #---------------------------------------------------------------------#
     
290.         "input_shape"       : [512, 512],
     
291.         #---------------------------------------------------------------------#
     
292.         #   只有得分大于置信度的预测框会被保留下来
     
293.         #---------------------------------------------------------------------#
     
294.         "confidence"        : 0.5,
     
295.         #---------------------------------------------------------------------#
     
296.         #   非极大抑制所用到的nms_iou大小
     
297.         #---------------------------------------------------------------------#
     
298.         "nms_iou"           : 0.3,
     
299.         #-------------------------------#
     
300.         #   是否使用Cuda
     
301.         #   没有GPU可以设置成False
     
302.         #-------------------------------#
     
303.         "cuda"              : True,
     
304.     }
     
305. 
     
306.     @classmethod
     
307.     def get_defaults(cls, n):
     
308.         if n in cls._defaults:
     
309.             return cls._defaults[n]
     
310.         else:
     
311.             return "Unrecognized attribute name '" + n + "'"
     
312. 
     
313.     #---------------------------------------------------#
     
314.     #   初始化YOLO
     
315.     #---------------------------------------------------#
     
316.     def __init__(self, **kwargs):
     
317.         self.__dict__.update(self._defaults)
     
318.         for name, value in kwargs.items():
     
319.             setattr(self, name, value)
     
320.             self._defaults[name] = value
     
321.             
     
322.         #---------------------------------------------------#
     
323.         #   获得种类和先验框的数量
     
324.         #---------------------------------------------------#
     
325.         self.class_names, self.num_classes  = get_classes(self.classes_path)
     
326.         self.anchors, self.num_anchors      = get_anchors(self.anchors_path)
     
327.         self.bbox_util                      = DecodeBox(self.anchors, self.num_classes, (self.input_shape[0], self.input_shape[1]), self.anchors_mask)
     
328. 
     
329.         #---------------------------------------------------#
     
330.         #   画框设置不同的颜色
     
331.         #---------------------------------------------------#
     
332.         hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
     
333.         self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
     
334.         self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
     
335.         self.generate()
     
336.         
     
337.         # show_config(**self._defaults)
     
338. 
     
339.     #---------------------------------------------------#
     
340.     #   生成模型
     
341.     #---------------------------------------------------#
     
342.     def generate(self, onnx=False):
     
343.         #---------------------------------------------------#
     
344.         #   建立yolov3模型，载入yolov3模型的权重
     
345.         #---------------------------------------------------#
     
346.         self.net    = YoloBody(self.anchors_mask, self.num_classes)
     
347.         device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
348.         # self.net.load_state_dict(torch.load(self.model_path, map_location=device))
     
349.         # self.net    = self.net.eval()
     
350.         print('{} model, anchors, and classes loaded.'.format(self.model_path))
     
351.         # if not onnx:
     
352.         #     if self.cuda:
     
353.         #         self.net = nn.DataParallel(self.net)
     
354.         #         self.net = self.net.cuda()
     
355.         self.net = torch.jit.load(self.model_path, map_location=device)
     
356.         self.net    = self.net.eval()
     
357. 
     
358.     #---------------------------------------------------#
     
359.     #   检测图片
     
360.     #---------------------------------------------------#
     
361.     def detect_image(self, image, crop = False, count = False):
     
362.         image_shape = np.array(np.shape(image)[0:2])
     
363.         #---------------------------------------------------------#
     
364.         #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
     
365.         #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
     
366.         #---------------------------------------------------------#
     
367.         image       = cvtColor(image)
     
368.         #---------------------------------------------------------#
     
369.         #   给图像增加灰条，实现不失真的resize
     
370.         #   也可以直接resize进行识别
     
371.         #---------------------------------------------------------#
     
372.         image_data  = resize_image(image, (self.input_shape[1],self.input_shape[0]))
     
373.         #---------------------------------------------------------#
     
374.         #   添加上batch_size维度
     
375.         #---------------------------------------------------------#
     
376.         image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
     
377. 
     
378.         with torch.no_grad():
     
379.             images = torch.from_numpy(image_data)
     
380.             if self.cuda:
     
381.                 images = images.cuda()
     
382.             #---------------------------------------------------------#
     
383.             #   将图像输入网络当中进行预测！
     
384.             #---------------------------------------------------------#
     
385.             outputs = self.net(images)
     
386.             outputs = self.bbox_util.decode_box(outputs)
     
387.             #---------------------------------------------------------#
     
388.             #   将预测框进行堆叠，然后进行非极大抑制
     
389.             #---------------------------------------------------------#
     
390.             results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
     
391.                         image_shape, conf_thres = self.confidence, nms_thres = self.nms_iou)
     
392.                                                    
     
393.             if results[0] is None:
     
394.                 return image
     
395. 
     
396.             top_label   = np.array(results[0][:, 6], dtype = 'int32')
     
397.             top_conf    = results[0][:, 4] * results[0][:, 5]
     
398.             top_boxes   = results[0][:, :4]
     
399.         #---------------------------------------------------------#
     
400.         #   设置字体与边框厚度
     
401.         #---------------------------------------------------------#
     
402.         font        = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
     
403.         thickness   = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1))
     
404.         #---------------------------------------------------------#
     
405.         #   图像绘制
     
406.         #---------------------------------------------------------#
     
407.         for i, c in list(enumerate(top_label)):
     
408.             predicted_class = self.class_names[int(c)]
     
409.             box             = top_boxes[i]
     
410.             score           = top_conf[i]
     
411. 
     
412.             top, left, bottom, right = box
     
413. 
     
414.             top     = max(0, np.floor(top).astype('int32'))
     
415.             left    = max(0, np.floor(left).astype('int32'))
     
416.             bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
     
417.             right   = min(image.size[0], np.floor(right).astype('int32'))
     
418. 
     
419.             label = '{} {:.2f}'.format(predicted_class, score)
     
420.             draw = ImageDraw.Draw(image)
     
421.             label_size = draw.textsize(label, font)
     
422.             label = label.encode('utf-8')
     
423.             print(label, top, left, bottom, right)
     
424.             
     
425.             if top - label_size[1] >= 0:
     
426.                 text_origin = np.array([left, top - label_size[1]])
     
427.             else:
     
428.                 text_origin = np.array([left, top + 1])
     
429. 
     
430.             for i in range(thickness):
     
431.                 draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
     
432.             draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c])
     
433.             draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
     
434.             del draw
     
435. 
     
436.         return image
     
437. 
     
438. if __name__ == "__main__":
     
439.     yolo = YOLO()
     
440.     #----------------------------------------------------------------------------------------------------------#
     
441.     #   mode用于指定测试的模式：
     
442.     #   'predict'           表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
     
443.     #-------------------------------------------------------------------------#
     
444.     #   crop                指定了是否在单张图片预测后对目标进行截取
     
445.     #   count               指定了是否进行目标的计数
     
446.     #   crop、count仅在mode='predict'时有效
     
447.     #-------------------------------------------------------------------------#
     
448.     crop            = False
     
449.     count           = False
     
450. 
     
451.     '''
     
452.     1、如果想要进行检测完的图片的保存，利用r_image.save("img.jpg")即可保存，直接在predict.py里进行修改即可。
     
453.     2、如果想要获得预测框的坐标，可以进入yolo.detect_image函数，在绘图部分读取top，left，bottom，right这四个值。
     
454.     3、如果想要利用预测框截取下目标，可以进入yolo.detect_image函数，在绘图部分利用获取到的top，left，bottom，right这四个值
     
455.     在原图上利用矩阵的方式进行截取。
     
456.     4、如果想要在预测图上写额外的字，比如检测到的特定目标的数量，可以进入yolo.detect_image函数，在绘图部分对predicted_class进行判断，
     
457.     比如判断if predicted_class == 'car': 即可判断当前目标是否为车，然后记录数量即可。利用draw.text即可写字。
     
458.     '''
     
459.     img = r'./img/street.jpg'
     
460.     img = r'./img/1.jpg'
     
461.     image = Image.open(img)
     
462. 
     
463.     r_image = yolo.detect_image(image, crop = crop, count=count)
     
464.     r_image.show()
     
465. 
     
466.    

复制代码

Halcom

cpp，torch::jit::IValue = module.forward(input)；
此时是tuple类型，而不是tensor。yolo特有的吧

1. // 输入初始化
   
2. std::vector<torch::jit::IValue> input;
   
3. input.emplace_back(tensor_image);
   
4. 
   
5. auto outputs = model.forward(input).toTuple();
   
6. // 提取三个head
   
7. std::vector<at::Tensor> output(3);
   
8. output[0] = outputs->elements()[0].toTensor().to(at::kCPU);
   
9. output[1] = outputs->elements()[1].toTensor().to(at::kCPU);
   
10. output[2] = outputs->elements()[2].toTensor().to(at::kCPU);

复制代码

Halcom

【1】低版本libtorch可用
(6条消息) libtorch (pytorch c++) 教程（七）_AllentDan的博客-CSDN博客

【2】高版本libtorch可用
https://developer.aliyun.com/article/1161305

【3】tensor块操作
(6条消息) libtorch的常见tensor操作_libtorch创建tensor_Mr.Q的博客-CSDN博客