def forward(self, x, boxes, features):
mask_prob = x
keypoints_features = x
scores = None
# print(mask_prob.shape)
if self.keypointer:
mask_prob, scores = self.keypointer(x, boxes)
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)
# print([yy.shape for yy in y])
keypoints_features = keypoints_features.split(boxes_per_image, dim=0)
# print(feature.shape)
scores = scores.split(boxes_per_image, dim=0)
results = []
# print(boxes_per_image)
for prob, box, score, feature in zip(mask_prob, boxes, scores,
keypoints_features):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
prob = PersonKeypoints(prob, box.size)
prob.add_field("logits", score)
prob.add_field("feature", feature)
bbox.add_field("keypoints", prob)
results.append(bbox)
return results
def forward(self, x, edges, boxes):
graph_prob = x
scores = None
if self.keypointer:
graph_prob, scores = self.keypointer(x, edges, boxes)
assert len(boxes) == 1, "Only non-batched inference supported for now"
boxes_per_image = [box.bbox.size(0) for box in boxes]
graph_prob = graph_prob.split(boxes_per_image, dim=0)
scores = scores.split(boxes_per_image, dim=0)
results = []
for prob, box, score in zip(graph_prob, boxes, scores):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
prob = PersonKeypoints(prob, box.size)
prob.add_field("logits", score)
bbox.add_field("keypoints", prob)
results.append(bbox)
return results