def forward(self, x, boxes):
"""
Arguments:
x (Tensor): the mask logits
boxes (list[BoxList]): bounding boxes that are used as
reference, one for ech image
Returns:
results (list[BoxList]): one BoxList for each image, containing
the extra field mask
"""
mask_prob = x.sigmoid()
# select masks coresponding to the predicted classes
num_masks = x.shape[0]
labels = [bbox.get_field("labels") for bbox in boxes]
labels = torch.cat(labels)
index = torch.arange(num_masks, device=labels.device)
mask_prob = mask_prob[index, labels][:, None]
boxes_per_image = [len(box) for box in boxes]
mask_prob = mask_prob.split(boxes_per_image, dim=0)
results = []
for prob, box in zip(mask_prob, boxes):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
bbox_scores = bbox.get_field("scores")
bbox.add_field("mask", prob.cpu().numpy())
bbox.add_field("mask_scores", bbox_scores.cpu().numpy())
results.append(bbox)
return results
def forward(self, x, boxes):
"""
Arguments:
x (Tensor): the mask logits
boxes (list[BoxList]): bounding boxes that are used as
reference, one for ech image
Returns:
results (list[BoxList]): one BoxList for each image, containing
the extra field mask
"""
parsing_prob = x
parsing_prob = F.softmax(parsing_prob, dim=1)
boxes_per_image = [len(box) for box in boxes]
parsing_prob = parsing_prob.split(boxes_per_image, dim=0)
results = []
for prob, box in zip(parsing_prob, boxes):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
bbox_scores = bbox.get_field("scores")
bbox.add_field("parsing", prob.cpu().numpy())
bbox.add_field("parsing_scores", bbox_scores.cpu().numpy())
results.append(bbox)
return results
def forward(self, boxes, pred_parsingiou):
num_parsings = pred_parsingiou.shape[0]
index = torch.arange(num_parsings, device=pred_parsingiou.device)
parsingious = pred_parsingiou[index, 0]
parsingious = [parsingious]
results = []
for parsingiou, box in zip(parsingious, boxes):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
bbox_scores = bbox.get_field("scores")
parsing_scores = torch.sqrt(bbox_scores * parsingiou)
bbox.add_field("parsing_scores", parsing_scores.cpu().numpy())
prob = bbox.get_field("parsing")
bbox.add_field("parsing", prob.cpu().numpy())
results.append(bbox)
return results
def forward(self, boxes, pred_maskiou, labels):
num_masks = pred_maskiou.shape[0]
index = torch.arange(num_masks, device=labels.device)
maskious = pred_maskiou[index, labels]
maskious = [maskious]
results = []
for maskiou, box in zip(maskious, boxes):
bbox = BoxList(box.bbox, box.size, mode="xyxy")
for field in box.fields():
bbox.add_field(field, box.get_field(field))
bbox_scores = bbox.get_field("scores")
mask_scores = bbox_scores * maskiou
bbox.add_field("mask_scores", mask_scores.cpu().numpy())
prob = bbox.get_field("mask")
bbox.add_field("mask", prob.cpu().numpy())
results.append(bbox)
return results
