facebookresearch / facebookresearch/sam2

How To Retrieve Encoder Output Embeddings When Using Automatic Mask Generator

Open
#442 0 comments 0 reactions 0 assignees View on GitHub
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.
![image](https://github.com/user-attachments/assets/6a1297b1-71e9-4b06-8600-c81873a739d9)

# 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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.