import torch import torch.nn as nn from utils.nms import multiclass_nms from .yolov3_backbone import build_backbone from .yolov3_neck import build_neck from .yolov3_fpn import build_fpn from .yolov3_head import build_head # YOLOv3 class YOLOv3(nn.Module): def __init__(self, cfg, device, num_classes=20, conf_thresh=0.01, topk=100, nms_thresh=0.5, trainable=False): super(YOLOv3, self).__init__() # ------------------- Basic parameters ------------------- self.cfg = cfg # 模型配置文件 self.device = device # cuda或者是cpu self.num_classes = num_classes # 类别的数量 self.trainable = trainable # 训练的标记 self.conf_thresh = conf_thresh # 得分阈值 self.nms_thresh = nms_thresh # NMS阈值 self.topk = topk # topk self.stride = [8, 16, 32] # 网络的输出步长 # ------------------- Anchor box ------------------- self.num_levels = 3 self.num_anchors = len(cfg['anchor_size']) // self.num_levels self.anchor_size = torch.as_tensor( cfg['anchor_size'] ).view(self.num_levels, self.num_anchors, 2) # [S, A, 2] # ------------------- Network Structure ------------------- ## 主干网络 self.backbone, feats_dim = build_backbone( cfg['backbone'], trainable&cfg['pretrained']) ## 颈部网络: SPP模块 self.neck = build_neck(cfg, in_dim=feats_dim[-1], out_dim=feats_dim[-1]) feats_dim[-1] = self.neck.out_dim ## 颈部网络: 特征金字塔 self.fpn = build_fpn(cfg=cfg, in_dims=feats_dim, out_dim=int(256*cfg['width'])) self.head_dim = self.fpn.out_dim ## 检测头 self.non_shared_heads = nn.ModuleList( [build_head(cfg, head_dim, head_dim, num_classes) for head_dim in self.head_dim ]) ## 预测层 self.obj_preds = nn.ModuleList( [nn.Conv2d(head.reg_out_dim, 1 * self.num_anchors, kernel_size=1) for head in self.non_shared_heads ]) self.cls_preds = nn.ModuleList( [nn.Conv2d(head.cls_out_dim, self.num_classes * self.num_anchors, kernel_size=1) for head in self.non_shared_heads ]) self.reg_preds = nn.ModuleList( [nn.Conv2d(head.reg_out_dim, 4 * self.num_anchors, kernel_size=1) for head in self.non_shared_heads ]) # --------- Network Initialization ---------- self.init_yolo() def init_yolo(self): # Init yolo for m in self.modules(): if isinstance(m, nn.BatchNorm2d): m.eps = 1e-3 m.momentum = 0.03 # Init bias init_prob = 0.01 bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob)) # obj pred for obj_pred in self.obj_preds: b = obj_pred.bias.view(self.num_anchors, -1) b.data.fill_(bias_value.item()) obj_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) # cls pred for cls_pred in self.cls_preds: b = cls_pred.bias.view(self.num_anchors, -1) b.data.fill_(bias_value.item()) cls_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) def generate_anchors(self, level, fmp_size): """ fmp_size: (List) [H, W] """ fmp_h, fmp_w = fmp_size # [KA, 2] anchor_size = self.anchor_size[level] # generate grid cells anchor_y, anchor_x = torch.meshgrid([torch.arange(fmp_h), torch.arange(fmp_w)]) anchor_xy = torch.stack([anchor_x, anchor_y], dim=-1).float().view(-1, 2) # [HW, 2] -> [HW, KA, 2] -> [M, 2] anchor_xy = anchor_xy.unsqueeze(1).repeat(1, self.num_anchors, 1) anchor_xy = anchor_xy.view(-1, 2).to(self.device) # [KA, 2] -> [1, KA, 2] -> [HW, KA, 2] -> [M, 2] anchor_wh = anchor_size.unsqueeze(0).repeat(fmp_h*fmp_w, 1, 1) anchor_wh = anchor_wh.view(-1, 2).to(self.device) anchors = torch.cat([anchor_xy, anchor_wh], dim=-1) return anchors def decode_boxes(self, level, anchors, reg_pred): """ 将txtytwth转换为常用的x1y1x2y2形式。 """ # 计算预测边界框的中心点坐标和宽高 pred_ctr = (torch.sigmoid(reg_pred[..., :2]) + anchors[..., :2]) * self.stride[level] pred_wh = torch.exp(reg_pred[..., 2:]) * anchors[..., 2:] # 将所有bbox的中心带你坐标和宽高换算成x1y1x2y2形式 pred_x1y1 = pred_ctr - pred_wh * 0.5 pred_x2y2 = pred_ctr + pred_wh * 0.5 pred_box = torch.cat([pred_x1y1, pred_x2y2], dim=-1) return pred_box def post_process(self, obj_preds, cls_preds, reg_preds, anchors): """ Input: obj_preds: List(Tensor) [[H x W, 1], ...] cls_preds: List(Tensor) [[H x W, C], ...] reg_preds: List(Tensor) [[H x W, 4], ...] anchors: List(Tensor) [[H x W, 2], ...] """ all_scores = [] all_labels = [] all_bboxes = [] for level, (obj_pred_i, cls_pred_i, reg_pred_i, anchor_i) \ in enumerate(zip(obj_preds, cls_preds, reg_preds, anchors)): # (H x W x KA x C,) scores_i = (torch.sqrt(obj_pred_i.sigmoid() * cls_pred_i.sigmoid())).flatten() # Keep top k top scoring indices only. num_topk = min(self.topk, reg_pred_i.size(0)) # torch.sort is actually faster than .topk (at least on GPUs) predicted_prob, topk_idxs = scores_i.sort(descending=True) topk_scores = predicted_prob[:num_topk] topk_idxs = topk_idxs[:num_topk] # filter out the proposals with low confidence score keep_idxs = topk_scores > self.conf_thresh scores = topk_scores[keep_idxs] topk_idxs = topk_idxs[keep_idxs] anchor_idxs = torch.div(topk_idxs, self.num_classes, rounding_mode='floor') labels = topk_idxs % self.num_classes reg_pred_i = reg_pred_i[anchor_idxs] anchor_i = anchor_i[anchor_idxs] # decode box: [M, 4] bboxes = self.decode_boxes(level, anchor_i, reg_pred_i) all_scores.append(scores) all_labels.append(labels) all_bboxes.append(bboxes) scores = torch.cat(all_scores) labels = torch.cat(all_labels) bboxes = torch.cat(all_bboxes) # threshold keep_idxs = scores.gt(self.conf_thresh) scores = scores[keep_idxs] labels = labels[keep_idxs] bboxes = bboxes[keep_idxs] # to cpu & numpy scores = scores.cpu().numpy() labels = labels.cpu().numpy() bboxes = bboxes.cpu().numpy() # nms scores, labels, bboxes = multiclass_nms( scores, labels, bboxes, self.nms_thresh, self.num_classes, False) return bboxes, scores, labels @torch.no_grad() def inference(self, x): # 主干网络 pyramid_feats = self.backbone(x) # 颈部网络 pyramid_feats[-1] = self.neck(pyramid_feats[-1]) # 特征金字塔 pyramid_feats = self.fpn(pyramid_feats) # 检测头 all_anchors = [] all_obj_preds = [] all_cls_preds = [] all_reg_preds = [] for level, (feat, head) in enumerate(zip(pyramid_feats, self.non_shared_heads)): cls_feat, reg_feat = head(feat) # [1, C, H, W] obj_pred = self.obj_preds[level](reg_feat) cls_pred = self.cls_preds[level](cls_feat) reg_pred = self.reg_preds[level](reg_feat) # anchors: [M, 2] fmp_size = cls_pred.shape[-2:] anchors = self.generate_anchors(level, fmp_size) # [1, AC, H, W] -> [H, W, AC] -> [M, C] obj_pred = obj_pred[0].permute(1, 2, 0).contiguous().view(-1, 1) cls_pred = cls_pred[0].permute(1, 2, 0).contiguous().view(-1, self.num_classes) reg_pred = reg_pred[0].permute(1, 2, 0).contiguous().view(-1, 4) all_obj_preds.append(obj_pred) all_cls_preds.append(cls_pred) all_reg_preds.append(reg_pred) all_anchors.append(anchors) # post process bboxes, scores, labels = self.post_process( all_obj_preds, all_cls_preds, all_reg_preds, all_anchors) return bboxes, scores, labels def forward(self, x): if not self.trainable: return self.inference(x) else: bs = x.shape[0] # 主干网络 pyramid_feats = self.backbone(x) # 颈部网络 pyramid_feats[-1] = self.neck(pyramid_feats[-1]) # 特征金字塔 pyramid_feats = self.fpn(pyramid_feats) # 检测头 all_fmp_sizes = [] all_obj_preds = [] all_cls_preds = [] all_box_preds = [] for level, (feat, head) in enumerate(zip(pyramid_feats, self.non_shared_heads)): cls_feat, reg_feat = head(feat) # [B, C, H, W] obj_pred = self.obj_preds[level](reg_feat) cls_pred = self.cls_preds[level](cls_feat) reg_pred = self.reg_preds[level](reg_feat) fmp_size = cls_pred.shape[-2:] # generate anchor boxes: [M, 4] anchors = self.generate_anchors(level, fmp_size) # [B, AC, H, W] -> [B, H, W, AC] -> [B, M, C] obj_pred = obj_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, 1) cls_pred = cls_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, self.num_classes) reg_pred = reg_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, 4) # decode bbox box_pred = self.decode_boxes(level, anchors, reg_pred) all_obj_preds.append(obj_pred) all_cls_preds.append(cls_pred) all_box_preds.append(box_pred) all_fmp_sizes.append(fmp_size) # output dict outputs = {"pred_obj": all_obj_preds, # List [B, M, 1] "pred_cls": all_cls_preds, # List [B, M, C] "pred_box": all_box_preds, # List [B, M, 4] 'fmp_sizes': all_fmp_sizes, # List 'strides': self.stride, # List } return outputs