代码目录（代码包含大量注释，有疑问的可以看注释）

一、思路

1.1 YOLO对图像缩放算法的实现

1）先做一个过拟合版本的YOLOv3，再大批量增加数据，则我们需要先取我们进行训练的数据，在这里我取了4张猫和狗的图片，因为网络的输入是416x416，所以我们要先把所有的数据按比例缩放到要求的大小。

①thumbnail等比例缩放：

from PIL import Imagedef narrow_image(im):"""缩放图片成416*416:param im:  原图:return: 原图W、H、缩放比例、缩放后的416*416图片"""img = Image.new(mode="RGB", size=(416, 416), color=(128, 128, 128))     #生成416*416的灰色图片W,H = im.size             #获取原始图片宽高max_side = max(W,H)           #最大边长min_side = min(W,H)           #最短边长scale = 416/max_side          #得到缩放比例Z = max_side-min_side         #得到最大边和最短边差值im.thumbnail((416,416))   #按照最大边长缩放图片为416*416if W > H :               #如果宽大于高，填充高xy = (0,int((Z*scale)/2))else:                    #如果高大于宽，填充宽xy = (int((Z*scale)/2),0)img.paste(im, xy)       #把图片粘贴到灰色图区域，起始点为xyreturn W,H,scale,imgif __name__ == '__main__':import oslistpath=os.listdir("path")for path in listpath:image = Image.open(os.path.join("path",path))W,H,scale,img=narrow_image(image)img.save("images416/{}".format(path))# print(W,H,scale,img)

②resize实现的缩放：

from PIL import Image
import osdef pad_image(image, target_size):iw, ih = image.size  # 原始图像的尺寸w, h = target_size  # 目标图像的尺寸scale = min(float(w) / float(iw), float(h) / float(ih))  # 转换的最小比例# 保证长或宽，至少一个符合目标图像的尺寸nw = int(iw * scale)nh = int(ih * scale)image = image.resize((nw, nh), Image.BICUBIC)  # 采用双三次插值算法缩小图new_image = Image.new('RGB', target_size, (128, 128, 128))  # 生成灰色图像new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))  # 将图像填充为中间图像，两侧为灰色的样式# new_image.show()return new_imageif __name__ == '__main__':listpath= os.listdir("models")for path in listpath:image = Image.open(os.path.join("models",path))size = (416,416)new_images = pad_image(image,size)new_images.save("images416/{}".format(path))

做好的图像数据如下：

1.2 做数据集（json/xml解析）

2）有了转化为416x416的图像数据，我们就可以来进行框图打标签了，在这里我选择的是精灵标注助手，标注好之后，选择导出（一般选择xml和json），我这里选择json格式导出，得到导出的数据outputs如下所示：

{"path":"C:\\Users\\wdvbji\\Desktop\\images416\\1.jpg","outputs":{"object":[{"name":"1","bndbox":{"xmin":17,"ymin":71,"xmax":215,"ymax":295}},{"name":"0","bndbox":{"xmin":214,"ymin":114,"xmax":399,"ymax":321}}]},"time_labeled":1588260234463,"labeled":true,"size":{"width":416,"height":416,"depth":3}}
{"path":"C:\\Users\\wdvbji\\Desktop\\images416\\2.jpg","outputs":{"object":[{"name":"0","bndbox":{"xmin":4,"ymin":125,"xmax":233,"ymax":410}},{"name":"1","bndbox":{"xmin":141,"ymin":21,"xmax":411,"ymax":408}}]},"time_labeled":1588260263416,"labeled":true,"size":{"width":416,"height":416,"depth":3}}
{"path":"C:\\Users\\wdvbji\\Desktop\\images416\\3.jpg","outputs":{"object":[{"name":"1","bndbox":{"xmin":41,"ymin":102,"xmax":235,"ymax":356}},{"name":"0","bndbox":{"xmin":214,"ymin":200,"xmax":407,"ymax":353}}]},"time_labeled":1588260279045,"labeled":true,"size":{"width":416,"height":416,"depth":3}}
{"path":"C:\\Users\\wdvbji\\Desktop\\images416\\4.jpg","outputs":{"object":[{"name":"1","bndbox":{"xmin":10,"ymin":42,"xmax":249,"ymax":391}},{"name":"0","bndbox":{"xmin":209,"ymin":161,"xmax":399,"ymax":395}}]},"time_labeled":1588260305125,"labeled":true,"size":{"width":416,"height":416,"depth":3}}

①代码解析json格式，得到我们需要的标签：（json2txt.py）

import os
import json
from cfg import *LABEL_FILE_PATH = "data/labels.txt"
IMG_BASE_DIR = "data/images416"with open(LABEL_FILE_PATH,"w+") as f:for json_file in os.listdir("data/outputs"):with open(os.path.join("data/outputs",json_file),"r",encoding="utf-8") as jf:j = json.load(jf)print(j)pic = j["path"].split("\\")[-1]boxes = j["outputs"]["object"]wh= j["size"]w,h = wh["width"],wh["height"]w_scale , h_scale = w/IMG_WIDTH , h/IMG_HEIGHTf.write(pic)for box in boxes:bndbox = box["bndbox"]_x1, _y1, _x2, _y2 = bndbox["xmin"], bndbox["ymin"], bndbox["xmax"], bndbox["ymax"]_w, _h = _x2 - _x1, _y2 - _y1_w0_5, _h0_5 = _w / 2, _h / 2_cx, _cy = _x1 + _w0_5, _y1 + _h0_5x1, y1, w, h = int(_cx / w_scale), int(_cy / h_scale), int(_w / w_scale), int(_h / h_scale)f.write(" {} {} {} {} {}".format(str(box["name"]), x1, y1, w, h))f.write("\n")f.flush()

②代码解析json格式，得到我们需要的标签：（parse_xml.py）

from xml.dom.minidom import parse
xml_doc = r"./data/1.xml"dom = parse(xml_doc)
root = dom.documentElement
img_name = root.getElementsByTagName("path")[0].childNodes[0].data
img_size= root.getElementsByTagName("size")[0]
objects = root.getElementsByTagName("object")
img_w = img_size.getElementsByTagName("width")[0].childNodes[0].data
img_h = img_size.getElementsByTagName("height")[0].childNodes[0].data
img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
print(img_name)
print(img_w,img_h,img_c)
for box in objects:cls_name = box.getElementsByTagName("name")[0].childNodes[0].datax1 = int(box.getElementsByTagName("xmin")[0].childNodes[0].data)y1 = int(box.getElementsByTagName("ymin")[0].childNodes[0].data)x2 = int(box.getElementsByTagName("xmax")[0].childNodes[0].data)y2 = int(box.getElementsByTagName("ymax")[0].childNodes[0].data)print(cls_name,x1,y1,x2,y2)cls_name = box.getElementsByTagName("name")[1].childNodes[0].datax1 = int(box.getElementsByTagName("xmin")[1].childNodes[0].data)y1 = int(box.getElementsByTagName("ymin")[1].childNodes[0].data)x2 = int(box.getElementsByTagName("xmax")[1].childNodes[0].data)y2 = int(box.getElementsByTagName("ymax")[1].childNodes[0].data)print(cls_name,x1,y1,x2,y2)

得到的标签labels.txt为：

大工告成！当然我们可以检测一下我们的标签是否是正确的：（测试很好，可以下一步了）

from PIL import Image
from PIL import ImageDraw
import osIMG_DIR = r"data/images416"
ANO_DIR = r"data/labels.txt"img = Image.open(os.path.join(IMG_DIR, "2.jpg"))
imgDraw = ImageDraw.Draw(img)
imgDraw.rectangle((4,  125,  233,  410),outline="red")
img.show()

1.3 配置文件（anchor box等）

3）在这个部分，我们需要有coco数据集聚类得到的anchor box（也可以使用K-Means自己聚类自己的数据集）、anchor box面积的计算方式、图像宽高、类别等。

IMG_HEIGHT = 416
IMG_WIDTH = 416
CLASS_NUM = 2"anchor box是对coco数据集聚类获得"
ANCHORS_GROUP_KMEANS = {52: [[10,13],  [16,30],  [33,23]],26: [[30,61],  [62,45],  [59,119]],13: [[116,90],  [156,198],  [373,326]]}ANCHORS_GROUP = {13: [[360, 360], [360, 180], [180, 360]],26: [[180, 180], [180, 90], [90, 180]],52: [[90, 90], [90, 45], [45, 90]]}ANCHORS_GROUP_AREA = {13: [x * y for x, y in ANCHORS_GROUP[13]],26: [x * y for x, y in ANCHORS_GROUP[26]],52: [x * y for x, y in ANCHORS_GROUP[52]],
}
if __name__ == '__main__':for feature_size, anchors in ANCHORS_GROUP.items():print(feature_size)print(anchors)for feature_size, anchor_area in ANCHORS_GROUP_AREA.items():print(feature_size)print(anchor_area)

1.4 造数据集（Mydataset.py）

类别使用one_hot编码，属于当前类的编码为1即可。

import torch
from torch.utils.data import Dataset,DataLoader
import torchvision
import numpy as np
import cfg
import os
from PIL import Image
import math# 标签路径
LABEL_FILE_PATH = "data/person_label.txt"
# 图片路径
IMG_BASE_DIR = "data"
# 预处理变tensor
transforms = torchvision.transforms.Compose([torchvision.transforms.ToTensor()
])
# 对类别进行one-hot编码
def one_hot(cls_num, i):b = np.zeros(cls_num)b[i] = 1.return bclass MyDataset(Dataset):def __init__(self):# with open()自动调用close，不占内存with open(LABEL_FILE_PATH) as f:# 调用readlines()一次读取所有内容并按行返回listself.dataset = f.readlines()def __len__(self):# 返回数据集的长度return len(self.dataset)def __getitem__(self, index):# 定义标签的格式为字典labels = {}# 根据索引取到数据line = self.dataset[index]# 切割数据，返回字符串strs = line.split()# 打开图片_img_data = Image.open(os.path.join(IMG_BASE_DIR, strs[0]))# 经处理得到图片数据img_data = transforms(_img_data)# 得到所有框# _boxes = np.array(float(x) for x in strs[1:])_boxes = np.array(list(map(float, strs[1:])))# 5个元素分为一组boxes = np.split(_boxes, len(_boxes) // 5)# 遍历尺寸和建议框for feature_size, anchors in cfg.ANCHORS_GROUP.items():# 每个尺寸下建议框用零矩阵填充labels[feature_size] = np.zeros(shape=(feature_size, feature_size, 3, 5 + cfg.CLASS_NUM))# 遍历类别和中心点、宽高for box in boxes:cls, cx, cy, w, h = box# 得到中心点的偏移量和索引cx_offset, cx_index = math.modf(cx * feature_size / cfg.IMG_WIDTH)cy_offset, cy_index = math.modf(cy * feature_size / cfg.IMG_WIDTH)# 遍历每个尺寸下的每个建议框for i, anchor in enumerate(anchors):# 建议框的面积anchor_area = cfg.ANCHORS_GROUP_AREA[feature_size][i]# 宽高的偏移量p_w, p_h = w / anchor[0], h / anchor[1]# 实际框的面积p_area = w * h# 计算置信度（用IOU代表），最小面积/最大面积得到的值在[0,1]之间iou = min(p_area, anchor_area) / max(p_area, anchor_area)# 每个尺寸下每个建议框的标签labels[feature_size][int(cy_index), int(cx_index), i] = np.array([iou, cx_offset, cy_offset, np.log(p_w), np.log(p_h), *one_hot(cfg.CLASS_NUM, int(cls))])#10,ireturn labels[13], labels[26], labels[52], img_data# 返回的是每个尺寸下的标签和图片数据
if __name__ == '__main__':x=one_hot(10,2)print(x)data = MyDataset()dataloader = DataLoader(data,3,shuffle=True)# for i,x in enumerate(dataloader):#     print(x[0].shape)#     print(x[1].shape)#     print(x[2].shape)#     print(x[3].shape)for target_13, target_26, target_52, img_data in dataloader:print(target_13.shape)print(target_26.shape)print(target_52.shape)print(img_data.shape)

1.5 搭建网络（model.py）

这部分根据网络结构图完成代码。
网络结构梳理：
1）卷积_BN_激活，我们可以做成一个类（反复使用，简化代码），LeakyReLU的斜率为0.1；
2）下采样层做成一个类，kernel_size=3，stride=2，padding=1；
2）残差层做成一个类，先缩小为通道的一半（kernel_size=1，stride=1，padding=0），再扩通道到原来的大小（kernel_size=3，stride=1，padding=1）
3）上采样层做成一个类，选择最邻近插值法，比例为2；
4）卷积块做成一个类，1x1（kernel_size=1，stride=1，padding=0），3x3（kernel_size=3，stride=1，padding=1）
主网络梳理：
①主干52：1个卷积层，1个下采样层，1个残差层，1个下采样层，2个残差层，1个下采样层，8个残差层
②主干26：1个下采样层，8个残差层
③主干13：1个下采样层，4个残差层
④卷积块13：1024–>512
⑤侦测13：3x3（kernel_size=3，stride=1，padding=1），1x1（kernel_size=1，stride=1，padding=0）
⑥上采样26：512–>256（kernel_size=1，stride=1，padding=0）+上采样层
⑦卷积块26：768–>256
⑧侦测26：3x3（kernel_size=3，stride=1，padding=1），1x1（kernel_size=1，stride=1，padding=0）
⑨上采样52：256–>128（kernel_size=1，stride=1，padding=0）+上采样层
⑩卷积块52：384–>128
⑪侦测52：3x3（kernel_size=3，stride=1，padding=1），1x1（kernel_size=1，stride=1，padding=0）

import torch
import torch.nn as nn
import torch.nn.functional as Fclass ConvolutionalLayer(nn.Module):def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=False):super().__init__()self.sub_module = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),nn.BatchNorm2d(out_channels),nn.LeakyReLU(0.1))def forward(self, x):return self.sub_module(x)class DownsampleLayer(nn.Module):def __init__(self, in_channels, out_channels):super().__init__()self.sub_module = nn.Sequential(ConvolutionalLayer(in_channels, out_channels, 3, 2, 1))def forward(self, x):return self.sub_module(x)class ResidualLayer(nn.Module):def __init__(self, in_channels):super().__init__()self.sub_moduel = nn.Sequential(ConvolutionalLayer(in_channels, in_channels // 2, 1, 1, 0),ConvolutionalLayer(in_channels // 2, in_channels, 3, 1, 1))def forward(self, x):return x + self.sub_moduel(x)class UpsampleLayer(nn.Module):def __init__(self):super().__init__()def forward(self, x):return F.interpolate(x, scale_factor=2, mode="nearest")class ConvolutionalSet(nn.Module):def __init__(self, in_channels, out_channels):super().__init__()self.sub_module = nn.Sequential(ConvolutionalLayer(in_channels, out_channels, 1, 1, 0),ConvolutionalLayer(out_channels, in_channels, 3, 1, 1),ConvolutionalLayer(in_channels, out_channels, 1, 1, 0),ConvolutionalLayer(out_channels, in_channels, 3, 1, 1),ConvolutionalLayer(in_channels, out_channels, 1, 1, 0),)def forward(self, x):return self.sub_module(x)class MainNet(nn.Module):def __init__(self):super().__init__()self.trunk52 = nn.Sequential(ConvolutionalLayer(3, 32, 3, 1, 1),DownsampleLayer(32, 64),ResidualLayer(64),DownsampleLayer(64, 128),ResidualLayer(128),ResidualLayer(128),DownsampleLayer(128, 256),ResidualLayer(256),ResidualLayer(256),ResidualLayer(256),ResidualLayer(256),ResidualLayer(256),ResidualLayer(256),ResidualLayer(256),ResidualLayer(256),)self.trunk26 = nn.Sequential(DownsampleLayer(256, 512),ResidualLayer(512),ResidualLayer(512),ResidualLayer(512),ResidualLayer(512),ResidualLayer(512),ResidualLayer(512),ResidualLayer(512),ResidualLayer(512))self.trunk13 = nn.Sequential(DownsampleLayer(512, 1024),ResidualLayer(1024),ResidualLayer(1024),ResidualLayer(1024),ResidualLayer(1024))self.convset_13 = nn.Sequential(ConvolutionalSet(1024, 512))self.detection_13 = nn.Sequential(ConvolutionalLayer(512, 1024, 3, 1, 1),ConvolutionalLayer(1024, 21, 1, 1, 0))self.up_26 = nn.Sequential(ConvolutionalLayer(512,256,1,1,0),UpsampleLayer())self.convset_26 = nn.Sequential(ConvolutionalSet(768,256))self.detcetion_26 = nn.Sequential(ConvolutionalLayer(256, 512, 3, 1, 1),nn.Conv2d(512, 21, 1, 1, 0))self.up_52 = nn.Sequential(ConvolutionalLayer(256, 128, 1, 1, 0),UpsampleLayer())self.convset_52 = nn.Sequential(ConvolutionalSet(384, 128))self.detcetion_52 = nn.Sequential(ConvolutionalLayer(128, 256, 3, 1, 1),nn.Conv2d(256, 21, 1, 1, 0))def forward(self,x):h_52 = self.trunk52(x)h_26 = self.trunk26(h_52)h_13 = self.trunk13(h_26)convset_out_13 = self.convset_13(h_13)detection_out_13 = self.detection_13(convset_out_13)up_out_26 = self.up_26(convset_out_13)route_out_26 = torch.cat((up_out_26,h_26),dim=1)convset_out_26 = self.convset_26(route_out_26)detection_out_26 = self.detcetion_26(convset_out_26)up_out_52 = self.up_52(convset_out_26)route_out_52 = torch.cat((up_out_52, h_52), dim=1)convset_out_52 = self.convset_52(route_out_52)detection_out_52 = self.detcetion_52(convset_out_52)return detection_out_13, detection_out_26, detection_out_52if __name__ == '__main__':trunk = MainNet()x = torch.randn([2, 3, 416, 416], dtype=torch.float32)# 测试网络y_13, y_26, y_52 = trunk(x)print(y_13.shape)print(y_26.shape)print(y_52.shape)print(y_13.view(-1, 3, 7, 13, 13).shape)

1.6 训练（Train.py）

from Mydataset import MyDataset
from model import *
import torch
from torch.utils.data import DataLoader
import osdef loss_fn(output, target, alpha):conf_loss_fn = nn.BCEWithLogitsLoss()crood_loss_fn = nn.MSELoss()cls_loss_fn = nn.CrossEntropyLoss()output = output.permute(0, 2, 3, 1)output = output.reshape(output.size(0), output.size(1), output.size(2), 3, -1)output = output.cpu().double()mask_obj = target[..., 0] > 0output_obj = output[mask_obj]target_obj = target[mask_obj]loss_obj_conf = conf_loss_fn(output_obj[:, 0], target_obj[:, 0])loss_obj_crood = crood_loss_fn(output_obj[:, 1:5], target_obj[:, 1:5])loss_obj_cls = conf_loss_fn(output_obj[:, 5:], target_obj[:, 5:])loss_obj = loss_obj_conf + loss_obj_crood + loss_obj_clsmask_noobj = target[..., 0] == 0output_nonbj = output[mask_noobj]target_noobj = target[mask_noobj]loss_noobj = conf_loss_fn(output_nonbj[:, 0], target_noobj[:, 0])loss = alpha * loss_obj + (1 - alpha) * loss_noobjreturn lossclass Trainer:def __init__(self):self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")self.net = MainNet().to(self.device)self.opt = torch.optim.Adam(self.net.parameters())def train(self):save_path = "models/net_yolo.pth"if not os.path.exists("models"):os.mkdir("models")if os.path.exists(save_path):self.net.load_state_dict(torch.load(save_path))else:print("NO Param")BATCH_SIZE = 4EPOCH = 0mydataset = MyDataset()train_loader = DataLoader(mydataset, batch_size=BATCH_SIZE, shuffle=True)while True:for target_13, target_26, target_52, img_data in train_loader:img_data = img_data.to(self.device)output_13, output_26, output_52 = self.net(img_data)loss_13 = loss_fn(output_13, target_13, 0.9)loss_26 = loss_fn(output_26, target_26, 0.9)loss_52 = loss_fn(output_52, target_52, 0.9)loss = loss_13 + loss_26 + loss_52self.opt.zero_grad()loss.backward()self.opt.step()if EPOCH % 10 == 0:torch.save(self.net.state_dict(),save_path)print("save {}".format(EPOCH))print(loss.item())EPOCH += 1if __name__ == '__main__':t = Trainer()t.train()

1.7 工具iou,nms（tool.py）

import numpy as np
import torchdef ious(box, boxes, isMin = False):box_area = (box[3] - box[1]) * (box[4] - box[2])area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 4] - boxes[:, 2])xx1 = torch.max(box[1], boxes[:, 1])yy1 = torch.max(box[2], boxes[:, 2])xx2 = torch.min(box[3], boxes[:, 3])yy2 = torch.min(box[4], boxes[:, 4])w = torch.clamp(xx2 - xx1, min=0)h = torch.clamp(yy2 - yy1, min=0)inter = w * hovr2 = inter/ (box_area + area - inter)return ovr2def nms(boxes, thresh=0.3, isMin = True):if boxes.shape[0] == 0:return np.array([])_boxes = boxes[(-boxes[:, 0]).argsort()]r_boxes = []while _boxes.shape[0] > 1:a_box = _boxes[0]b_boxes = _boxes[1:]r_boxes.append(a_box)index = np.where(ious(a_box, b_boxes,isMin) < thresh)_boxes = b_boxes[index]if _boxes.shape[0] > 0:r_boxes.append(_boxes[0])return torch.stack(r_boxes)if __name__ == '__main__':# a = np.array([1,1,11,11])# bs = np.array([[1,1,10,10],[11,11,20,20]])# print(iou(a,bs))bs = torch.tensor([[1, 1, 10, 10, 40,8], [1, 1, 9, 9, 10,9], [9, 8, 13, 20, 15,3], [6, 11, 18, 17, 13,2]])# print(bs[:,3].argsort())print(nms(bs))

1.8 侦测（detector.py）

from model import MainNet
import cfg
import torch
import numpy as np
import PIL.Image as pimg
import PIL.ImageDraw as draw
import toolclass Detector(torch.nn.Module):def __init__(self,save_path):super(Detector, self).__init__()# 实例化网络self.net = MainNet()# 加载保存的参数self.net.load_state_dict(torch.load(save_path))# 固定BN,Dropoutself.net.eval()def forward(self, input, thresh, anchors):# 输入图像传到网络 得到预测output_13, output_26, output_52 = self.net(input)# 预测，阈值传到过滤器，得到索引和置信度，偏移量，类别idxs_13, vecs_13 = self._filter(output_13, thresh)# 索引和15个值经过分析，得到13的预测（置信度，偏移量，类别）boxes_13 = self._parse(idxs_13, vecs_13, 32, anchors[13])idxs_26, vecs_26 = self._filter(output_26, thresh)boxes_26 = self._parse(idxs_26, vecs_26, 16, anchors[26])idxs_52, vecs_52 = self._filter(output_52, thresh)boxes_52 = self._parse(idxs_52, vecs_52, 8, anchors[52])return torch.cat([boxes_13, boxes_26, boxes_52], dim=0)def _filter(self, output, thresh):# 输出形状是（NCHW），变形为（NHWC）output = output.permute(0, 2, 3, 1)# NHWC变形为NHW3(5+类别)output = output.reshape(output.size(0), output.size(1), output.size(2), 3, -1)# 变形后，最后一个维度的第0个元素是置信度，置信度大于阈值的为有物体的掩码mask = output[..., 0] > thresh# 根据掩码获取索引idxs = mask.nonzero()# 花式索引取到符合有物体的(5+类别)个值vecs = output[mask]return idxs, vecs# 返回索引和(5+类别)个值def _parse(self, idxs, vecs, t, anchors):# 建议框变tensoranchors = torch.Tensor(anchors)# 得到NHW3中3的索引a = idxs[:, 3]# 置信度为第0个元素confidence = vecs[:, 0]# 类别为第5个元素到最后_classify = vecs[:, 5:]# 如果类别的长度为0，返回一个空tensor，否则取最大值的索引if len(_classify) == 0:classify = torch.Tensor([])else:classify = torch.argmax(_classify, dim=1).float()# 实际框中心点坐标 = (索引+相对于当前格子左上角的偏移量)*(图像大小/分的格子数)cy = (idxs[:, 1].float() + vecs[:, 2]) * tcx = (idxs[:, 2].float() + vecs[:, 1]) * t# 实际框宽高 = 建议框的宽高*e^(宽高的偏移量)w = anchors[a, 0] * torch.exp(vecs[:, 3])h = anchors[a, 1] * torch.exp(vecs[:, 4])# 算出左上右下坐标，并堆叠起来x1 = cx - w / 2y1 = cy - h / 2x2 = x1 + wy2 = y1 + hout = torch.stack([confidence,x1,y1,x2,y2,classify], dim=1)return outif __name__ == '__main__':save_path = "models/net_yolo.pth"detector = Detector(save_path)# y = detector(torch.randn(3, 3, 416, 416), 0.3, cfg.ANCHORS_GROUP)# print(y.shape)img1 = pimg.open(r'data\image\1.jpg')img = img1.convert('RGB')img = np.array(img) / 255img = torch.Tensor(img)img = img.unsqueeze(0)img = img.permute(0, 3, 1, 2)img = img.cuda()out_value = detector(img, 0.3, cfg.ANCHORS_GROUP)boxes = []for j in range(10):classify_mask = (out_value[..., -1] == j)_boxes = out_value[classify_mask]boxes.append(tool.nms(_boxes))for box in boxes:try:img_draw = draw.ImageDraw(img1)c,x1, y1, x2, y2 = box[0, 0:5]print(c,x1, y1, x2, y2)img_draw.rectangle((x1, y1, x2, y2))except:continueimg1.show()

二、总结

如果在小数据量上，过拟合测试效果不错，那么我们就可以加大数据量进行训练 。