[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

examples/sam-architecture.ipynb

@@ -19,8 +19,8 @@
     }
    ],
    "source": [
+    "from tiatoolbox.models.architecture.sam import SAM\n",
     "from tiatoolbox.models.engine.general_segmentor import GeneralSegmentor\n",
-    "from tiatoolbox.models.architecture.sam import SAM, SAMPrompts \n",
     "\n",
     "# abc = GeneralSegmentor(model=SAM())\n",
     "# prompts = SAMPrompts([[100,100]])\n",
@@ -42,14 +42,15 @@
    ],
    "source": [
     "from pathlib import Path\n",
+    "\n",
     "model = GeneralSegmentor(SAM())\n",
     "\n",
     "glands = \"slides/glands.png\"\n",
     "slides = \"slides/sample_wsi.svs\"\n",
-    "glands_prompts = [(370,270),(300,400)]\n",
-    "slides_prompts = [[5792,6018]]\n",
-    "slides_location = (5745,5972)\n",
-    "slides_size = (200,114)"
+    "glands_prompts = [(370, 270), (300, 400)]\n",
+    "slides_prompts = [[5792, 6018]]\n",
+    "slides_location = (5745, 5972)\n",
+    "slides_size = (200, 114)"
    ]
   },
   {
@@ -74,8 +75,9 @@
    ],
    "source": [
     "prompts = model.create_prompts(slides_prompts)\n",
-    "output = model.predict(slides, prompts, \"cpu\", \"overlays\", slides_location,slides_size, 0.5)\n",
-    "\n"
+    "output = model.predict(\n",
+    "    slides, prompts, \"cpu\", \"overlays\", slides_location, slides_size, 0.5\n",
+    ")"
    ]
   },
   {
@@ -84,9 +86,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
-    "save_path = model.to_annotation(output[0][1], output[0][2], Path(\"overlays/sample_wsi\"))\n",
-    "\n"
+    "save_path = model.to_annotation(output[0][1], output[0][2], Path(\"overlays/sample_wsi\"))"
    ]
   }
  ],

examples/tiaviz-test.ipynb

@@ -3,13 +3,10 @@
 from pathlib import Path
 from typing import Callable
 
-import numpy as np
 import pytest
-import torch
 
 from tiatoolbox.models import SAM
 from tiatoolbox.models.architecture.sam import SAMPrompts
-from tiatoolbox.models.architecture import fetch_pretrained_weights
 from tiatoolbox.utils import imread
 
 ON_GPU = False
@@ -32,21 +29,21 @@ 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]])
+    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]]
+    )
 
     model = SAM()
 
     # load pretrained weights
-    #pretrained = torch.load(weights_path, map_location="cpu")
-    #model.load_state_dict(pretrained)
+    # pretrained = torch.load(weights_path, map_location="cpu")
+    # model.load_state_dict(pretrained)
 
-    _ = model.infer_batch(model, img, on_gpu=ON_GPU) # no prompts
+    _ = 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)
-
-

tests/models/test_arch_sam.py

@@ -8,11 +8,11 @@
 from .architecture.mapde import MapDe
 from .architecture.micronet import MicroNet
 from .architecture.nuclick import NuClick
-from .architecture.sccnn import SCCNN
 from .architecture.sam import SAM
+from .architecture.sccnn import SCCNN
+from .engine.general_segmentor import GeneralSegmentor
 from .engine.multi_task_segmentor import MultiTaskSegmentor
 from .engine.nucleus_instance_segmentor import NucleusInstanceSegmentor
-from .engine.general_segmentor import GeneralSegmentor
 from .engine.patch_predictor import (
     IOPatchPredictorConfig,
     PatchDataset,
@@ -27,17 +27,17 @@
 )
 
 __all__ = [
+    "SAM",
     "SCCNN",
+    "GeneralSegmentor",
     "HoVerNet",
     "HoVerNetPlus",
     "IDaRS",
     "MapDe",
     "MicroNet",
     "MultiTaskSegmentor",
     "NuClick",
-    "SAM",
     "NucleusInstanceSegmentor",
     "PatchPredictor",
     "SemanticSegmentor",
-    "GeneralSegmentor",
 ]

tiatoolbox/models/__init__.py

@@ -2,30 +2,24 @@
 
 from __future__ import annotations
 
-from collections import OrderedDict
-
 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 tiatoolbox.models.models_abc import ModelABC
-
+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 sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
+
+from tiatoolbox.models.models_abc import ModelABC
 
 
-class SAMPrompts():
+class SAMPrompts:
     """Structure of prompts for SAM."""
-    def __init__(self, point_coords = None, point_labels = None, box_coords = None):
+
+    def __init__(self, point_coords=None, point_labels=None, box_coords=None):
         self.point_coords = None if point_coords == [] else point_coords
         self.box_coords = None if box_coords == [] else box_coords
-        if(point_coords and point_labels is None):
+        if point_coords and point_labels is None:
             self.point_labels = [1] * len(point_coords)
-        else: 
+        else:
             self.point_labels = point_labels
 
 
@@ -44,15 +38,15 @@ def __init__(
             self.model = build_sam2_hf(model_hf_path, device="cpu")
         else:
             self.model = build_sam2(model_cfg_path, checkpoint_path)
-    
+
         self.predictor = SAM2ImagePredictor(self.model)
         self.generator = SAM2AutomaticMaskGenerator(self.model)
 
     def forward(self: SAM, image: np.ndarray, prompts: SAMPrompts = None) -> np.ndarray:
         """Torch method, this contains logic for using layers defined in init."""
         mask = self.generate_mask(self, image, prompts)
         return mask
-    
+
     @staticmethod
     def infer_batch(
         model: torch.nn.Module,
@@ -78,7 +72,9 @@ def infer_batch(
         model.eval()
         model = model.to(device)
 
-        if isinstance(batch_data, torch.Tensor): # Move the tensor to the CPU if it's a PyTorch tensor
+        if isinstance(
+            batch_data, torch.Tensor
+        ):  # Move the tensor to the CPU if it's a PyTorch tensor
             batch_data = batch_data.to(device).type(torch.float32)
             batch_data = batch_data.cpu().numpy()
 
@@ -92,7 +88,7 @@ def infer_batch(
     def encode_image(self, image: np.ndarray) -> np.ndarray:
         """Encodes the image for feature extraction."""
         self.predictor.set_image(image)
-    
+
     @staticmethod
     def generate_mask(self, features: np.ndarray, prompts: SAMPrompts) -> np.ndarray:
         """Generates a segmentation mask using SAM 2, optionally guided by a prompt."""
@@ -109,24 +105,33 @@ def generate_mask(self, features: np.ndarray, prompts: SAMPrompts) -> np.ndarray
             scores = np.around(scores[sorted_ind], 2)
         else:
             masks = self.generator.generate(features)
-            scores = np.array([mask['predicted_iou'] for mask in masks])
+            scores = np.array([mask["predicted_iou"] for mask in masks])
         return masks, scores
-    
+
     @staticmethod
     def load_weights(self, checkpoint_path: str) -> None:
         """Loads model weights from specified checkpoint."""
-        self.model.load_state_dict(torch.load(checkpoint_path, map_location=self.device))
+        self.model.load_state_dict(
+            torch.load(checkpoint_path, map_location=self.device)
+        )
 
     @staticmethod
     def preproc(image: np.ndarray) -> np.ndarray:
         """Pre-processes images - Converts them into a format accepted by SAM (HWC) from NCHW."""
-        if isinstance(image, torch.Tensor): # Move the tensor to the CPU if it's a PyTorch tensor
+        if isinstance(
+            image, torch.Tensor
+        ):  # Move the tensor to the CPU if it's a PyTorch tensor
             image = image.cpu().numpy()
-        
+
         # Handle different shapes
-        if image.ndim == 4 and image.shape == (1,512,512,3):  # Case 1: (N, H, W, C)
+        if image.ndim == 4 and image.shape == (1, 512, 512, 3):  # Case 1: (N, H, W, C)
             image = np.squeeze(image, axis=0)  # Remove batch dimension
-        elif image.ndim == 4 and image.shape == (1,3,512,512):  # Case 2: (N, C, H, W)
+        elif image.ndim == 4 and image.shape == (
+            1,
+            3,
+            512,
+            512,
+        ):  # Case 2: (N, C, H, W)
             image = np.squeeze(image, axis=0)  # Remove batch dimension
             image = np.transpose(image, (1, 2, 0))  # (C, H, W) -> (H, W, C)
 
@@ -136,4 +141,4 @@ def preproc(image: np.ndarray) -> np.ndarray:
     @staticmethod
     def postproc(image: np.ndarray) -> np.ndarray:
         """Define the post-processing of this class of model."""
-        return image
+        return image