diff --git a/.github/workflows/testing.yaml b/.github/workflows/testing.yaml index 8ae7e4d4e..b05da241f 100644 --- a/.github/workflows/testing.yaml +++ b/.github/workflows/testing.yaml @@ -46,6 +46,9 @@ jobs: files: | **/*.py + - name: Install libboost + if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }} + run: sudo apt install libboost-dev - name: Setup Python if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }} uses: actions/setup-python@v4 @@ -107,6 +110,9 @@ jobs: with: files: | **/*.py + - name: Install libboost + if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }} + run: sudo apt install libboost-dev - name: Setup Python if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }} uses: actions/setup-python@v4 @@ -155,6 +161,9 @@ jobs: with: files: | **/*.py + - name: Install libboost + if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }} + run: sudo apt install libboost-dev - name: Setup Python if: ${{ steps.changed-py-files.outputs.any_changed == 'true' }} uses: actions/setup-python@v4 @@ -207,6 +216,9 @@ jobs: with: files: | **/*.rst + - name: Install libboost + if: ${{ steps.changed-py-files.outputs.any_changed == 'true' || steps.changed-docs-files.outputs.any_changed == 'true'}} + run: sudo apt install libboost-dev - name: Setup Python if: ${{ steps.changed-py-files.outputs.any_changed == 'true' || steps.changed-docs-files.outputs.any_changed == 'true'}} uses: actions/setup-python@v4 diff --git a/.github/workflows/testing_integration.yaml b/.github/workflows/testing_integration.yaml index bfb0c575e..2c42a17db 100644 --- a/.github/workflows/testing_integration.yaml +++ b/.github/workflows/testing_integration.yaml @@ -25,7 +25,6 @@ jobs: private-key: ${{ secrets.APP_PEM }} # owner is required, otherwise the creds will fail the checkout step owner: ${{ github.repository_owner }} - - name: Checkout from GitHub uses: actions/checkout@v4 with: @@ -33,6 +32,8 @@ jobs: submodules: true ssh-key: ${{ secrets.git_ssh_key }} token: ${{ steps.app_token.outputs.token }} + - name: Install libboost + run: sudo apt install libboost-dev - name: Setup Python uses: actions/setup-python@v4 with: diff --git a/docker/Dockerfile.all.p310 b/docker/Dockerfile.all.p310 index 57130b607..8cd25aa12 100644 --- a/docker/Dockerfile.all.p310 +++ b/docker/Dockerfile.all.p310 @@ -4,7 +4,7 @@ FROM nvidia/cuda:12.2.0-runtime-ubuntu20.04 ENV DEBIAN_FRONTEND="noninteractive" RUN apt-get update \ - && apt-get install -y git build-essential wget curl vim ffmpeg libsm6 libxext6 software-properties-common unixodbc-dev \ + && apt-get install -y git build-essential wget curl vim ffmpeg libsm6 libxext6 software-properties-common unixodbc-dev libboost-dev \ && add-apt-repository ppa:deadsnakes/ppa -y \ && apt-get update \ && apt-get install -y --no-install-recommends python3.10 python3-pip python3.10-dev \ @@ -13,9 +13,6 @@ RUN apt-get update \ && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \ && python3.10 get-pip.py \ && pip install posix-ipc \ - && apt-get --purge autoremove -y build-essential \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* \ && mkdir /opt/cue \ && cd /opt/cue \ && wget https://github.com/cue-lang/cue/releases/download/v0.4.3/cue_v0.4.3_linux_amd64.tar.gz \ @@ -27,4 +24,9 @@ WORKDIR /opt/zetta_utils ADD pyproject.toml /opt/zetta_utils/ RUN pip install '.[modules]' COPY . /opt/zetta_utils/ + +RUN apt-get --purge autoremove -y build-essential \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + RUN zetta --help diff --git a/docker/Dockerfile.all.p311 b/docker/Dockerfile.all.p311 index a12e17b6b..2fdfd1b77 100644 --- a/docker/Dockerfile.all.p311 +++ b/docker/Dockerfile.all.p311 @@ -4,7 +4,7 @@ FROM nvidia/cuda:12.2.0-runtime-ubuntu20.04 ENV DEBIAN_FRONTEND="noninteractive" RUN apt-get update \ - && apt-get install -y git build-essential wget curl vim ffmpeg libsm6 libxext6 software-properties-common unixodbc-dev \ + && apt-get install -y git build-essential wget curl vim ffmpeg libsm6 libxext6 software-properties-common unixodbc-dev libboost-dev \ && add-apt-repository ppa:deadsnakes/ppa -y \ && apt-get update \ && apt-get install -y --no-install-recommends python3.11 python3-pip python3.11-dev \ @@ -13,9 +13,6 @@ RUN apt-get update \ && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \ && python3.11 get-pip.py \ && pip install posix-ipc \ - && apt-get --purge autoremove -y build-essential \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* \ && mkdir /opt/cue \ && cd /opt/cue \ && wget https://github.com/cue-lang/cue/releases/download/v0.4.3/cue_v0.4.3_linux_amd64.tar.gz \ @@ -27,4 +24,9 @@ WORKDIR /opt/zetta_utils ADD pyproject.toml /opt/zetta_utils/ RUN pip install '.[modules]' COPY . /opt/zetta_utils/ + +RUN apt-get --purge autoremove -y build-essential \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + RUN zetta --help diff --git a/pyproject.toml b/pyproject.toml index 7b47aef0d..6f2f58b5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,6 +118,8 @@ segmentation = [ "onnx >= 1.13.0", "onnxruntime-gpu >= 1.13.1", "scikit-learn >= 1.2.2", + "lsds @ git+https://github.com/ZettaAI/lsd.git@cebe976", + "abiss @ git+https://github.com/ZettaAI/abiss.git@1d1fc27", ] tensor-ops = [ "zetta_utils[tensor_typing]", diff --git a/tests/unit/mazepa_layer_processing/segmentation/test_watershed.py b/tests/unit/mazepa_layer_processing/segmentation/test_watershed.py new file mode 100644 index 000000000..42e086599 --- /dev/null +++ b/tests/unit/mazepa_layer_processing/segmentation/test_watershed.py @@ -0,0 +1,24 @@ +import pytest +import torch + +from zetta_utils.mazepa_layer_processing.segmentation import watershed + + +@pytest.mark.parametrize( + "fragments_in_xy", + [False, True], +) +def test_ws_dummy_data_lsd(fragments_in_xy): + affs = torch.zeros(3, 8, 8, 8) + ret = watershed.watershed_from_affinities(affs, method="lsd", fragments_in_xy=fragments_in_xy) + assert ret.shape == (1, 8, 8, 8) + + +@pytest.mark.parametrize( + "size_threshold", + [0, 200], +) +def test_ws_dummy_data_abiss(size_threshold): + affs = torch.zeros(3, 8, 8, 8) + ret = watershed.watershed_from_affinities(affs, method="abiss", size_threshold=size_threshold) + assert ret.shape == (1, 8, 8, 8) diff --git a/zetta_utils/mazepa_layer_processing/segmentation/__init__.py b/zetta_utils/mazepa_layer_processing/segmentation/__init__.py index a9ac823b7..bcbeeadda 100644 --- a/zetta_utils/mazepa_layer_processing/segmentation/__init__.py +++ b/zetta_utils/mazepa_layer_processing/segmentation/__init__.py @@ -1 +1,2 @@ -from . import masks +from . import masks, watershed +from zetta_utils.internal import segmentation diff --git a/zetta_utils/mazepa_layer_processing/segmentation/watershed.py b/zetta_utils/mazepa_layer_processing/segmentation/watershed.py new file mode 100644 index 000000000..b2a37f2d0 --- /dev/null +++ b/zetta_utils/mazepa_layer_processing/segmentation/watershed.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from typing import Literal + +import abiss +import einops +import numpy as np +import torch +from lsd.post.fragments import watershed_from_affinities as _watershed_lsd + +from zetta_utils import builder + + +def _run_watershed_abiss( + affs: torch.Tensor, + aff_threshold_low: float = 0.01, + aff_threshold_high: float = 0.99, + size_threshold: int = 0, + # agglomeration_threshold: float = 0.0, +) -> torch.Tensor: + """ + Args: + affs: + Affinity tensor in float32 with values [0.0, 1.0]. + + aff_threshold_low, aff_threshold_high: + Low and high watershed thresholds. + + size_threshold: + If greater than 0, perform single-linkage merging as a subsequent + step. + + agglomeration_threshold: + If greater than 0.0, perform agglomeration as a subsequent + step with this threshold. + """ + affs = torch.nn.functional.pad(affs, (1, 1, 1, 1, 1, 1)) # abiss requires 1px padding + affs = einops.rearrange(affs, "C X Y Z -> X Y Z C") # channel last + ret = abiss.watershed( + affs=affs.numpy(), + aff_threshold_low=aff_threshold_low, + aff_threshold_high=aff_threshold_high, + size_threshold=size_threshold, + # agglomeration_threshold=agglomeration_threshold, + ) + ret = ret[1:-1, 1:-1, 1:-1] + ret = np.expand_dims(ret, axis=0) + return ret + + +def _run_watershed_lsd( + affs: torch.Tensor, + fragments_in_xy: bool = False, + min_seed_distance: int = 10, + affs_in_xyz: bool = True, +) -> torch.Tensor: + """ + Args: + affs: + Affinity tensor in either float32 or uint8. + + fragments_in_xy: + Produce supervoxels in xy. + + min_seed_distance: + Controls distance between seeds in voxels. + """ + """ + TODO: + - add supervoxel filtering based on average aff value + - add option to also perform agglomeration + """ + affs_np = einops.rearrange(affs, "C X Y Z -> C Z Y X").numpy() + if affs_in_xyz: + # aff needs to be zyx + affs_np = np.flip(affs_np, 0) + if affs_np.dtype == np.uint8: + max_affinity_value = 255.0 + affs_np = affs_np.astype(np.float32) + else: + max_affinity_value = 1.0 + + ret, _ = _watershed_lsd( + affs=affs_np, + max_affinity_value=max_affinity_value, + fragments_in_xy=fragments_in_xy, + min_seed_distance=min_seed_distance, + ) + ret = einops.rearrange(ret, "Z Y X -> X Y Z") + ret = np.expand_dims(ret, axis=0) + return ret + + +@builder.register("watershed_from_affinities") +def watershed_from_affinities( + affs: torch.Tensor, # in CXYZ + method: Literal["abiss", "lsd"], + **kwargs, +) -> torch.Tensor: + """ + Produce supervoxels by running watershed on aff data. Optionally perform + agglomeration and output segmentation. + """ + if method == "lsd": + seg = _run_watershed_lsd(affs, **kwargs) + elif method == "abiss": + seg = _run_watershed_abiss(affs, **kwargs) + """ + TODO: write a wrapper for multi-chunk watershed that performs: + - relabel supervoxels based on chunkid & chunk size + - add supervoxel filtering based on mask + - store a list of supervoxels within a chunk to a database + """ + return seg