Repository: kornia/kornia

Issue:
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.



Discussion:

Resolving patch:
diff --git a/kornia/contrib/models/sam/model.py b/kornia/contrib/models/sam/model.py
--- 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/a509aac54fdd7af59f843135f2f7cee307283c88/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
--- 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:


Test patch:
diff --git a/test/contrib/models/test_sam.py b/test/contrib/models/test_sam.py
--- 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))