Encode the current image and its prediction into a memory feature.
(
self,
current_vision_feats,
feat_sizes,
pred_masks_high_res,
is_mask_from_pts,
)
| 662 | return pix_feat_with_mem |
| 663 | |
| 664 | def _encode_new_memory( |
| 665 | self, |
| 666 | current_vision_feats, |
| 667 | feat_sizes, |
| 668 | pred_masks_high_res, |
| 669 | is_mask_from_pts, |
| 670 | ): |
| 671 | """Encode the current image and its prediction into a memory feature.""" |
| 672 | B = current_vision_feats[-1].size(1) # batch size on this frame |
| 673 | C = self.hidden_dim |
| 674 | H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size |
| 675 | # top-level feature, (HW)BC => BCHW |
| 676 | pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W) |
| 677 | if self.non_overlap_masks_for_mem_enc and not self.training: |
| 678 | # optionally, apply non-overlapping constraints to the masks (it's applied |
| 679 | # in the batch dimension and should only be used during eval, where all |
| 680 | # the objects come from the same video under batch size 1). |
| 681 | pred_masks_high_res = self._apply_non_overlapping_constraints( |
| 682 | pred_masks_high_res |
| 683 | ) |
| 684 | # scale the raw mask logits with a temperature before applying sigmoid |
| 685 | binarize = self.binarize_mask_from_pts_for_mem_enc and is_mask_from_pts |
| 686 | if binarize and not self.training: |
| 687 | mask_for_mem = (pred_masks_high_res > 0).float() |
| 688 | else: |
| 689 | # apply sigmoid on the raw mask logits to turn them into range (0, 1) |
| 690 | mask_for_mem = torch.sigmoid(pred_masks_high_res) |
| 691 | # apply scale and bias terms to the sigmoid probabilities |
| 692 | if self.sigmoid_scale_for_mem_enc != 1.0: |
| 693 | mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc |
| 694 | if self.sigmoid_bias_for_mem_enc != 0.0: |
| 695 | mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc |
| 696 | maskmem_out = self.memory_encoder( |
| 697 | pix_feat, mask_for_mem, skip_mask_sigmoid=True # sigmoid already applied |
| 698 | ) |
| 699 | maskmem_features = maskmem_out["vision_features"] |
| 700 | maskmem_pos_enc = maskmem_out["vision_pos_enc"] |
| 701 | |
| 702 | return maskmem_features, maskmem_pos_enc |
| 703 | |
| 704 | def track_step( |
| 705 | self, |
no test coverage detected