| 113 | |
| 114 | |
| 115 | def sample_positive_negative(labels, positive_count, total_count): |
| 116 | # Sample positive and negative proposals |
| 117 | positive = torch.where(labels >= 1)[0] |
| 118 | negative = torch.where(labels == 0)[0] |
| 119 | num_pos = positive_count |
| 120 | num_pos = min(positive.numel(), num_pos) |
| 121 | num_neg = total_count - num_pos |
| 122 | num_neg = min(negative.numel(), num_neg) |
| 123 | perm_positive_idxs = torch.randperm(positive.numel(), |
| 124 | device=positive.device)[:num_pos] |
| 125 | perm_negative_idxs = torch.randperm(negative.numel(), |
| 126 | device=negative.device)[:num_neg] |
| 127 | pos_idxs = positive[perm_positive_idxs] |
| 128 | neg_idxs = negative[perm_negative_idxs] |
| 129 | sampled_pos_idx_mask = torch.zeros_like(labels, dtype=torch.bool) |
| 130 | sampled_neg_idx_mask = torch.zeros_like(labels, dtype=torch.bool) |
| 131 | sampled_pos_idx_mask[pos_idxs] = True |
| 132 | sampled_neg_idx_mask[neg_idxs] = True |
| 133 | return sampled_neg_idx_mask, sampled_pos_idx_mask |
| 134 | |
| 135 | |
| 136 | def clamp_boxes_to_image_boundary(boxes, image_shape): |