first commit

ssd/modeling/box_head/__init__.py

@@ -0,0 +1,8 @@
+from ssd.modeling import registry
+from .box_head import SSDBoxHead
+
+__all__ = ['build_box_head', 'SSDBoxHead']
+
+
+def build_box_head(cfg):
+    return registry.BOX_HEADS[cfg.MODEL.BOX_HEAD.NAME](cfg)

ssd/modeling/box_head/box_head.py

@@ -0,0 +1,49 @@
+from torch import nn
+import torch.nn.functional as F
+
+from ssd.modeling import registry
+from ssd.modeling.anchors.prior_box import PriorBox
+from ssd.modeling.box_head.box_predictor import make_box_predictor
+from ssd.utils import box_utils
+from .inference import PostProcessor
+from .loss import MultiBoxLoss
+
+
+@registry.BOX_HEADS.register('SSDBoxHead')
+class SSDBoxHead(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.predictor = make_box_predictor(cfg)
+        self.loss_evaluator = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO) #caculate loss
+        self.post_processor = PostProcessor(cfg)
+        self.priors = None
+
+    def forward(self, features, targets=None):
+        cls_logits, bbox_pred = self.predictor(features)  #返回所有卷积结果的tensor合并
+        if self.training:
+            return self._forward_train(cls_logits, bbox_pred, targets)
+        else:
+            return self._forward_test(cls_logits, bbox_pred)
+
+    def _forward_train(self, cls_logits, bbox_pred, targets):
+        gt_boxes, gt_labels = targets['boxes'], targets['labels']
+        reg_loss, cls_loss = self.loss_evaluator(cls_logits, bbox_pred, gt_labels, gt_boxes)
+        loss_dict = dict(
+            reg_loss=reg_loss,
+            cls_loss=cls_loss,
+        )
+        detections = (cls_logits, bbox_pred)
+        return detections, loss_dict
+
+    def _forward_test(self, cls_logits, bbox_pred):
+        if self.priors is None:
+            self.priors = PriorBox(self.cfg)().to(bbox_pred.device)
+        scores = F.softmax(cls_logits, dim=2)  # 计算分类损失
+        boxes = box_utils.convert_locations_to_boxes(
+            bbox_pred, self.priors, self.cfg.MODEL.CENTER_VARIANCE, self.cfg.MODEL.SIZE_VARIANCE
+        )  #将卷积的结果转化成boxex格式 bbox_pred前面卷积生成的是边框的四个坐标形式xyxy  self.priors xywh
+        boxes = box_utils.center_form_to_corner_form(boxes) #变成xyxy
+        detections = (scores, boxes)
+        detections = self.post_processor(detections)
+        return detections, {}

ssd/modeling/box_head/box_predictor.py

@@ -0,0 +1,71 @@
+import torch
+from torch import nn
+
+from ssd.layers import SeparableConv2d
+from ssd.modeling import registry
+
+
+class BoxPredictor(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.cls_headers = nn.ModuleList()
+        self.reg_headers = nn.ModuleList()
+        for level, (boxes_per_location, out_channels) in enumerate(zip(cfg.MODEL.PRIORS.BOXES_PER_LOCATION, cfg.MODEL.BACKBONE.OUT_CHANNELS)):#对要求的feature map进行分类和回归卷积层建立
+            self.cls_headers.append(self.cls_block(level, out_channels, boxes_per_location))
+            self.reg_headers.append(self.reg_block(level, out_channels, boxes_per_location))
+        self.reset_parameters()
+
+    def cls_block(self, level, out_channels, boxes_per_location):
+        raise NotImplementedError
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        raise NotImplementedError
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, features):
+        cls_logits = []
+        bbox_pred = []
+        for feature, cls_header, reg_header in zip(features, self.cls_headers, self.reg_headers):
+            cls_logits.append(cls_header(feature).permute(0, 2, 3, 1).contiguous()) #计算每一层的分类卷积结果  contiguous将结果在内存中变成连续的
+            bbox_pred.append(reg_header(feature).permute(0, 2, 3, 1).contiguous()) #计算每一层的边框回归的卷积结果
+
+        batch_size = features[0].shape[0]
+        cls_logits = torch.cat([c.view(c.shape[0], -1) for c in cls_logits], dim=1).view(batch_size, -1, self.cfg.MODEL.NUM_CLASSES)
+        bbox_pred = torch.cat([l.view(l.shape[0], -1) for l in bbox_pred], dim=1).view(batch_size, -1, 4)
+
+        return cls_logits, bbox_pred
+
+
+@registry.BOX_PREDICTORS.register('SSDBoxPredictor')
+class SSDBoxPredictor(BoxPredictor):
+    def cls_block(self, level, out_channels, boxes_per_location):
+        return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+@registry.BOX_PREDICTORS.register('SSDLiteBoxPredictor')
+class SSDLiteBoxPredictor(BoxPredictor):
+    def cls_block(self, level, out_channels, boxes_per_location):
+        num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
+        if level == num_levels - 1:
+            return nn.Conv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=1)
+        return SeparableConv2d(out_channels, boxes_per_location * self.cfg.MODEL.NUM_CLASSES, kernel_size=3, stride=1, padding=1)
+
+    def reg_block(self, level, out_channels, boxes_per_location):
+        num_levels = len(self.cfg.MODEL.BACKBONE.OUT_CHANNELS)
+        if level == num_levels - 1:
+            return nn.Conv2d(out_channels, boxes_per_location * 4, kernel_size=1)
+        return SeparableConv2d(out_channels, boxes_per_location * 4, kernel_size=3, stride=1, padding=1)
+
+
+def make_box_predictor(cfg):
+    return registry.BOX_PREDICTORS[cfg.MODEL.BOX_HEAD.PREDICTOR](cfg)  #将括号中的函数当作参数进行调用
+#

ssd/modeling/box_head/inference.py

@@ -0,0 +1,65 @@
+import torch
+
+from ssd.structures.container import Container
+from ssd.utils.nms import boxes_nms
+
+
+class PostProcessor: #传进来的是边框的类别打分和边框boxes
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.width = cfg.INPUT.IMAGE_SIZE
+        self.height = cfg.INPUT.IMAGE_SIZE
+
+    def __call__(self, detections):  #detections传进来的参数 是一个三维 10*8732*21 类别 相当于每个box都对每个类别进行预测 10*8732*4 相当于对  10 是10幅图片
+        batches_scores, batches_boxes = detections
+        device = batches_scores.device
+        batch_size = batches_scores.size(0)
+        results = []
+        for batch_id in range(batch_size):  #从0到9
+            processed_boxes = []
+            processed_scores = []
+            processed_labels = []
+
+            per_img_scores, per_img_boxes = batches_scores[batch_id], batches_boxes[batch_id]  # (N, #CLS) (N, 4)取一副图片的数据
+            #print(per_img_scores.data)
+            for class_id in range(1, per_img_scores.size(1)):  # skip background
+
+                scores = per_img_scores[:, class_id]
+                mask = scores > self.cfg.TEST.CONFIDENCE_THRESHOLD
+                scores = scores[mask]
+                if scores.size(0) == 0:
+                    continue
+                boxes = per_img_boxes[mask, :]
+                boxes[:, 0::2] *= self.width
+                boxes[:, 1::2] *= self.height
+
+                keep = boxes_nms(boxes, scores, self.cfg.TEST.NMS_THRESHOLD, self.cfg.TEST.MAX_PER_CLASS)
+
+                nmsed_boxes = boxes[keep, :]
+                nmsed_labels = torch.tensor([class_id] * keep.size(0), device=device)
+                nmsed_scores = scores[keep]
+
+                processed_boxes.append(nmsed_boxes)
+                processed_scores.append(nmsed_scores)
+                processed_labels.append(nmsed_labels)
+
+            if len(processed_boxes) == 0:
+                processed_boxes = torch.empty(0, 4)
+                processed_labels = torch.empty(0)
+                processed_scores = torch.empty(0)
+            else:
+                processed_boxes = torch.cat(processed_boxes, 0)
+                processed_labels = torch.cat(processed_labels, 0)
+                processed_scores = torch.cat(processed_scores, 0)
+
+            if processed_boxes.size(0) > self.cfg.TEST.MAX_PER_IMAGE > 0:
+                processed_scores, keep = torch.topk(processed_scores, k=self.cfg.TEST.MAX_PER_IMAGE)
+                processed_boxes = processed_boxes[keep, :]
+                processed_labels = processed_labels[keep]
+
+            container = Container(boxes=processed_boxes, labels=processed_labels, scores=processed_scores)  #结果输出 处理好的边框坐标 类别 分数
+            container.img_width = self.width
+            container.img_height = self.height
+            results.append(container)
+        return results

ssd/modeling/box_head/loss.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from ssd.utils import box_utils
+
+
+class MultiBoxLoss(nn.Module):
+    def __init__(self, neg_pos_ratio):
+        """Implement SSD MultiBox Loss.
+
+        Basically, MultiBox loss combines classification loss
+         and Smooth L1 regression loss.
+        """
+        super(MultiBoxLoss, self).__init__()
+        self.neg_pos_ratio = neg_pos_ratio
+
+    def forward(self, confidence, predicted_locations, labels, gt_locations):
+        """Compute classification loss and smooth l1 loss.
+
+        Args:
+            confidence (batch_size, num_priors, num_classes): class predictions.
+            predicted_locations (batch_size, num_priors, 4): predicted locations.
+            labels (batch_size, num_priors): real labels of all the priors.
+            gt_locations (batch_size, num_priors, 4): real boxes corresponding all the priors.
+        """
+        num_classes = confidence.size(2)
+        with torch.no_grad():
+            # derived from cross_entropy=sum(log(p))
+            loss = -F.log_softmax(confidence, dim=2)[:, :, 0]
+            mask = box_utils.hard_negative_mining(loss, labels, self.neg_pos_ratio)
+
+        confidence = confidence[mask, :]
+        classification_loss = F.cross_entropy(confidence.view(-1, num_classes), labels[mask], reduction='sum')
+
+        pos_mask = labels > 0
+        predicted_locations = predicted_locations[pos_mask, :].view(-1, 4)
+        gt_locations = gt_locations[pos_mask, :].view(-1, 4)
+        smooth_l1_loss = F.smooth_l1_loss(predicted_locations, gt_locations, reduction='sum')
+        num_pos = gt_locations.size(0)
+        return smooth_l1_loss / num_pos, classification_loss / num_pos