Initial setup of SAM architecture

Author: mbasheer04

tiatoolbox/models/architecture/sam.py

@@ -0,0 +1,73 @@
+"""Define SAM architecture."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F  # noqa: N812
+from skimage import morphology
+from torch import nn
+
+from tiatoolbox.utils import misc
+
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+
+sam_model = build_sam2(model_cfg, checkpoint)
+
+predictor = SAM2ImagePredictor(sam_model)
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    predictor.set_image(<your_image>)
+    masks, _, _ = predictor.predict(<input_prompts>)
+
+class SAM(ModelABC):
+
+    def __init__(self: ModelABC) -> None:
+        """Initialize Abstract class ModelABC."""
+        super().__init__()
+        self._postproc = self.postproc
+        self._preproc = self.preproc
+
+    def forward(self: ModelABC, *args: tuple[Any, ...], **kwargs: dict) -> None:
+        """Torch method, this contains logic for using layers defined in init."""
+        ...  # pragma: no cover
+
+    def infer_batch(
+        model: torch.nn.Module,
+        batch_data: np.ndarray,
+        *,
+        on_gpu: bool,
+    ) -> None:
+        """Run inference on an input batch.
+
+        Contains logic for forward operation as well as I/O aggregation.
+
+        Args:
+            model (nn.Module):
+                PyTorch defined model.
+            batch_data (np.ndarray):
+                A batch of data generated by
+                `torch.utils.data.DataLoader`.
+            on_gpu (bool):
+                Whether to run inference on a GPU.
+
+        """
+        ...  # pragma: no cover
+
+    @staticmethod
+    def preproc(image: np.ndarray) -> np.ndarray:
+        """Define the pre-processing of this class of model."""
+        return image
+
+    @staticmethod
+    def postproc(image: np.ndarray) -> np.ndarray:
+        """Define the post-processing of this class of model."""
+        return image
+