Enable to easily load pretrained SAM weights

Author: sungchul2

otx/algorithms/visual_prompting/adapters/pytorch_lightning/models/visual_prompters/segment_anything.py

@@ -162,8 +162,15 @@ def replace_state_dict_keys(state_dict, revise_keys):
             # state_dict from args.load_from
             state_dict = replace_state_dict_keys(state_dict, revise_keys)
             self.load_state_dict(state_dict)
-        elif self.config.model.checkpoint:
+        elif self.config.model.checkpoint == "pretrained" or self.config.model.checkpoint is None:
+            # load SAM pretrained weights
+            state_dict = torch.hub.load_state_dict_from_url(CKPT_PATHS[self.config.model.backbone])
+            state_dict = replace_state_dict_keys(state_dict, revise_keys)
+            self.load_state_dict(state_dict)
+        else:
+            # load custom weights
             try:
+                # load checkpoint trained by pytorch lightning
                 self.load_from_checkpoint(self.config.model.checkpoint)
             except Exception:
                 if str(self.config.model.checkpoint).startswith("http"):

otx/algorithms/visual_prompting/configs/sam_vit_b/config.yaml

@@ -17,7 +17,7 @@ model:
   freeze_image_encoder: true
   freeze_prompt_encoder: true
   freeze_mask_decoder: false
-  checkpoint: https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
+  checkpoint: pretrained
 
 optimizer:
   name: Adam

tests/unit/algorithms/visual_prompting/adapters/pytorch_lightning/models/visual_prompters/test_segment_anything.py

@@ -51,8 +51,14 @@ def forward(self, *args, **kwargs):
 
 
 class TestSegmentAnything:
+    @pytest.fixture
+    def mocker_load_state_dict(self, mocker) -> None:
+        return mocker.patch(
+            "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.load_state_dict"
+        )
+
     @pytest.fixture(autouse=True)
-    def setup(self, monkeypatch) -> None:
+    def setup(self, mocker, monkeypatch) -> None:
         monkeypatch.setattr(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SAMImageEncoder",
             MockImageEncoder,
@@ -65,6 +71,12 @@ def setup(self, monkeypatch) -> None:
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SAMMaskDecoder",
             MockMaskDecoder,
         )
+        monkeypatch.setattr("torch.hub.load_state_dict_from_url", lambda *args, **kwargs: OrderedDict())
+        monkeypatch.setattr("torch.load", lambda *args, **kwargs: None)
+
+        self.mocker_load_from_checkpoint = mocker.patch(
+            "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.load_from_checkpoint"
+        )
 
         self.base_config = DictConfig(
             dict(
@@ -84,7 +96,7 @@ def setup(self, monkeypatch) -> None:
 
     @e2e_pytest_unit
     @pytest.mark.parametrize("backbone", ["vit_b", "resnet"])
-    def test_set_models(self, mocker, backbone: str) -> None:
+    def test_set_models(self, mocker, mocker_load_state_dict, backbone: str) -> None:
         """Test set_models."""
         mocker.patch(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.freeze_networks"
@@ -114,7 +126,7 @@ def test_set_models(self, mocker, backbone: str) -> None:
     @pytest.mark.parametrize("freeze_prompt_encoder", [True, False])
     @pytest.mark.parametrize("freeze_mask_decoder", [True, False])
     def test_freeze_networks(
-        self, mocker, freeze_image_encoder: bool, freeze_prompt_encoder: bool, freeze_mask_decoder: bool
+        self, mocker, mocker_load_state_dict, freeze_image_encoder: bool, freeze_prompt_encoder: bool, freeze_mask_decoder: bool
     ):
         """Test freeze_networks."""
         mocker.patch(
@@ -154,7 +166,7 @@ def test_freeze_networks(
 
     @e2e_pytest_unit
     @pytest.mark.parametrize("loss_type", ["sam", "medsam"])
-    def test_set_metrics(self, mocker, loss_type: str):
+    def test_set_metrics(self, mocker, mocker_load_state_dict, loss_type: str):
         """Test set_metrics."""
         mocker.patch(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.set_models"
@@ -231,8 +243,8 @@ def test_load_checkpoint_with_state_dict(self, mocker, is_backbone_arg: bool, st
             assert v == sam_state_dict[k]
 
     @e2e_pytest_unit
-    @pytest.mark.parametrize("checkpoint", [None, "checkpoint", "http://checkpoint"])
-    def test_load_checkpoint(self, mocker, monkeypatch, checkpoint: str):
+    @pytest.mark.parametrize("checkpoint", [None, "pretrained", "checkpoint.pth", "http://checkpoint"])
+    def test_load_checkpoint(self, mocker, mocker_load_state_dict, checkpoint: str):
         """Test load_checkpoint."""
         mocker.patch(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.freeze_networks"
@@ -241,21 +253,8 @@ def test_load_checkpoint(self, mocker, monkeypatch, checkpoint: str):
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.set_metrics"
         )
         if checkpoint is not None:
-            monkeypatch.setattr("torch.hub.load_state_dict_from_url", lambda *args, **kwargs: OrderedDict())
-            monkeypatch.setattr("torch.load", lambda *args, **kwargs: None)
-
-            mocker_load_state_dict = mocker.patch(
-                "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.load_state_dict"
-            )
             if checkpoint.startswith("http"):
-                mocker_load_from_checkpoint = mocker.patch(
-                    "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.load_from_checkpoint",
-                    side_effect=ValueError(),
-                )
-            else:
-                mocker_load_from_checkpoint = mocker.patch(
-                    "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.load_from_checkpoint"
-                )
+                self.mocker_load_from_checkpoint.side_effect = ValueError()
 
         config = self.base_config.copy()
         config.model.update(dict(checkpoint=checkpoint))
@@ -264,19 +263,19 @@ def test_load_checkpoint(self, mocker, monkeypatch, checkpoint: str):
 
         if checkpoint is None:
             assert True
-        elif checkpoint.startswith("http"):
+        elif checkpoint.startswith("http") or checkpoint == "pretrained":
             mocker_load_state_dict.assert_called_once()
         else:
-            mocker_load_from_checkpoint.assert_called_once()
+            self.mocker_load_from_checkpoint.assert_called_once()
 
     @e2e_pytest_unit
-    def test_forward(self) -> None:
-        """Test forward."""
+    def test_forward_train(self, mocker_load_state_dict) -> None:
+        """Test forward_train."""
         sam = SegmentAnything(config=self.base_config)
         images = torch.zeros((1))
         bboxes = torch.zeros((1))
 
-        results = sam.forward(images=images, bboxes=bboxes, points=None)
+        results = sam.forward_train(images=images, bboxes=bboxes, points=None)
         pred_masks, ious = results
 
         assert len(bboxes) == len(pred_masks) == len(ious)
@@ -285,7 +284,7 @@ def test_forward(self) -> None:
     @pytest.mark.parametrize(
         "loss_type,expected", [("sam", torch.tensor(2.4290099144)), ("medsam", torch.tensor(0.9650863409))]
     )
-    def test_training_step(self, mocker, loss_type: str, expected: Tensor) -> None:
+    def test_training_step(self, mocker, mocker_load_state_dict, loss_type: str, expected: Tensor) -> None:
         """Test training_step."""
         mocker.patch(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.forward",
@@ -311,7 +310,7 @@ def test_training_step(self, mocker, loss_type: str, expected: Tensor) -> None:
         assert torch.equal(results, expected)
 
     @e2e_pytest_unit
-    def test_training_epoch_end(self) -> None:
+    def test_training_epoch_end(self, mocker_load_state_dict) -> None:
         """Test training_epoch_end."""
         sam = SegmentAnything(config=self.base_config)
         for k, v in sam.train_metrics.items():
@@ -328,7 +327,7 @@ def test_training_epoch_end(self) -> None:
         assert sam.train_metrics["train_loss_iou"].compute().isnan()
 
     @e2e_pytest_unit
-    def test_validation_step(self, mocker) -> None:
+    def test_validation_step(self, mocker, mocker_load_state_dict) -> None:
         """Test validation_step."""
         mocker.patch(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.forward",
@@ -357,7 +356,7 @@ def test_validation_step(self, mocker) -> None:
         assert torch.equal(results["val_IoU"].compute(), torch.tensor(1.0))
 
     @e2e_pytest_unit
-    def test_validation_epoch_end(self) -> None:
+    def test_validation_epoch_end(self, mocker_load_state_dict) -> None:
         """Test validation_epoch_end."""
         sam = SegmentAnything(config=self.base_config)
         for k, v in sam.val_metrics.items():
@@ -377,7 +376,7 @@ def test_validation_epoch_end(self) -> None:
             (False, torch.Tensor([[False for _ in range(4)] for _ in range(4)])),
         ],
     )
-    def test_predict_step(self, mocker, return_logits: bool, expected: Tensor) -> None:
+    def test_predict_step(self, mocker, mocker_load_state_dict, return_logits: bool, expected: Tensor) -> None:
         """Test predict_step."""
         mocker.patch(
             "otx.algorithms.visual_prompting.adapters.pytorch_lightning.models.visual_prompters.segment_anything.SegmentAnything.forward",
@@ -414,7 +413,7 @@ def test_predict_step(self, mocker, return_logits: bool, expected: Tensor) -> No
         ],
     )
     def test_postprocess_masks(
-        self, input_size: Tuple[int], original_size: Tuple[int], is_predict: bool, expected: Tuple[int]
+        self, mocker_load_state_dict, input_size: Tuple[int], original_size: Tuple[int], is_predict: bool, expected: Tuple[int]
     ) -> None:
         """Test postprocess_masks."""
         sam = SegmentAnything(config=self.base_config)
@@ -433,7 +432,7 @@ def test_postprocess_masks(
             (Tensor([[0, 0, 0.3, 0.3, 0, 0]]), Tensor([[0, 0, 1, 1, 0, 0]]), Tensor([0.3888888359])),
         ],
     )
-    def test_calculate_dice_loss(self, inputs: Tensor, targets: Tensor, expected: Tensor) -> None:
+    def test_calculate_dice_loss(self, mocker_load_state_dict, inputs: Tensor, targets: Tensor, expected: Tensor) -> None:
         """Test calculate_dice_loss."""
         sam = SegmentAnything(config=self.base_config)
 
@@ -450,7 +449,7 @@ def test_calculate_dice_loss(self, inputs: Tensor, targets: Tensor, expected: Te
             (Tensor([[0, 0, 0.3, 0.3, 0, 0]]), Tensor([[0, 0, 1, 1, 0, 0]]), Tensor([0.0226361733])),
         ],
     )
-    def test_calculate_sigmoid_ce_focal_loss(self, inputs: Tensor, targets: Tensor, expected: Tensor) -> None:
+    def test_calculate_sigmoid_ce_focal_loss(self, mocker_load_state_dict, inputs: Tensor, targets: Tensor, expected: Tensor) -> None:
         """Test calculate_sigmoid_ce_focal_loss."""
         sam = SegmentAnything(config=self.base_config)
 
@@ -467,7 +466,7 @@ def test_calculate_sigmoid_ce_focal_loss(self, inputs: Tensor, targets: Tensor,
             (Tensor([[0, 0, 0.3, 0.3, 0, 0]]), Tensor([[0, 0, 1, 1, 0, 0]]), Tensor([0.0])),
         ],
     )
-    def test_calculate_iou(self, inputs: Tensor, targets: Tensor, expected: Tensor) -> None:
+    def test_calculate_iou(self, mocker_load_state_dict, inputs: Tensor, targets: Tensor, expected: Tensor) -> None:
         """Test calculate_iou."""
         sam = SegmentAnything(config=self.base_config)