How To Retrieve Encoder Output Embeddings When Using Automatic Mask Generator

[AlexandreBrown]

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.

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.

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 ?