facebookresearch / facebookresearch/EdgeTAM

Adding CoreML Export Support - Looking for Feedback

Open
#14 10 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
978
Forks
83
PR merge metrics
No merged PRs in 30d

Description

Before I make the pr, just want to make sure that my process makes sense:

Started by exploring how to handle EdgeTAM's dynamic prompts (points, boxes, masks) in CoreML. The main challenge was that CoreML prefers static input shapes, but EdgeTAM needs to handle variable numbers of points and different prompt types.

Initially tried exporting the entire model as one piece, but quickly realized that wouldn't work well for mobile - the prompt handling was too dynamic and CoreML's tracing couldn't handle the conditional logic properly.

Ended up splitting the model into 3 parts:
- Image encoder (~9.6MB) - handles the heavy lifting of processing the 1024x1024 input image
- Prompt encoder (~2MB) - processes user inputs with fixed slots (up to 4 points, 1 box, 1 mask)
- Mask decoder (~8MB) - generates the final masks

The tricky part was getting the prompt encoder right. Had to dig into the SAM2 code to understand how prompts are embedded, then create wrapper classes that make the dynamic parts static enough for CoreML.

Let me know if this makes sense! happy to create the pr soon :)

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.