import torch
from torch import nn
class KeypointPostProcessor(nn.Module):
def __init__(self, keypointer=None):
super(KeypointPostProcessor, self).__init__()
self.keypointer = keypointer # xy_preds(x,1,y), end_scores
def forward(self, x, boxes):
mask_prob = x
scores = None
if self.keypointer:
mask_prob, scores = self.keypointer(x, boxes) # mask_prob和 boxes →→ mask_prob, scores
assert len(boxes) == 1, "Only non-batched inference supported for now"
boxes_per_image = [box.bbox.size(0) for box in boxes]
mask_prob = mask_prob.split(boxes_per_image, dim=0)
scores = scores.split(boxes_per_image, dim=0)
results = []
for prob, box, score in zip(mask_prob, boxes, scores):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box
maskrcnn_benchmark理解记录——modeling\roi_heads\keypoint_head\inference.py
最新推荐文章于 2022-11-19 09:24:51 发布
![](https://img-home.csdnimg.cn/images/20240711042549.png)