torch-em 0.9.0__tar.gz

torch_em-0.9.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Constantin Pape
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_em-0.9.0/PKG-INFO

@@ -0,0 +1,27 @@
+Metadata-Version: 2.4
+Name: torch_em
+Version: 0.9.0
+Home-page: https://github.com/constantinpape/torch-em
+Author: Constantin Pape
+License: MIT
+License-File: LICENSE
+Requires-Dist: bioimage-cpp
+Requires-Dist: bioimageio.core
+Requires-Dist: h5py
+Requires-Dist: imageio
+Requires-Dist: kornia
+Requires-Dist: mrcfile
+Requires-Dist: natsort
+Requires-Dist: python-elf>=0.9.1
+Requires-Dist: requests
+Requires-Dist: scikit-image
+Requires-Dist: tensorboard
+Requires-Dist: tifffile
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: tqdm
+Dynamic: author
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist

torch_em-0.9.0/README.md

@@ -0,0 +1,80 @@
+[![DOC](https://shields.mitmproxy.org/badge/docs-pdoc.dev-brightgreen.svg)](https://constantinpape.github.io/torch-em/torch_em.html)
+[![Build Status](https://github.com/constantinpape/torch-em/workflows/test/badge.svg)](https://github.com/constantinpape/torch-em/actions)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5108853.svg)](https://doi.org/10.5281/zenodo.5108853)
+[![Anaconda-Server Badge](https://anaconda.org/conda-forge/torch_em/badges/version.svg)](https://anaconda.org/conda-forge/torch_em)
+
+# torch-em
+
+Deep-learning based semantic and instance segmentation for 3D Electron Microscopy and other bioimage analysis problems based on PyTorch. Any feedback is highly appreciated, just open an issue!
+
+**Important:** We recently migrated `torch_em` to [bioimage-cpp](https://github.com/computational-cell-analytics/bioimage-cpp). This enables installing `torch_em` via pip (in addition to conda) but changes dependencies and may still lead to some bugs. Please report any problems you encounter in an issue; you can pin `torch_em < 0.9` to install a version based on previous dependencies.
+
+Highlights:
+- Functional API with sensible defaults to train a state-of-the-art segmentation model with a few lines of code.
+- Differentiable augmentations on GPU and CPU thanks to [kornia](https://github.com/kornia/kornia).
+- Off-the-shelf logging with [tensorboard](https://www.tensorflow.org/tensorboard) or [wandb](https://wandb.ai/site).
+- Export trained models to [bioimage.io](https://bioimage.io/#/) model format with one function call to deploy them in [ilastik](https://www.ilastik.org/documentation/nn/nn) or [deepimageJ](https://deepimagej.github.io/deepimagej/).
+
+Design:
+- All parameters are specified in code, no configuration files.
+- No callback logic; to extend the core functionality inherit from `torch_em.trainer.DefaultTrainer` instead.
+- All data-loading is lazy to support training on large datasets.
+
+`torch_em` can be installed via conda: `conda install -c conda-forge torch_em`.
+Find an example script for how to train a 2D U-Net with it below and check out the [documentation](https://constantinpape.github.io/torch-em/torch_em.html) for more details.
+
+```python
+# Train a 2d U-Net for foreground and boundary segmentation of nuclei, using data from
+# https://github.com/mpicbg-csbd/stardist/releases/download/0.1.0/dsb2018.zip
+
+import torch_em
+from torch_em.model import UNet2d
+from torch_em.data.datasets import get_dsb_loader
+
+model = UNet2d(in_channels=1, out_channels=2)
+
+# Transform to convert from instance segmentation labels to foreground and boundary probabilties.
+label_transform = torch_em.transform.BoundaryTransform(add_binary_target=True, ndim=2)
+
+# Create the training and validation data loader.
+data_path = "./dsb"  # The training data will be downloaded and saved here.
+train_loader = get_dsb_loader(
+    data_path, 
+    patch_shape=(1, 256, 256),
+    batch_size=8,
+    split="train",
+    download=True,
+    label_transform=label_transform,
+)
+val_loader = get_dsb_loader(
+    data_path, 
+    patch_shape=(1, 256, 256),
+    batch_size=8,
+    split="test",
+    label_transform=label_transform,
+)
+
+# The trainer handles the details of the training process.
+# It will save checkpoints in "checkpoints/dsb-boundary-model"
+# and the tensorboard logs in "logs/dsb-boundary-model".
+trainer = torch_em.default_segmentation_trainer(
+    name="dsb-boundary-model",
+    model=model,
+    train_loader=train_loader,
+    val_loader=val_loader,
+    learning_rate=1e-4,
+)
+trainer.fit(iterations=5000)  # Fit for 5000 iterations.
+
+# Export the trained model to the bioimage.io model format.
+from glob import glob
+import imageio
+from torch_em.util import export_bioimageio_model
+
+# Load one of the images to use as reference image.
+# Crop it to a shape that is guaranteed to fit the network.
+test_im = imageio.imread(glob(f"{data_path}/test/images/*.tif")[0])[:256, :256]
+
+# Export the model.
+export_bioimageio_model("./checkpoints/dsb-boundary-model", "./bioimageio-model", test_im)
+```

torch_em-0.9.0/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=64.0", "wheel"]
+build-backend = "setuptools.build_meta"

torch_em-0.9.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_em-0.9.0/setup.py

@@ -0,0 +1,45 @@
+import runpy
+from setuptools import setup, find_packages
+
+__version__ = runpy.run_path("torch_em/__version__.py")["__version__"]
+
+
+requires = [
+    "bioimage-cpp",
+    "bioimageio.core",
+    "h5py",
+    "imageio",
+    "kornia",
+    "mrcfile",
+    "natsort",
+    "python-elf >=0.9.1",
+    "requests",
+    "scikit-image",
+    "tensorboard",
+    "tifffile",
+    "torch",
+    "torchvision",
+    "tqdm",
+]
+
+
+setup(
+    name="torch_em",
+    packages=find_packages(exclude=["test"]),
+    version=__version__,
+    author="Constantin Pape",
+    install_requires=requires,
+    url="https://github.com/constantinpape/torch-em",
+    license="MIT",
+    entry_points={
+        "console_scripts": [
+            "torch_em.train_2d_unet = torch_em.cli:train_2d_unet",
+            "torch_em.train_3d_unet = torch_em.cli:train_3d_unet",
+            "torch_em.predict = torch_em.cli:predict",
+            "torch_em.predict_with_tiling = torch_em.cli:predict_with_tiling",
+            "torch_em.export_bioimageio_model = torch_em.util.modelzoo:main",
+            "torch_em.validate_checkpoint = torch_em.util.validation:main",
+            # "torch_em.submit_slurm = torch_em.util.submit_slurm:main",
+        ]
+    }
+)

torch_em-0.9.0/test/data/test_dataset_wrapper.py

@@ -0,0 +1,36 @@
+import os
+import unittest
+
+from torch_em.data import SegmentationDataset
+from torch_em.util.test import create_segmentation_test_data
+
+
+class TestDatasetWrapper(unittest.TestCase):
+    path = "./data.h5"
+    raw_key = "raw"
+    label_key = "labels"
+    shape = (32,) * 3
+
+    def setUp(self):
+        create_segmentation_test_data(self.path, self.raw_key, self.label_key, shape=self.shape, chunks=(32,) * 3)
+
+    def tearDown(self):
+        os.remove(self.path)
+
+    def test_wrap_dataset(self):
+        from torch_em.data import DatasetWrapper
+
+        patch_shape = (32, 32, 32)
+        ds = SegmentationDataset(self.path, self.raw_key, self.path, self.label_key, patch_shape=patch_shape)
+        wrapped_ds = DatasetWrapper(ds, lambda xy: (xy[0][:, :10, :10, :10], xy[1][:, :20, :20, :20]))
+
+        expected_shape_x = (1, 10, 10, 10)
+        expected_shape_y = (1, 20, 20, 20)
+        for i in range(3):
+            x, y = wrapped_ds[i]
+            self.assertEqual(x.shape, expected_shape_x)
+            self.assertEqual(y.shape, expected_shape_y)
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_em-0.9.0/test/data/test_image_collection_dataset.py

@@ -0,0 +1,108 @@
+import os
+import tempfile
+import unittest
+from glob import glob
+from shutil import rmtree
+
+import tifffile
+import numpy as np
+from torch_em.util.test import create_image_collection_test_data
+
+
+class TestImageCollectionDataset(unittest.TestCase):
+    folder = "./data"
+    n_images = 10
+
+    def setUp(self):
+        create_image_collection_test_data(self.folder,
+                                          n_images=self.n_images,
+                                          min_shape=(256, 256),
+                                          max_shape=(512, 512))
+
+    def tearDown(self):
+        rmtree(self.folder, ignore_errors=True)
+
+    def test_dataset(self):
+        from torch_em.data import ImageCollectionDataset
+        patch_shape = (256, 256)
+
+        raw_paths = glob(os.path.join(self.folder, "images", "*.tif"))
+        label_paths = glob(os.path.join(self.folder, "labels", "*.tif"))
+        ds = ImageCollectionDataset(raw_paths, label_paths,
+                                    patch_shape=patch_shape)
+        self.assertEqual(len(ds), self.n_images)
+        self.assertEqual(ds._ndim, 2)
+
+        expected_shape = (1,) + patch_shape
+        for i in range(10):
+            x, y = ds[i]
+            self.assertEqual(x.shape, expected_shape)
+            self.assertEqual(y.shape, expected_shape)
+
+    def test_dataset_with_sampler(self):
+        from torch_em.data import ImageCollectionDataset
+        from torch_em.data.sampler import MinIntensitySampler
+        patch_shape = (256, 256)
+
+        raw_paths = glob(os.path.join(self.folder, "images", "*.tif"))
+        label_paths = glob(os.path.join(self.folder, "labels", "*.tif"))
+        sampler = MinIntensitySampler(min_intensity=0.51, p_reject=0.9)
+        ds = ImageCollectionDataset(raw_paths, label_paths,
+                                    patch_shape=patch_shape,
+                                    sampler=sampler)
+        self.assertEqual(len(ds), self.n_images)
+        self.assertEqual(ds._ndim, 2)
+
+        expected_shape = (1,) + patch_shape
+        for i in range(10):
+            x, y = ds[i]
+            self.assertEqual(x.shape, expected_shape)
+            self.assertEqual(y.shape, expected_shape)
+
+
+def generate_sample_data(folder, n_images, image_shape, label_shape):
+    im_folder = os.path.join(folder, "images")
+    label_folder = os.path.join(folder, "labels")
+    os.makedirs(im_folder)
+    os.makedirs(label_folder)
+    for i in range(n_images):
+        raw = np.empty(image_shape, dtype=np.uint8)
+        label = np.ones(label_shape, dtype=np.float32)
+        tifffile.imwrite(os.path.join(im_folder, f"test_{i}.tif"), raw)
+        tifffile.imwrite(os.path.join(label_folder, f"test_{i}.tif"), label)
+
+
+class TestImageCollectionDatasetWithChannels(unittest.TestCase):
+    def test_channel_end(self):
+        from torch_em.data import ImageCollectionDataset
+
+        patch_shape = (256, 256)
+
+        with tempfile.TemporaryDirectory() as td:
+            generate_sample_data(td, 10, (256, 256, 2), (256, 256))
+            raw_paths = glob(os.path.join(td, "images", "*.tif"))
+            label_paths = glob(os.path.join(td, "labels", "*.tif"))
+
+            ds = ImageCollectionDataset(raw_paths, label_paths,
+                                        patch_shape=patch_shape)
+            self.assertEqual(len(ds), 10)
+            self.assertEqual(ds._get_sample(0)[0].shape[0], 2)
+
+    def test_channel_begin(self):
+        from torch_em.data import ImageCollectionDataset
+
+        patch_shape = (256, 256)
+
+        with tempfile.TemporaryDirectory() as td:
+            generate_sample_data(td, 10, (2, 256, 256), (256, 256))
+            raw_paths = glob(os.path.join(td, "images", "*.tif"))
+            label_paths = glob(os.path.join(td, "labels", "*.tif"))
+
+            ds = ImageCollectionDataset(raw_paths, label_paths,
+                                        patch_shape=patch_shape)
+            self.assertEqual(len(ds), 10)
+            self.assertEqual(ds._get_sample(0)[0].shape[0], 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_em-0.9.0/test/data/test_raw_dataset.py

@@ -0,0 +1,173 @@
+import os
+import unittest
+
+import h5py
+import numpy as np
+
+from torch_em.data.raw_dataset import RawDataset, RawDatasetWithMasks
+
+
+class TestRawDatasetBase:
+    """Base class to hold shared tests for RawDataset and subclasses."""
+
+    path = "./data.h5"
+    dataset_class = None
+
+    def tearDown(self):
+        os.remove(self.path)
+
+    def create_default_data(self, key, shape=None, chunks=None):
+        shape = (128,) * 3 if shape is None else shape
+        chunks = tuple(min(32, sh) for sh in shape) if chunks is None else chunks
+        with h5py.File(self.path, "a") as f:
+            f.create_dataset(key, data=np.random.rand(*shape), chunks=chunks)
+        return shape, chunks
+
+    def test_dataset_3d(self):
+        raw_key = "raw"
+        shape, chunks = self.create_default_data(raw_key)
+        patch_shape = chunks
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape)
+        self.assertEqual(ds.raw.shape, shape)
+        self.assertEqual(ds._ndim, 3)
+        expected_shape = (1,) + patch_shape
+        for i in range(10):
+            self.assertEqual(ds[i].shape, expected_shape)
+
+    def test_dataset_2d(self):
+        raw_key = "raw"
+        shape, chunks = self.create_default_data(raw_key)
+        patch_shape = (1, 32, 32)
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape, ndim=2)
+        self.assertEqual(ds.raw.shape, shape)
+        self.assertEqual(ds._ndim, 2)
+        expected_shape = patch_shape
+        for i in range(10):
+            self.assertEqual(ds[i].shape, expected_shape)
+
+    def test_dataset_4d(self):
+        raw_key = "raw"
+        shape = (12, 64, 64, 64)
+        self.create_default_data(raw_key, shape=shape)
+        patch_shape = (3, 32, 32, 32)
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape)
+        self.assertEqual(ds.raw.shape, shape)
+        self.assertEqual(ds._ndim, 4)
+        expected_shape = patch_shape
+        for i in range(10):
+            self.assertEqual(ds[i].shape, expected_shape)
+
+    def test_roi(self):
+        raw_key = "raw"
+        self.create_default_data(raw_key)
+        patch_shape = (32, 32, 32)
+        roi = np.s_[32:96, 32:96, 32:96]
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape, roi=roi)
+        roi_shape = 3 * (64,)
+        self.assertEqual(ds.raw.shape, roi_shape)
+        expected_shape = (1,) + patch_shape
+        for i in range(10):
+            self.assertEqual(tuple(ds[i].shape), expected_shape)
+
+    def test_with_channels(self):
+        raw_key = "raw"
+        self.create_default_data(raw_key, shape=(3, 128, 128, 128), chunks=(1, 32, 32, 32))
+        patch_shape = (32, 32, 32)
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape, ndim=3, with_channels=True)
+        self.assertEqual(ds._ndim, 3)
+        expected_raw_shape = (3,) + patch_shape
+        for i in range(10):
+            self.assertEqual(ds[i].shape, expected_raw_shape)
+
+
+class TestRawDataset(TestRawDatasetBase, unittest.TestCase):
+    dataset_class = RawDataset
+
+
+class TestRawDatasetWithMasks(TestRawDatasetBase, unittest.TestCase):
+    dataset_class = RawDatasetWithMasks
+    mrc_path = "./mask.mrc"
+    
+    def tearDown(self):
+        super().tearDown()
+
+        if os.path.exists(self.mrc_path):
+            os.remove(self.mrc_path)
+    
+    def create_default_mrc_data(self, data):
+        import mrcfile
+        with mrcfile.new(self.mrc_path, overwrite=True) as f:
+            f.set_data(data.astype(np.float32))
+
+    # subclass-specific tests, designed for a 3d dataset
+    def test_sample_mask(self):
+        from torch_em.data.sampler import MinForegroundSampler
+
+        raw_key, mask_key = "raw", "mask"
+        shape = (128, 128, 128)
+        chunks = (32, 32, 32)
+        
+        # define central z slices 40:80 as the sample ROI, above and below is empty
+        sample_mask = np.zeros(shape, dtype=np.float32)
+        sample_mask[40:80, :, :] = 1.0
+
+        with h5py.File(self.path, "a") as f:
+            f.create_dataset(raw_key, data=sample_mask, chunks=chunks)
+            f.create_dataset(mask_key, data=sample_mask, chunks=chunks)
+
+        patch_shape = chunks
+
+        min_fraction = 0.95
+        sampler = MinForegroundSampler(min_fraction=min_fraction)
+
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape, sampler=sampler, 
+                                sample_mask_path=self.path, sample_mask_key=mask_key)
+
+        self.assertEqual(ds.raw.shape, shape)
+        self.assertEqual(ds.sample_mask.shape, shape)
+        self.assertEqual(ds._ndim, 3)
+
+        expected_shape = (1,) + patch_shape
+        for i in range(10):
+            patch = ds[i]
+            self.assertEqual(patch.shape, expected_shape)
+
+            # check that patches are sampled inside the ROI defined by sample mask
+            # note that the raw data and the sample mask are identical, so this checks if the
+            # positive samples from the sample mask are correct
+            self.assertGreaterEqual(patch.mean().item(), min_fraction)
+
+    def test_bg_mask(self):
+        raw_key, mask_key = "raw", "mask"
+
+        shape, chunks = self.create_default_data(raw_key)
+        self.create_default_data(mask_key)
+
+        patch_shape = chunks
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape, 
+                                bg_mask_path=self.path, bg_mask_key=mask_key)
+
+        self.assertEqual(ds.raw.shape, shape)
+        self.assertEqual(ds.bg_mask.shape, shape)
+        self.assertEqual(ds._ndim, 3)
+        expected_shape = (2,) + patch_shape
+
+        for i in range(10):
+
+            patch = ds[i]
+            self.assertEqual(patch.shape, expected_shape)
+
+    def test_bg_mask_mrc(self):
+        raw_key, mask_key = "raw", None
+
+        shape, chunks = self.create_default_data(raw_key)
+        self.create_default_mrc_data(np.zeros(shape))
+
+        patch_shape = chunks
+        ds = self.dataset_class(self.path, raw_key, patch_shape=patch_shape, 
+                                bg_mask_path=self.mrc_path, bg_mask_key=None)
+        
+        self.assertEqual(ds.bg_mask.shape, shape)
+        
+if __name__ == "__main__":
+    unittest.main()

torch_em-0.9.0/test/data/test_sampler.py

@@ -0,0 +1,35 @@
+import unittest
+import numpy as np
+
+
+class TestSampler(unittest.TestCase):
+    def test_min_foreground_sampler(self):
+        from torch_em.data import MinForegroundSampler
+        threshold = 0.5
+        sampler = MinForegroundSampler(min_fraction=threshold)
+
+        shape = (32, 32)
+        n_samples = 100
+        for _ in range(n_samples):
+            data = np.random.rand(*shape)
+            labels = (np.random.rand(*shape) > 0.5).astype("float32")
+            accept = sampler(data, labels)
+            expected = (labels.sum() / labels.size) > threshold
+            self.assertEqual(accept, expected)
+
+    def test_min_intensity_sampler(self):
+        from torch_em.data import MinIntensitySampler
+        threshold = 0.5
+        sampler = MinIntensitySampler(min_intensity=threshold)
+        shape = (32, 32)
+
+        n_samples = 100
+        for _ in range(n_samples):
+            data = np.random.rand(*shape)
+            accept = sampler(data)
+            expected = np.median(data) > threshold
+            self.assertEqual(accept, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()