This repository provides a command-line tool for performing semantic segmentation on H&E-stained histopathology images using deep learning models (U-Net variants).
It processes large images via patch extraction, runs model inference, stitches predictions and exports the final segmentation mask as an OME-TIFF file.
-
Supports multiple architectures:
- U-Net (U-Net paper) (trained model DOI)
- U-NeXt (Influenced from the paper) (trained model DOI)
- Context U-Net (CU-Net) (Influenced from the paper) (trained model DOI)
-
Automatic model download from zenodo
-
Patch-based inference for large images
-
Mask stitching
-
OME-TIFF output
-
Optional GPU acceleration (CUDA)
All models are trained for a 7-class segmentation task.
- 0: background
- 1: neutrophil
- 2: epithelial
- 3: lymphocyte
- 4: plasma
- 5: eosinophil
- 6: connective
Clone the repository:
git clone <url>
cd stainsegmy
Create a conda environment:
conda env create -f environment.yml
conda activate stainsegmy-env- For GPU
pip install torch==2.1.1 torchvision==0.16.1 \
--index-url https://download.pytorch.org/whl/cu118- For CPU
pip install torch==2.1.1 torchvision==0.16.1Or install as a package:
pip install .
- Image format: TIFF / OME-TIFF
- Shape: (C, Y, X) or (Y, X, C)
- Channels: 3 (RGB)
Run the segmentation pipeline via CLI:
python main.py \
--input path/to/image.tif \
--output path/to/output_folder \
--architecture U-Net \
--cuda
| Option | Description |
|---|---|
-i, --input |
Path to input image |
-o, --output |
Output directory |
-m, --model |
Path to model checkpoint (optional) |
--architecture |
Model type: U-Net, U-NeXt, CU-Net |
-c / -nc |
Enable/disable CUDA |
-s / -ns |
Delete downloaded model after prediction |
--version |
Show version |
-
File:
Segmentation_mask.ome.tif -
Format: OME-TIFF
-
Includes:
- Class labels
This project is licensed under the MIT License. See LICENSE file for details.
This repository is written by Thusheera Kumarasekara and some parts of the project is adapted from the work of Dominik Molitor.