from __future__ import division import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import numpy as np import cv2 # 获取任意给定图像中存在的类别 def unique(tensor): tensor_np =tensor.cpu().numpy() unique_np = np.unique(tensor_np) unique_tensor = torch.from_numpy(unique_np) tensor_res = tensor.new(unique_tensor.shape) tensor_res.copy_(unique_tensor) return tensor_res # 计算两个边界框的IoU def bbox_iou(box1, box2): # 获取边框的坐标 b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3] b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3] # 获取交叉矩形的坐标 inter_rect_x1 = torch.max(b1_x1, b2_x1) inter_rect_y1 = torch.max(b1_y1, b2_y1) inter_rect_x2 = torch.min(b1_x2, b2_x2) inter_rect_y2 = torch.min(b1_y2, b2_y2) # 交叉面积 inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0) # 合并面积 b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) union_area = b1_area + b2_area - inter_area # IoU iou = inter_area / union_area return iou # 把检测特征图转换成二维张量，张量的每一行对应边界框的属性，5个参数：输出，输入图像的维度…… def predict_transform(prediction, inp_dim, anchors, num_classes, CUDA=True): batch_size = prediction.size(0) stride = inp_dim // prediction.size(2) grid_size = inp_dim // stride bbox_attrs = 5 + num_classes num_anchors = len(anchors) prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size) prediction = prediction.transpose(1,2).contiguous() prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs) # 锚点的维度与net块的h和w属性一致，输入图像的维度和检测图的维度之商就是步长，用检测特征图的步长分割锚点 anchors = [(a[0]/stride, a[1]/stride) for a in anchors] # 对（x，y）坐标和objectness分数执行Sigmoid函数操作 prediction[:,:,0] = torch.sigmoid(prediction[:,:,0]) prediction[:,:,1] = torch.sigmoid(prediction[:,:,1]) prediction[:,:,4] = torch.sigmoid(prediction[:,:,4]) # 将网格偏移添加到中心坐标预测中 grid = np.arange(grid_size) a,b = np.meshgrid(grid, grid) x_offset = torch.FloatTensor(a).view(-1,1) y_offset = torch.FloatTensor(b).view(-1,1) if CUDA: x_offset = x_offset.cuda() y_offset = y_offset.cuda() x_y_offset = torch.cat((x_offset, y_offset), 1).repeat(1,num_anchors).view(-1,2).unsqueeze(0) prediction[:,:,:2] += x_y_offset # 将锚点应用到边界框维度中 anchors = torch.FloatTensor(anchors) if CUDA: anchors = anchors.cuda() anchors = anchors.repeat(grid_size*grid_size, 1).unsqueeze(0) prediction[:,:,2:4] = torch.exp(prediction[:,:,2:4])*anchors # 将sigmoid激活函数应用到类别分数中 prediction[:,:,5:5 + num_classes] = torch.sigmoid((prediction[:,:,5:5 + num_classes])) # 将检测图的大小调整到与输入图像大小一致，乘以stride变量（边界框属性根据特征图大小而定） prediction[:,:,:4] *= stride return prediction # 加载类别，返回字典——将每个类别的索引映射到其名称的字符串 def load_classes(namesfile): fp = open(namesfile, "r") names = fp.read().split("\n")[:-1] return names # 输出满足objectness分数阈值和非极大值抑制(NMS)，得到真实检测结果 def write_results(prediction, confidence, num_classes, nms_conf=0.4): # 输入为预测结果，置信度，类别数，NMS阈值 # 低于objectness分数的每个边界框，其每个属性值都置0，即一整行。 conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2) prediction = prediction*conf_mask # 每个框的两个对焦坐标更容易计算两个框的IoU，故将(中心x，中心y，高度，宽度)属性转化成(左上角x，左上角y，右下角x，右下角y) box_a = prediction.new(prediction.shape) box_a[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2) box_a[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2) box_a[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2) box_a[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2) prediction[:,:,:4] = box_a[:,:,:4] batch_size = prediction.size(0) #output = prediction.new(1, prediction.size(2) + 1) write = False # 标识尚未初始化输出 # 在第一个维度即bacth上循环，一次完成一个图像的置信度阈值和NMS for ind in range(batch_size): # 获取图像,10647x85 image_pred = prediction[ind] # 每个边界框行有85个属性，其中80个类别分数，只取最大值的类别分数 # 获取具有最高分数的类及其索引 max_conf, max_conf_score = torch.max(image_pred[:,5:5+num_classes], 1) max_conf = max_conf.float().unsqueeze(1) max_conf_score = max_conf_score.float().unsqueeze(1) # 删除80个分类分数，增加最高分数类别的索引及最高分数 seq = (image_pred[:,:5], max_conf, max_conf_score) image_pred = torch.cat(seq, 1) # 删除objectness置信度小于阈值的置0条目,try-except处理无检测结果的情况，continue跳过对本图像的循环 non_zero_ind = torch.nonzero(image_pred[:,4]) try: image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7) # 7列 except: continue # PyTorch 0.4兼容 if image_pred_.shape[0] == 0: continue # 获得一个图像的所有种类 img_classes = unique(image_pred_[:,-1]) # 按类别执行NMS for cls in img_classes: # 得到一个类别的所有检测 cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1) class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze() image_pred_class = image_pred_[class_mask_ind].view(-1,7) # 对所有检测排序,按照objectness置信度 conf_sort_index = torch.sort(image_pred_class[:,4], descending=True)[1] image_pred_class = image_pred_class[conf_sort_index] idx = image_pred_class.size(0) # 对于每一个检测,执行NMS for i in range(idx): # 获取正在查看的box之后所有boxes的IoUs try: ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:]) except ValueError: # image_pred_class[i+1,:]返回空张量 break except IndexError: # image_pred_class移除部分后，idx索引越界 break # 清除IoU>阈值的检测 iou_mask = (ious < nms_conf).float().unsqueeze(1) image_pred_class[i+1:] *= iou_mask # 移除0条目 non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze() image_pred_class = image_pred_class[non_zero_ind].view(-1,7) batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind) seq = batch_ind, image_pred_class if not write: output = torch.cat(seq, 1) write = True else: out = torch.cat(seq, 1) output = torch.cat((output, out)) # 输出一个形状为Dx8的张量；其中D是所有图像中的「真实」检测结果，每个都用一行表示。 # 每一个检测结果都有8个属性，即该检测结果所属的batch中图像的索引、4个对角的坐标、objectness分数、有最大置信度的类别的分数、该类别的索引。 try: return output except: return 0 # 使用填充调整具有不变长宽性的图像 def letterbox_image(img, inp_dim): img_w, img_h = img.shape[1], img.shape[0] w, h = inp_dim new_w = int(img_w * min(w/img_w, h/img_h)) new_h = int(img_h * min(w/img_w, h/img_h)) resized_image = cv2.resize(img, (new_w,new_h), interpolation = cv2.INTER_CUBIC) canvas = np.full((inp_dim[1], inp_dim[0], 3), 128) canvas[(h-new_h)//2:(h-new_h)//2 + new_h,(w-new_w)//2:(w-new_w)//2 + new_w, :] = resized_image return canvas # 将numpy数组转换成PyTorch的的输入格式 # OpenCV将图像载入成numpy数组，颜色通道为BGR。 # PyTorch的图像输入格式是(batch x 通道 x 高度 x 宽度)，通道顺序RGB。 def prep_image(img, inp_dim): img = letterbox_image(img, (inp_dim, inp_dim)) # 转换格式大小 img = img[:,:,::-1].transpose((2,0,1)).copy() # BGR -> RGB(起止位置省略，步长为-1，负:从右往左)) | H x W x C -> C x H x W img = torch.from_numpy(img).float().div(255.0).unsqueeze(0) return img