(mask, num_samples)
| 15 | |
| 16 | |
| 17 | def sample_rays(mask, num_samples): |
| 18 | B, H, W = mask.shape |
| 19 | mask_unfold = mask.reshape(-1) |
| 20 | indices = torch.rand_like(mask_unfold).topk(num_samples)[1] |
| 21 | sampled_masks = (torch.zeros_like( |
| 22 | mask_unfold).scatter_(-1, indices, 1).reshape(B, H, W) > 0) |
| 23 | return sampled_masks |