def pp_predictions(self, inspred, carmask, pedmask, cyclistmask):
selector = torch.ones_like(inspred.scores)
if inspred.scores.shape[0] == 0:
return inspred
# Get indices sort by score
inspred = self.erase_srhink(inspred, carmask, cat=2)
inspred = self.erase_srhink(inspred, cyclistmask, cat=1)
inspred = self.erase_srhink(inspred, pedmask, cat=0)
for k in range(inspred.scores.shape[0]):
cat = inspred.pred_classes[k].item()
numpixel = torch.sum(inspred.pred_masks[k]).item()
if numpixel < self.minpixel[cat]:
selector[k] = 0
pp_inspred = Instances(image_size=inspred.image_size)
selector = selector == 1
pp_inspred.scores = inspred.scores[selector]
pp_inspred.pred_classes = inspred.pred_classes[selector]
pp_inspred.pred_boxes = inspred.pred_boxes[selector]
pp_inspred.pred_masks = inspred.pred_masks[selector]
return pp_inspred
def pp_predictions_simple(self, inspred):
selector = torch.ones_like(inspred.scores)
if inspred.scores.shape[0] == 0:
return inspred
for k in range(inspred.scores.shape[0]):
cat = inspred.pred_classes[k].item()
conf = inspred.scores[k].item()
numpixel = torch.sum(inspred.pred_masks).item()
if conf < self.catconfbar[cat]:
selector[k] = 0
if numpixel < self.minpixel[cat]:
selector[k] = 0
pp_inspred = Instances(image_size=inspred.image_size)
selector = selector == 1
pp_inspred.scores = inspred.scores[selector]
pp_inspred.pred_classes = inspred.pred_classes[selector]
pp_inspred.pred_boxes = inspred.pred_boxes[selector]
pp_inspred.pred_masks = inspred.pred_masks[selector]
return pp_inspred