Skip to content

Instantly share code, notes, and snippets.

@arunm8489

arunm8489/35

Last active June 5, 2020 05:09
Show Gist options
  • Save arunm8489/d04763f1c8650bbe330a14d569cb6952 to your computer and use it in GitHub Desktop.
Save arunm8489/d04763f1c8650bbe330a14d569cb6952 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
35
#perform NMS
#get the detections with one particular class
cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
# taking the non zero indexes
class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
image_pred_class = image_pred_[class_mask_ind]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment