Input: - memory: bs, \sum{hw}, d_model - memory_padding_mask: bs, \sum{hw} - spatial_shapes: nlevel, 2 - learnedwh: 2 Output: - output_memory: bs, \sum{hw}, d_model - output_proposals: bs, \sum{hw}, 4
(memory: Tensor,
memory_padding_mask: Tensor,
spatial_shapes: Tensor,
learnedwh=None)
| 30 | |
| 31 | |
| 32 | def gen_encoder_output_proposals(memory: Tensor, |
| 33 | memory_padding_mask: Tensor, |
| 34 | spatial_shapes: Tensor, |
| 35 | learnedwh=None): |
| 36 | """ |
| 37 | Input: |
| 38 | - memory: bs, \sum{hw}, d_model |
| 39 | - memory_padding_mask: bs, \sum{hw} |
| 40 | - spatial_shapes: nlevel, 2 |
| 41 | - learnedwh: 2 |
| 42 | Output: |
| 43 | - output_memory: bs, \sum{hw}, d_model |
| 44 | - output_proposals: bs, \sum{hw}, 4 |
| 45 | """ |
| 46 | N_, S_, C_ = memory.shape |
| 47 | base_scale = 4.0 |
| 48 | proposals = [] |
| 49 | _cur = 0 |
| 50 | for lvl, (H_, W_) in enumerate(spatial_shapes): |
| 51 | mask_flatten_ = memory_padding_mask[:, _cur:(_cur + H_ * W_)].view( |
| 52 | N_, H_, W_, 1) |
| 53 | valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1) |
| 54 | valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1) |
| 55 | |
| 56 | grid_y, grid_x = torch.meshgrid( |
| 57 | torch.linspace(0, |
| 58 | H_ - 1, |
| 59 | H_, |
| 60 | dtype=torch.float32, |
| 61 | device=memory.device), |
| 62 | torch.linspace(0, |
| 63 | W_ - 1, |
| 64 | W_, |
| 65 | dtype=torch.float32, |
| 66 | device=memory.device)) |
| 67 | grid = torch.cat( |
| 68 | [grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) # H_, W_, 2 |
| 69 | |
| 70 | scale = torch.cat([valid_W.unsqueeze(-1), |
| 71 | valid_H.unsqueeze(-1)], 1).view(N_, 1, 1, 2) |
| 72 | grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale |
| 73 | |
| 74 | if learnedwh is not None: |
| 75 | wh = torch.ones_like(grid) * learnedwh.sigmoid() * (2.0**lvl) |
| 76 | else: |
| 77 | wh = torch.ones_like(grid) * 0.05 * (2.0**lvl) |
| 78 | proposal = torch.cat((grid, wh), -1).view(N_, -1, 4) |
| 79 | proposals.append(proposal) |
| 80 | _cur += (H_ * W_) |
| 81 | # import pdb; pdb.set_trace() |
| 82 | output_proposals = torch.cat(proposals, 1) |
| 83 | output_proposals_valid = ((output_proposals > 0.01) & |
| 84 | (output_proposals < 0.99)).all(-1, keepdim=True) |
| 85 | output_proposals = torch.log(output_proposals / |
| 86 | (1 - output_proposals)) # unsigmoid |
| 87 | output_proposals = output_proposals.masked_fill( |
| 88 | memory_padding_mask.unsqueeze(-1), float('inf')) |
| 89 | output_proposals = output_proposals.masked_fill(~output_proposals_valid, |