一、预测框解码
三个Head输出特征图进行解码
inputs=[[batch_size, 255, 13, 13],
[batch_size, 255, 26, 26,
[batch_size, 255, 52, 52]]]
def decode_box(self, inputs):outputs = []#-----------------------------------------------## 输入的input一共有三个，他们的shape分别是# batch_size, 255, 13, 13# batch_size, 255, 26, 26# batch_size, 255, 52, 52#-----------------------------------------------#for i, input in enumerate(inputs):#这里每次取出一个特征图进行处理，第一个为[batch_size, 255, 13, 13]batch_size = input.size(0)input_height = input.size(2) #13input_width = input.size(3) #13#-----------------------------------------------## 输入为416x416时# stride_h = stride_w = 32、16、8#-----------------------------------------------#stride_h = self.input_shape[0] / input_height #416/13=32stride_w = self.input_shape[1] / input_width #416/13=32#-------------------------------------------------## 此时获得的scaled_anchors大小是相对于特征层的#-------------------------------------------------##anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]#anchors=10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326#for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]#的意思为取出第[6,7,8]个anchor，即116,90, 156,198, 373,326#这三个anchor为当前13x13特征图的在原图上的anchor的宽和高#除以步长后即为在特征图上的宽和高，输出如下：#0:(3.625, 2.8125) 116,90/32#1:(4.875, 6.1875)#2:(11.65625, 10.1875)scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors[self.anchors_mask[i]]]#-----------------------------------------------## 输出prediction# batch_size, 3, 13, 13, 85# batch_size, 3, 26, 26, 85# batch_size, 3, 52, 52, 85#-----------------------------------------------#prediction = input.view(batch_size, len(self.anchors_mask[i]),self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()#-----------------------------------------------## 先验框的中心位置坐标的调整参数#-----------------------------------------------#x = torch.sigmoid(prediction[..., 0])y = torch.sigmoid(prediction[..., 1])#-----------------------------------------------## 先验框的宽高调整参数#-----------------------------------------------#w = prediction[..., 2]h = prediction[..., 3]#-----------------------------------------------## 获得置信度，是否有物体#-----------------------------------------------#conf = torch.sigmoid(prediction[..., 4])#-----------------------------------------------## 种类置信度#-----------------------------------------------#pred_cls = torch.sigmoid(prediction[..., 5:])FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensorLongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor#----------------------------------------------------------## 生成13x13的网格，先验框的中心在网格左上角即（0,0），x轴向右，y轴向下# batch_size,3,13,13#----------------------------------------------------------#grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_height, 1).repeat(batch_size * len(self.anchors_mask[i]), 1, 1).view(x.shape).type(FloatTensor)grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_width, 1).t().repeat(batch_size * len(self.anchors_mask[i]), 1, 1).view(y.shape).type(FloatTensor)#----------------------------------------------------------## 按照网格格式在每个网格上生成先验框的宽高# batch_size,3,13,13#----------------------------------------------------------#anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)#----------------------------------------------------------## 利用预测结果对先验框进行调整# 首先调整先验框的中心，从先验框中心向右下角偏移# 再调整先验框的宽高。#----------------------------------------------------------#pred_boxes = FloatTensor(prediction[..., :4].shape)#取出每个网格的前四维信息，即预测中心点坐标和预测宽高pred_boxes[..., 0] = x.data + grid_x #每个网格内移动xpred_boxes[..., 1] = y.data + grid_y #每个网格内移动ypred_boxes[..., 2] = torch.exp(w.data) * anchor_w #调整宽pred_boxes[..., 3] = torch.exp(h.data) * anchor_h#----------------------------------------------------------## 将输出结果归一化成小数的形式#----------------------------------------------------------#_scale = torch.Tensor([input_width, input_height, input_width, input_height]).type(FloatTensor)#将处理后的中心的坐标，宽高和置信度，80个类别的概率拼接，总共85output = torch.cat((pred_boxes.view(batch_size, -1, 4) / _scale,conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)#output=([1, 507, 85]) 507=13x13x3(每个网格3个框，有13x13个网格)outputs.append(output.data) #将这个特征图的结果插入，然后在循环处理另两个特征图的#输出#([1, 507, 85])#([1, 2028, 85])#([1, 8112, 85])return outputs
二、非极大值抑制
将上面输出的output进行非极大值抑制
以下为函数传入参数
torch.cat(outputs, 1)：将output输出的第二维全部拼接，即507+2028+8112
#([1, 507, 85])
#([1, 2028, 85])
#([1, 8112, 85])
self.num_classes:类别数，此处使用coco数据集总共80类
self.input_shape：网络要求输入图像尺寸，416
image_shape：输入图片实际尺寸
self.letterbox_image：false不用考虑，
conf_thres：置信度阈值0.5，
nms_thres：非极大值抑制阈值
results = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,image_shape, self.letterbox_image, conf_thres = self.confidence, nms_thres = self.nms_iou)
以下为非极大值过程代码
def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):#----------------------------------------------------------## 将预测结果的格式转换成左上角右下角的格式。# prediction [batch_size, num_anchors, 85]，此处为([1, 10647, 85])#----------------------------------------------------------#box_corner = prediction.new(prediction.shape)#转换为先验框左上角和右下角坐标#左上角x坐标为：中心点x坐标-宽的一半box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2#左上角y坐标为：中心点y坐标-高的一半box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2#右下角x坐标为：中心点x坐标+宽的一半box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2#右下角y坐标为：中心点y坐标+高的一半box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2prediction[:, :, :4] = box_corner[:, :, :4]#初始化输出列表output = [None for _ in range(len(prediction))]for i, image_pred in enumerate(prediction):#去一个batch(如果只输入一张图片，只循环一次)#----------------------------------------------------------## 对种类预测部分取max。# class_conf [num_anchors, 1] 种类置信度# class_pred [num_anchors, 1] 种类#----------------------------------------------------------##取出每个网格中置信度最高的类别#class_conf=[10647,1] 最高的置信度#class_pred=[10647,1] 最高置信度代表的类别class_conf, class_pred = torch.max(image_pred[:, 5:5 + num_classes], 1, keepdim=True)#----------------------------------------------------------## 利用置信度进行第一轮筛选#----------------------------------------------------------#conf_mask = (image_pred[:, 4] * class_conf[:, 0] >= conf_thres).squeeze()#----------------------------------------------------------## 根据置信度进行预测结果的筛选，去掉置信度小于置信度阈值的框框#----------------------------------------------------------#image_pred = image_pred[conf_mask]class_conf = class_conf[conf_mask]class_pred = class_pred[conf_mask]if not image_pred.size(0):continue#-------------------------------------------------------------------------## detections [num_anchors, 7] num_anchors为过滤后剩下的框框# 7的内容为：左上角右下角坐标 目标框框置信度 种类置信度 种类类别# x1, y1, x2, y2, obj_conf, class_conf, class_pred#-------------------------------------------------------------------------#detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)#------------------------------------------## 获得预测结果中包含的所有种类，每个种类可能有多个框框#------------------------------------------#unique_labels = detections[:, -1].cpu().unique()if prediction.is_cuda:unique_labels = unique_labels.cuda()detections = detections.cuda()for c in unique_labels:#对每个种类的多个框框进行处理，取出最好的那个框框#------------------------------------------## 获得某一类得分筛选后全部的预测结果#------------------------------------------#detections_class = detections[detections[:, -1] == c]#------------------------------------------## 使用官方自带的非极大抑制处理剩下的框#------------------------------------------#keep = nms(detections_class[:, :4],detections_class[:, 4] * detections_class[:, 5],nms_thres)max_detections = detections_class[keep]# nms原理如下：# # 按照存在物体的置信度排序# _, conf_sort_index = torch.sort(detections_class[:, 4]*detections_class[:, 5], descending=True)# detections_class = detections_class[conf_sort_index]# # 进行非极大抑制# max_detections = []# while detections_class.size(0):# # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉# max_detections.append(detections_class[0].unsqueeze(0))# if len(detections_class) == 1:# break# ious = bbox_iou(max_detections[-1], detections_class[1:])# detections_class = detections_class[1:][ious < nms_thres]# # 堆叠# max_detections = torch.cat(max_detections).data# Add max detections to outputsoutput[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))if output[i] is not None:output[i] = output[i].cpu().numpy()box_xy, box_wh = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]output[i][:, :4] = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)return output #此处为所有框的信息
2.1 IOU代码
以下为非极大值抑制中IOU计算代码
def box_area(boxes)return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])def box_iou(boxes2, boxes1)area1 = box_area(boxes1) #其余框的面积area2 = box_area(boxes2) #置信度最大框的面积lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # 交集左上角坐标rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # 交集右下角坐标wh = (rb - lt).clamp(min=0) #小于0的为0 clamp 钳；夹钳；#交集inter = wh[:, :, 0] * wh[:, :, 1]#iou = 交集/并集iou = inter / (area1[:, None] + area2 - inter)return iou
#------------------------------------------## 获得某一类得分筛选后全部的预测结果#------------------------------------------##取出每个框框代表的类别top_label = np.array(results[0][:, 6], dtype = 'int32')#取出每个框框代表的置信度top_conf = results[0][:, 4] * results[0][:, 5]#取出每个框框代表的左上角和右下角坐标top_boxes = results[0][:, :4]
若有收获，就点个赞吧
0 人点赞
