66 lines
2.9 KiB
Python
Executable File
66 lines
2.9 KiB
Python
Executable File
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
|