Improvement documentation and testing (#70)

* Update to uv setup.

* improve test coverage

* improve documentation

* update uv.lock

* add more docstrings

* update documentation

* update documentation

* update tests

* make rfest optional

* update readme

* update workflow

* update dependecies

* move mb test file

* update workflow

* update Python version requirements in workflow and project configuration

Author: jonathanoesterle

.github/workflows/python-app.yml

@@ -12,33 +12,33 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [ "3.8", "3.9", "3.10", "3.11" ]
+        python-version: [ "3.10", "3.11", "3.12" ]
 
     steps:
-      # Checkout
       - name: Checkout repo
-        uses: actions/checkout@v2
-        # Setup Python
-      - name: Set Up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
-        # Display Python version
+
       - name: Display Python version
         run: python -c "import sys; print(sys.version)"
-        # Install dependencies
+
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
-          pip install flake8 pytest
-          if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-        # Run tests
+          uv pip install --system flake8 pytest
+          uv pip install --system -e .
+
       - name: Lint with flake8
         run: |
-          # stop the build if there are Python syntax errors or undefined names
           flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
           flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+
       - name: Test with pytest
         run: |
           pytest -v

.gitignore

@@ -6,6 +6,7 @@ __pycache__/
 # Files
 *.pickle
 *.pkl
+test_data/*
 
 # C extensions
 *.so
@@ -119,4 +120,4 @@ djconfig/
 
 # Exclude user stuff
 user/
-!user/README.md
+!user/README.md

README.md

@@ -21,17 +21,18 @@ Download the package:
 git clone https://github.com/eulerlab/djimaging.git
 ````
 
-Install the package e.g. using pip as an editable package:
+Install the package as an editable package using [uv](https://docs.astral.sh/uv/):
 
 ```bash
 cd djimaging
-pip install -r requirements.txt
-pip install -e .
+uv pip install -e .
 ```
 
-If you want to use autorois, also install the following:
-```
-pip install -r requirements-autorois.txt
+If you want to use autorois or receptive fields, install with the optional extras:
+```bash
+uv pip install -e ".[autorois]"
+uv pip install -e ".[rf]"
+uv pip install -e ".[autorois,rf]"
 ```
 
 > ❗ To test if the package was successfully installed, e.g.

djimaging/autorois/autorois_utils.py

@@ -10,20 +10,58 @@
 
 
 @lru_cache(2)
-def _get_neighbor_kernel(device: torch.device):
+def _get_neighbor_kernel(device: "torch.device") -> "torch.Tensor":
+    """Return a 4-connected neighbor kernel (cross-shaped) as a cached tensor.
+
+    Parameters
+    ----------
+    device : torch.device
+        The device on which the kernel tensor will be placed.
+
+    Returns
+    -------
+    torch.Tensor
+        A 4D tensor of shape (1, 1, 3, 3) with the 4-connected neighbor pattern.
+    """
     return torch.tensor([[[[0, 1, 0], [1, 0, 1], [0, 1, 0]]]],
                         dtype=torch.float32,
                         device=device)
 
 
 @lru_cache(2)
-def _get_diagonal_neighbor_kernel(device: torch.device):
+def _get_diagonal_neighbor_kernel(device: "torch.device") -> "torch.Tensor":
+    """Return an 8-connected neighbor kernel (full 3x3 except center) as a cached tensor.
+
+    Parameters
+    ----------
+    device : torch.device
+        The device on which the kernel tensor will be placed.
+
+    Returns
+    -------
+    torch.Tensor
+        A 4D tensor of shape (1, 1, 3, 3) with the 8-connected neighbor pattern.
+    """
     return torch.tensor([[[[1, 1, 1], [1, 0, 1], [1, 1, 1]]]],
                         dtype=torch.float32,
                         device=device)
 
 
-def _run_convolution(mask: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
+def _run_convolution(mask: "torch.Tensor", kernel: "torch.Tensor") -> "torch.Tensor":
+    """Convolve a 2D boolean mask with the given kernel and return a boolean output.
+
+    Parameters
+    ----------
+    mask : torch.Tensor
+        2D boolean tensor to convolve.
+    kernel : torch.Tensor
+        4D convolution kernel of shape (1, 1, kH, kW).
+
+    Returns
+    -------
+    torch.Tensor
+        2D boolean tensor of the same shape as ``mask``.
+    """
     mask_float = mask.to(torch.float32)
     mask4d = mask_float.unsqueeze(dim=0).unsqueeze(dim=0)
     conv_out = torch.nn.functional.conv2d(mask4d, kernel, padding=1)[0, 0]
@@ -32,7 +70,26 @@ def _run_convolution(mask: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
     return conv_out_bool
 
 
-def make_mask_ids_consecutive(mask: np.array) -> np.array:
+def make_mask_ids_consecutive(mask: np.ndarray) -> np.ndarray:
+    """Remap ROI mask values to a consecutive integer range starting at 0.
+
+    Parameters
+    ----------
+    mask : np.ndarray
+        Integer array where 0 is background and positive integers are ROI IDs.
+        Must contain 0 as the minimum value (cellpose format).
+
+    Returns
+    -------
+    np.ndarray
+        Mask with the same shape as ``mask`` but with IDs remapped to
+        ``{0, 1, 2, ..., n_rois}``.
+
+    Raises
+    ------
+    AssertionError
+        If the minimum value of ``mask`` is not 0.
+    """
     values = np.unique(mask)
     assert values.min() == 0, "Mask has to be in cellpose format"
     if (values.max() + 1) == values.shape[0]:
@@ -46,27 +103,44 @@ def make_mask_ids_consecutive(mask: np.array) -> np.array:
 
 
 def create_mask(
-        cell_probs: torch.Tensor,
-        offsets: torch.Tensor,
-        center_mask: torch.Tensor,
+        cell_probs: "torch.Tensor",
+        offsets: "torch.Tensor",
+        center_mask: "torch.Tensor",
         cell_prob_threshold: float = 0.5,
         kernel_size: int = 5,
         center_prob_threshold: float = 0.1,
         max_number_of_cells: int = 99999,
-) -> torch.Tensor:
-    """
-    Create a roi mask from the outputs of the instance unet model.
+) -> "torch.Tensor":
+    """Create a ROI mask from the outputs of the instance UNet model.
+
     See https://github.com/bowenc0221/panoptic-deeplab/blob/master/segmentation/model/post_processing/instance_post_processing.py
-    for another implementation
-
-    Args:
-        cell_probs: probability that a pixel contains a cell, shape: [dim_x, dim_y]
-        offsets: offset to the cell center in x and y direction for each pixel, shape [2, dim_x, dim_y]
-        center_mask: 'probability' that a pixel is the center of the cell, shape: [dim_x, dim_y]
-        cell_prob_threshold: threshold that we consider a pixel to be a cell
-        kernel_size: mentioned parameter in section 3.2 of https://arxiv.org/pdf/1911.10194.pdf, changed to 5
-        center_prob_threshold: probability threshold that to determine whether a pixel is a center
-        max_number_of_cells: maximum number of cells per mask
+    for another implementation.
+
+    Parameters
+    ----------
+    cell_probs : torch.Tensor
+        Probability that a pixel contains a cell, shape: (dim_x, dim_y).
+    offsets : torch.Tensor
+        Offset to the cell center in x and y direction for each pixel,
+        shape: (2, dim_x, dim_y).
+    center_mask : torch.Tensor
+        Probability that a pixel is the center of a cell, shape: (dim_x, dim_y).
+    cell_prob_threshold : float, optional
+        Threshold above which a pixel is considered a cell. Default is 0.5.
+    kernel_size : int, optional
+        Max-pooling kernel size used to find center peaks (Section 3.2 of
+        https://arxiv.org/pdf/1911.10194.pdf). Default is 5.
+    center_prob_threshold : float, optional
+        Probability threshold for a pixel to be considered a center.
+        Default is 0.1.
+    max_number_of_cells : int, optional
+        Maximum number of cells per mask. Default is 99999.
+
+    Returns
+    -------
+    torch.Tensor
+        Integer tensor of the same spatial shape as ``cell_probs`` where each
+        pixel is labelled with its instance ID (0 = background).
     """
     # determine pixels that are cells
     pixel_is_cell = cell_probs > cell_prob_threshold
@@ -123,11 +197,19 @@ def create_mask(
     return roi_mask_consecutive
 
 
-def clean_roi_mask(roi_mask: torch.Tensor) -> torch:
-    """"
-    (1) filter pixel of a cell that don't have a horizontal and vertical neighboring pixel
-    (2) throw out cells that have less than 3 pixels
-    Possible Todo: throw out pixels of a cell that are horizontal or vertical neighbors of another cell
+def clean_roi_mask(roi_mask: "torch.Tensor") -> None:
+    """Remove isolated pixels and small ROIs from a ROI mask in-place.
+
+    Steps applied:
+    (1) Remove pixels of a cell that have no horizontal or vertical neighbour
+        belonging to the same cell.
+    (2) Remove cells that have fewer than 3 pixels.
+
+    Parameters
+    ----------
+    roi_mask : torch.Tensor
+        Integer 2D tensor where 0 is background and positive integers are ROI
+        IDs. Modified in-place.
     """
     # Filter pixels that don't have any horizontal or vertical neighboring pixel
     kernel = _get_neighbor_kernel(roi_mask.device)