Added docstrings

Author: mbasheer04

tests/models/test_arch_sam.py

@@ -4,10 +4,10 @@
 from typing import Callable
 
 from tiatoolbox.models import SAM
-from tiatoolbox.models.architecture.sam import SAMPrompts
+from tiatoolbox.utils import env_detection as toolbox_env
 from tiatoolbox.utils import imread
 
-ON_GPU = False  # TODO: Use Environment variable to set this to True
+ON_GPU = toolbox_env.has_gpu()
 
 # Test pretrained Model =============================
 
@@ -26,17 +26,14 @@ def test_functional_sam(
     # test inference
     # create prompts
 
-    prompts1 = SAMPrompts(point_coords=[[64, 64]])
-    prompts2 = SAMPrompts(point_coords=[[64, 64]], point_labels=[1])
-    prompts3 = SAMPrompts(box_coords=[[64, 64, 128, 128]])
-    prompts4 = SAMPrompts(
-        point_coords=[[64, 64]], point_labels=[1], box_coords=[[64, 64, 128, 128]]
-    )
+    points1 = [[[64, 64]]]
+    points2 = [[[64, 64], [128, 128]]]
+    boxes1 = [[[64, 64, 128, 128]]]
 
     model = SAM()
 
     _ = model.infer_batch(model, img, on_gpu=ON_GPU)  # no prompts
-    _ = model.infer_batch(model, img, prompts=prompts1, on_gpu=ON_GPU)
-    _ = model.infer_batch(model, img, prompts=prompts2, on_gpu=ON_GPU)
-    _ = model.infer_batch(model, img, prompts=prompts3, on_gpu=ON_GPU)
-    _ = model.infer_batch(model, img, prompts=prompts4, on_gpu=ON_GPU)
+    _ = model.infer_batch(model, img, points1, on_gpu=ON_GPU)
+    _ = model.infer_batch(model, img, points2, on_gpu=ON_GPU)
+    _ = model.infer_batch(model, img, box_coords=boxes1, on_gpu=ON_GPU)
+    _ = model.infer_batch(model, img, points2, boxes1, on_gpu=ON_GPU)

tiatoolbox/models/architecture/sam.py

@@ -9,6 +9,7 @@
 from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
 from sam2.build_sam import build_sam2, build_sam2_hf
 from sam2.sam2_image_predictor import SAM2ImagePredictor
+from segment_anything import SamAutomaticMaskGenerator, SamPredictor, sam_model_registry
 
 from tiatoolbox.models.models_abc import ModelABC
 
@@ -17,33 +18,112 @@
 
 
 class SAM(ModelABC):
-    """SAM architecture."""
+    """Segment Anything Model (SAM) Architecture.
+
+    Meta AI's zero-shot segmentation model.
+    SAM is used for interactive general-purpose segmentation.
+
+    Currently supports both SAM and SAM2, each of which require
+    different model checkpoints and configuration files.
+
+    SAM accepts an RGB image patch along with a list of point and bounding
+    box coordinates as prompts.
+
+    Args:
+        model_type (str):
+            Model type. Currently supported: vit_b, vit_l, vit_h.
+            Required for SAM.
+        checkpoint_path (str):
+            Path to the model checkpoint.
+            Required for both SAM and SAM2.
+        model_cfg_path (str):
+            Path to the model configuration file.
+            Required for SAM2.
+        model_hf_path (str):
+            Huggingface path for the pretrained SAM2 model.
+            If provided, it will override the checkpoint_path and model_cfg_path.
+            Default is "facebook/sam2-hiera-tiny".
+        device (str):
+            Device to run inference on.
+        use_sam2 (bool):
+            Whether to use SAM2 or not. Default is True.
+
+    Examples:
+        >>> # instantiate SAM with checkpoint path and model type
+        >>> sam = SAM(
+        ...     model_type="vit_b",
+        ...     checkpoint_path="path/to/sam_checkpoint.pth"
+        ...     use_sam2=False
+        ... )
+        >>> # instantiate SAM2 with checkpoint and config path
+        >>> sam2 = SAM(
+        ...     checkpoint_path="path/to/sam2_checkpoint.pth",
+        ...     model_cfg_path="path/to/sam2_config.yaml"
+        ... )
+        >>> # instantiate SAM2 with Huggingface path
+        >>> sam2 = SAM(
+        ...     model_hf_path="facebook/sam2-hiera-tiny"
+        ... )
+    """
 
     def __init__(
         self: SAM,
-        model_hf_path: str | None = "facebook/sam2-hiera-tiny",
+        model_type: str | None = None,
         checkpoint_path: str | None = None,
         model_cfg_path: str | None = None,
+        model_hf_path: str = "facebook/sam2-hiera-tiny",
+        *,
+        device: str = "cpu",
+        use_sam2: bool = True,
     ) -> None:
         """Initialize :class:`SAM`."""
         super().__init__()
+        self.use_sam2 = use_sam2
         self.net_name = "SAM"
 
-        if checkpoint_path is None or model_cfg_path is None:
-            self.model = build_sam2_hf(model_hf_path, device="cpu")
+        if self.use_sam2:
+            # Load SAM2
+            if checkpoint_path is None or model_cfg_path is None:
+                self.model = build_sam2_hf(model_hf_path, device=device)
+            else:
+                self.model = build_sam2(model_cfg_path, checkpoint_path)
+            self.predictor = SAM2ImagePredictor(self.model)
+            self.generator = SAM2AutomaticMaskGenerator(self.model)
         else:
-            self.model = build_sam2(model_cfg_path, checkpoint_path)
-
-        self.predictor = SAM2ImagePredictor(self.model)
-        self.generator = SAM2AutomaticMaskGenerator(self.model)
+            # Load original SAM
+            if checkpoint_path is None:
+                msg = "You must provide a checkpoint path for SAM."
+                raise ValueError(msg)
+            self.model = sam_model_registry[model_type](checkpoint=checkpoint_path).to(
+                device
+            )
+            self.predictor = SamPredictor(self.model)
+            self.generator = SamAutomaticMaskGenerator(self.model)
 
     def forward(
         self: SAM,
         imgs: list,
         point_coords: list[list[IntPair]] | None = None,
         box_coords: list[list[IntBounds]] | None = None,
     ) -> np.ndarray:
-        """Torch method, this contains logic for using layers defined in init."""
+        """Torch method. Defines forward pass on each image in the batch.
+
+        Note: This architecture only uses a single layer, so only one forward pass
+        is needed.
+
+        Args:
+            imgs (list):
+                List of images to process, of the shape NHWC.
+            point_coords (list):
+                List of point coordinates for each image.
+            box_coords (list):
+                List of bounding box coordinates for each image.
+
+        Returns:
+            list:
+                List of masks and scores for each image.
+
+        """
         batch_masks, batch_scores = [], []
 
         for i, image in enumerate(imgs):
@@ -96,8 +176,10 @@ def infer_batch(
             batch_data (list):
                 A batch of data generated by
                 `torch.utils.data.DataLoader`.
-            prompts (SAMPrompts):
-                Prompts for SAM model.
+            point_coords (list):
+                Point coordinates for each image in the batch.
+            box_coords (list):
+                Bounding box coordinates for each image in the batch.
             device (str):
                 Device to run inference on.
 
@@ -115,7 +197,6 @@ def _encode_image(self: SAM, image: np.ndarray) -> np.ndarray:
         """Encodes the image for feature extraction."""
         self.predictor.set_image(image)
 
-    @staticmethod
     def load_weights(self: SAM, checkpoint_path: str) -> None:
         """Loads model weights from specified checkpoint."""
         self.model.load_state_dict(