facebookresearch / segment-anything

The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.
Apache License 2.0
47.5k stars 5.62k forks source link

torchScript Export #241

Open paul-ml75 opened 1 year ago

paul-ml75 commented 1 year ago

Hi!

I truly appreciate the OXXN export script, which is fantastic. Would it be possible to add a TorchScript export as well? This addition would be extremely helpful for utilizing Segment-Anything with various third-party software.

Thank you,

Paul

leonardcser commented 1 year ago

Hi Paul, I've played a bit around with this and you can find a starting point here: https://github.com/csia-pme/djl-image-sam

paul-ml75 commented 1 year ago

Hi Leonardscer!

that is absolutely awesome! having a look right now!

Paul

darshats commented 1 year ago

Hi leonardcser, Thanks for the helpful starting points! I was able to work with it.

One issue I see is that by wrapping the model it is cumbersome to use as part of the automatic_mask_generator provided by the facebook team. This script is very useful as it does a lot of heavy lifting of implementing moving windows (crops) and then deduping masks.

So I changed the sam_predictor to only script out image encoder, prompt encoder and mask decoder separately. With that we can continue to use the automatic_mask_generator as is. Just fyi.