first commit

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