facebookresearch / facebookresearch/sam2
How To Retrieve Encoder Output Embeddings When Using Automatic Mask Generator
- Dominant language
- Jupyter Notebook
- Stars
- 19.9k
- Forks
- 2.5k
- PR merge metrics
- No merged PRs in 30d
Description
Hi,
I would like to retrieve the image embeddings produced by the encoder just like we can when running `SAM2ImagePredictor` as mentioned [in this issue](https://github.com/facebookresearch/sam2/issues/11#issuecomment-2265172163).
# Issue
After looking at the code, it seems like even if we we use `crop_n_layers=0`, there is a call to `reset_predictor` which will erase the imag embeddings making it impossible to re-use them for other downstream use cases.

# Workaround
The current workaround I found is a bit of a hack, I override the method `reset_predictor` of the predictor to be my custom function that saves the embeddings before executing the normal `reset_predictor` function which will set them to `None`.
```python
class MaskGeneratorWrapper:
def __init__(self, mask_generator):
original_reset_predictor = mask_generator.predictor.__class__.reset_predictor
self.img_embed = {}
def custom_reset_predictor(cls):
self.img_embed = cls._features
original_reset_predictor(cls)
self.mask_generator = mask_generator
self.mask_generator.predictor.reset_predictor = custom_reset_predictor.__get__(mask_generator.predictor)
def predict(self, image):
return self.mask_generator.generate(image)
mask_generator_wrapper = MaskGeneratorWrapper(mask_generator)
predictions = mask_generator_wrapper.predict(rearrange(image, "c h w -> h w c").cpu().numpy())
```
# Proposal
Would it be possible to add a flag to `SAM2AutomaticMaskGenerator` or to its `generate` method so that the image embeddings are preserved ?
Contributor guide
Assessment
This issue has not been assessed yet.