Repository: kornia/kornia
Pull request: Add pre-trained flag for SAM
#### Changes
<!-- Please include a summary of the change and which issue is fixed. -->
<!-- Please also include relevant motivation and context. -->
<!-- List any dependencies that are required for this change. -->

Fixes #2452 

#### Type of change
<!-- Please delete options that are not relevant. -->
- [ ] 📚  Documentation Update
- [ ] 🧪 Tests Cases
- [ ] 🐞 Bug fix (non-breaking change which fixes an issue)
- [x] 🔬 New feature (non-breaking change which adds functionality)
- [ ] 🚨 Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] 📝 This change requires a documentation update


#### Checklist

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] Did you update CHANGELOG in case of a major change?


Problem statement:
Add ability to use `pretrained=True` for SAM
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Add ability to use `pretrained=True` for SAM (and MobileSAM).

```python
prompter = VisualPrompter(SamConfig('vit_b', pretrained=True))
```

In other words, when there is `pretrained=True`, we will fill the checkpoint path with default values.

## Motivation

<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

Currently, to use pretrained SAM, we have to specify the checkpoint url to the config (except using the default ViT-H backbone). Adding a `pretrained` flag would make the experience better for the users.

Hints from issue discussion:

Commit message: add pre-trained for SAM

Patch to review:
diff --git a/kornia/contrib/models/sam/model.py b/kornia/contrib/models/sam/model.py
index 84a1455d06..e9addb4aef 100644
--- a/kornia/contrib/models/sam/model.py
+++ b/kornia/contrib/models/sam/model.py
@@ -8,6 +8,7 @@
 """
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any
@@ -55,6 +56,7 @@ class SamConfig:
 
     model_type: str | int | SamModelType | None = None
     checkpoint: str | None = None
+    pretrained: bool = False
 
     encoder_embed_dim: int | None = None
     encoder_depth: int | None = None
@@ -172,8 +174,20 @@ def from_config(config: SamConfig) -> Sam:
         else:
             raise NotImplementedError('Unexpected config. The model_type should be provide or the encoder configs.')
 
-        if config.checkpoint:
-            model.load_checkpoint(config.checkpoint)
+        checkpoint = config.checkpoint
+        if config.pretrained:
+            if checkpoint is None:
+                checkpoint = {
+                    SamModelType.vit_b: "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+                    SamModelType.vit_l: "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
+                    SamModelType.vit_h: "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+                    SamModelType.mobile_sam: "https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt",
+                }[model_type]
+            else:
+                warnings.warn("checkpoint is not None. pretrained=True is ignored")
+
+        if checkpoint:
+            model.load_checkpoint(checkpoint)
 
         return model
 
diff --git a/kornia/contrib/visual_prompter.py b/kornia/contrib/visual_prompter.py
index caa83adcff..f25ab22b1d 100644
--- a/kornia/contrib/visual_prompter.py
+++ b/kornia/contrib/visual_prompter.py
@@ -52,9 +52,7 @@ class VisualPrompter:
 
     def __init__(
         self,
-        config: SamConfig = SamConfig(
-            model_type='vit_h', checkpoint='https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth'
-        ),
+        config: SamConfig = SamConfig(model_type='vit_h', pretrained=True),
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
diff --git a/test/contrib/models/test_sam.py b/test/contrib/models/test_sam.py
index 738f923228..344649bdf0 100644
--- a/test/contrib/models/test_sam.py
+++ b/test/contrib/models/test_sam.py
@@ -91,3 +91,7 @@ def test_dynamo(self, device, torch_optimizer):
 
         self.assert_close(expected[0].logits, actual[0].logits)
         self.assert_close(expected[0].scores, actual[0].scores)
+
+    @pytest.mark.parametrize("model_type", ["vit_b", "mobile_sam"])
+    def test_pretrained(self, model_type):
+        Sam.from_config(SamConfig(model_type, pretrained=True))


Code review comments:
diff_hunk: @@ -172,8 +174,20 @@ def from_config(config: SamConfig) -> Sam:
         else:
             raise NotImplementedError('Unexpected config. The model_type should be provide or the encoder configs.')
 
-        if config.checkpoint:
-            model.load_checkpoint(config.checkpoint)
+        checkpoint = config.checkpoint
+        if config.pretrained:
+            if checkpoint is None:
+                checkpoint = {
+                    SamModelType.vit_b: "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+                    SamModelType.vit_l: "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
+                    SamModelType.vit_h: "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+                    SamModelType.mobile_sam: "https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt",
path: kornia/contrib/models/sam/model.py
text: @user1:
shouldn't we use the permanent link instead?

https://github.com/ChaoningZhang/MobileSAM/blob/01ea8d0f5590082f0c1ceb0a3e2272593f20154b/weights/mobile_sam.pt