├── .dockerignore ├── .github ├── FUNDING.yml └── workflows │ └── python-app.yml ├── .gitignore ├── .pylintrc ├── .readthedocs.yaml ├── .travis.yml ├── CHANGELOG.md ├── CODE_OF_CONDUCT.md ├── Dockerfile ├── LICENSE ├── MANIFEST.in ├── README.md ├── chunkflow ├── __init__.py ├── __main__.py ├── __version__.py ├── chunk │ ├── __init__.py │ ├── affinity_map │ │ ├── __init__.py │ │ └── base.py │ ├── base.py │ ├── image │ │ ├── __init__.py │ │ ├── adjust_grey.py │ │ └── base.py │ ├── probability_map.py │ ├── segmentation.py │ └── validate.py ├── flow │ ├── __init__.py │ ├── base.py │ ├── divid_conquer │ │ ├── __init__.py │ │ ├── inferencer.py │ │ ├── patch │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── identity.py │ │ │ ├── patch_mask.py │ │ │ ├── pytorch.py │ │ │ ├── pznet.py │ │ │ └── universal.py │ │ └── transform.py │ ├── downsample_upload.py │ ├── flow.py │ ├── load_pngs.py │ ├── load_precomputed.py │ ├── log_summary.py │ ├── mask.py │ ├── mesh.py │ ├── mesh_manifest.py │ ├── napari.py │ ├── neuroglancer.py │ ├── plugin.py │ ├── save_pngs.py │ ├── save_precomputed.py │ ├── setup_env.py │ └── view.py ├── lib │ ├── __init__.py │ ├── aws │ │ ├── __init__.py │ │ └── sqs_queue.py │ ├── cartesian_coordinate.py │ ├── flow.py │ ├── gala │ │ ├── __init__.py │ │ └── evaluate.py │ ├── igneous │ │ ├── __init__.py │ │ ├── downsample.py │ │ ├── downsample_scales.py │ │ └── tasks.py │ ├── region_of_interest.py │ └── utils.py ├── plugins │ ├── agglomerate.py │ ├── aggregate_skeleton_fragments.py │ ├── aws │ │ ├── cloud_watch.py │ │ └── test_cloud_watch.py │ ├── clahe.py │ ├── cutout_dvid_label.py │ ├── czann_inference.py │ ├── detect_points.py │ ├── fill_segmentation_holes.py │ ├── gaussian_filter.py │ ├── inverse.py │ ├── load_mrc.py │ ├── load_n5.py │ ├── load_nrrd.py │ ├── load_tensorstore.py │ ├── mapto01.py │ ├── median_filter.py │ ├── napari_pixel_classifier.py │ ├── print_max_id.py │ ├── requirements.txt │ ├── rescale_intensity.py │ ├── save_nrrd.py │ ├── skeletonize.py │ ├── statistics.py │ ├── stretch_intensity.py │ ├── synapse │ │ ├── adjust_pre.py │ │ ├── detect_duplicate_post.py │ │ ├── detect_post_synapses.py │ │ ├── detect_pre_synapses.py │ │ └── find_tbar_object.py │ └── transpose.py ├── point_cloud.py ├── pyproject.toml ├── synapses.py └── volume.py ├── conftest.py ├── count_operators.sh ├── dev-environments.txt ├── distributed ├── kubernetes │ ├── .dockerignore │ ├── README.md │ ├── connect_to_cluster.sh │ └── deploy.yml └── restapi │ ├── server.py │ ├── task.py │ └── worker.py ├── docker ├── base │ └── Dockerfile └── inference │ └── Dockerfile.pytorch ├── docs ├── Makefile ├── README.md ├── logo │ ├── CMYK_print │ │ ├── Chunkflow_logo_CMYK.eps │ │ ├── Chunkflow_logo_CMYK.jpg │ │ ├── Chunkflow_logo_CMYK.pdf │ │ └── Chunkflow_logo_CMYK.png │ └── RGB_web │ │ ├── Chunkflow_logo_RBG.eps │ │ ├── Chunkflow_logo_RBG.jpg │ │ ├── Chunkflow_logo_RBG.png │ │ └── Chunkflow_logo_RBG.svg ├── make.bat ├── requirements.txt └── source │ ├── _static │ ├── custom.css │ └── image │ │ ├── cleft.png │ │ ├── cleft_label.png │ │ ├── cremi_image.png │ │ ├── image_aff.png │ │ ├── image_seg.png │ │ ├── log_summary.png │ │ ├── operator_list.png │ │ ├── random_image_in_cloudvolume_viewer.png │ │ ├── random_image_in_neuroglancer.png │ │ └── tasks_in_sqs.png │ ├── api.rst │ ├── conf.py │ ├── development.rst │ ├── index.rst │ ├── install.rst │ ├── introduction.rst │ └── tutorial.rst ├── examples ├── downsample │ ├── create_tasks.sh │ └── downsample.sh ├── inference │ ├── 01_download.sh │ ├── batchnorm3d_to_instancenorm3d.py │ ├── inference.sh │ ├── universal_identity.py │ ├── universal_pytorch.py │ └── universal_pznet.py └── mesh │ ├── create_tasks.sh │ ├── manifest.sh │ └── mesh.sh ├── requirements.txt ├── setup.py └── tests ├── __init__.py ├── chunk ├── __init__.py ├── affinity_map │ └── test_affinity_map.py ├── image │ ├── __init__.py │ ├── test_adjust_grey.py │ └── test_image.py ├── segmentation │ ├── __init__.py │ └── test_segmentation.py ├── test_chunk.py └── test_validate.py ├── command_line.sh ├── data ├── levels │ ├── 1 │ │ └── 26774 │ └── normalized.npz └── log │ ├── 0-3_16384-16492_86294-88342_121142-123190.json │ └── 0-3_16384-16492_88342-90390_110902-112950.json ├── flow ├── __init__.py ├── divid_conquer │ ├── __init__.py │ ├── test_inferencer.py │ ├── test_patch_mask.py │ └── visualize_transform.py ├── test_downsample_upload.py ├── test_flow.py ├── test_load_precomputed.py ├── test_load_save.py ├── test_mesh.py ├── test_normalize_section_contrast.py └── test_save_precomputed.py ├── lib ├── aws │ ├── __init__.py │ └── test_sqs_queue.py ├── sp3_bboxes.txt ├── test_cartesian_coordinate.py └── test_synapses.py ├── requirements.txt ├── test_command_line.py └── test_volume.py /.dockerignore: -------------------------------------------------------------------------------- 1 | .cache 2 | .ropeproject 3 | docker 4 | .env 5 | .dat 6 | .json 7 | jwu 8 | dist 9 | build 10 | *.egg-info 11 | __pycache__ 12 | kubernetes 13 | bin/*.sh 14 | bin/*.yml 15 | *.ipynb 16 | -------------------------------------------------------------------------------- /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | # These are supported funding model platforms 2 | 3 | ko_fi: jingpengw 4 | -------------------------------------------------------------------------------- /.github/workflows/python-app.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a single version of Python 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: Python application 5 | 6 | on: 7 | push: 8 | branches: [ master ] 9 | pull_request: 10 | branches: [ master ] 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | 17 | steps: 18 | - uses: actions/checkout@v2 19 | - name: Set up Python 3.8 20 | uses: actions/setup-python@v2 21 | with: 22 | python-version: "3.8" 23 | - name: Install dependencies 24 | run: | 25 | python -m pip install --upgrade pip 26 | pip install flake8 pytest 27 | if [ -f requirements.txt ]; then pip install -r requirements.txt; fi 28 | - name: Lint with flake8 29 | run: | 30 | # stop the build if there are Python syntax errors or undefined names 31 | flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics 32 | # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide 33 | flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics 34 | - name: Test with pytest 35 | run: | 36 | pytest tests 37 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | chunkflow/plugins/chunkflow-plugins 2 | *.code-workspace 3 | 4 | *.pyc 5 | *.chkpt 6 | *__pycache__* 7 | *.swp 8 | *.dat 9 | *.sh 10 | *.so 11 | *secret.json 12 | .coverage 13 | dist/ 14 | bin/ 15 | jwu 16 | *.egg-info/ 17 | .ipynb_checkpoints/ 18 | .vscode/ 19 | *.idea/ 20 | docker/inference/build_docker.sh 21 | docker/inference/pytorch-emvision/ 22 | docker/inference/pytorch-model/ 23 | docs/build/ 24 | *generated/ 25 | build/ 26 | data/ 27 | *.so 28 | tmp/ 29 | docs/source/tutorial/ 30 | docs/source/_build/ 31 | *.DS_Store 32 | *.log 33 | .idea/ 34 | -------------------------------------------------------------------------------- /.pylintrc: -------------------------------------------------------------------------------- 1 | [TYPECHECK] 2 | 3 | # List of members which are set dynamically and missed by pylint inference 4 | # system, and so shouldn't trigger E1101 when accessed. Python regular 5 | # expressions are accepted. 6 | generated-members=numpy.*,torch.* 7 | 8 | [MASTER] 9 | disable= 10 | c0114,c0115,c0116 -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | # .readthedocs.yml 2 | # Read the Docs configuration file 3 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details 4 | 5 | # Required 6 | version: 2 7 | 8 | # Build documentation in the docs/ directory with Sphinx 9 | sphinx: 10 | configuration: docs/source/conf.py 11 | 12 | # Build documentation with MkDocs 13 | #mkdocs: 14 | # configuration: mkdocs.yml 15 | 16 | # Optionally build your docs in additional formats such as PDF and ePub 17 | # formats: all 18 | 19 | # Optionally set the version of Python and requirements required to build your docs 20 | python: 21 | version: 3.8 22 | install: 23 | # - requirements: requirements.txt 24 | - requirements: docs/requirements.txt 25 | system_packages: true 26 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - '3.8' 4 | - '3.9' 5 | - '3.10' 6 | 7 | #jobs: 8 | # allow_failures: 9 | # - python: '3.9' 10 | 11 | install: 12 | - sudo find /usr -name '*.pyc' -delete 13 | - sudo apt update 14 | - pip install . 15 | - python setup.py install 16 | - pip install pytest 17 | - pip install coveralls 18 | - pip install pytest-cov 19 | 20 | script: 21 | # unit tests 22 | - bash tests/command_line.sh 23 | - pytest --cov-append --cov=./chunkflow ./tests --verbose 24 | 25 | after_success: 26 | - coveralls 27 | 28 | deploy: 29 | provider: pypi 30 | skip_existing: true 31 | user: "__token__" 32 | password: 33 | secure: W20x3KuW8ziG7aAif+CwfDoCyoMMsylsnZka/YxHNfIv6FVcYLbY2Xb9PPqta2KE0oVObi8WI25ff7VGqfS3Rk2EX5IYrFz6KXhTzThtkUxP5r6geDm0X28+OCPlCiDH8xIjUoW92duHA1QQyqsekuClVCRsotT0r1wFBFLUVmHHxe7v5987HdEtc4kjCQtlnIkxlVSpe3taGHp5Ju7X9vGc0hDja7w9xRECsj4NHY39FxWVdwLijvKoulMB80CgCIHNHP9Q/fHJEYgVrvf4g9wL5z2HyoDn0vZuyl61ZwAf2XVFJp5zN+6UIU2csnm4C6Y1YOZyV4j7CcZiT78/zEy7nWOiPZcraPUiuIMZbsoLKCXLrrtdDrpa1rCsA1rcnguXJtNAXu9JqbpPsHEx/Ed8TPvPoCOapvbFTDVkdDEPXctvKciakvjd+x9FyJxP1UcFwDHVTvkFsz3++tezj6oPIczO2k1f+0Zjtz+ShpiMClzItKHwTsuoghgBK8Y7chW5YfE59N5a+XOaYenL1r33s6aaLCnnS4eLEHcZb1pK0JYBlPsScN7t/s0Z2Bgox+90O3NA7c9voyXr8nBdRUbNX+Vw3BWZqCYw5N7CNpnlvAga0Y7UUP2AvfwN01U7wEIYGoyA3ZLwGYr5q2UlR3MwPdUUauZ1IYEEf9d2Xl0= 34 | on: 35 | tags: true 36 | all_branches: true 37 | distributions: sdist 38 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Contributor Covenant Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to making participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies both within project spaces and in public spaces 49 | when an individual is representing the project or its community. Examples of 50 | representing a project or community include using an official project e-mail 51 | address, posting via an official social media account, or acting as an appointed 52 | representative at an online or offline event. Representation of a project may be 53 | further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the project team at jingpeng.wu@gmail.com. All 59 | complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 71 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 72 | 73 | [homepage]: https://www.contributor-covenant.org 74 | 75 | For answers to common questions about this code of conduct, see 76 | https://www.contributor-covenant.org/faq 77 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # backend: base | pytorch | pznet 2 | ARG BACKEND=pytorch 3 | 4 | FROM seunglab/chunkflow:${BACKEND} 5 | #FROM seunglab/pznet:latest 6 | 7 | target maintainer = "Jingpeng Wu" \ 8 | email = "jingpeng@princeton.edu" 9 | 10 | RUN mkdir -p $HOME/workspace/chunkflow 11 | 12 | # WORKDIR only works with ENV 13 | 14 | ENV LC_ALL C.UTF-8 15 | ENV LANG C.UTF-8 16 | ENV HOME /root 17 | 18 | WORKDIR $HOME/workspace/chunkflow 19 | COPY . . 20 | 21 | RUN apt-get update && apt-get install -y -qq --no-install-recommends \ 22 | apt-utils \ 23 | wget \ 24 | git \ 25 | build-essential \ 26 | parallel \ 27 | # test whether pip is working 28 | # there is an issue of pip: 29 | # https://github.com/laradock/laradock/issues/1496 30 | # we need this hash to solve this issue 31 | # && ln -sf /usr/bin/pip3 /usr/bin/pip \ 32 | # this do not work due to an issue in pip3 33 | # https://github.com/pypa/pip/issues/5240 34 | && pip install -U pip \ 35 | && hash -r pip \ 36 | && pip install --upgrade setuptools \ 37 | && pip install numpy setuptools cython --no-cache-dir \ 38 | && pip install -U protobuf scipy brotlipy --no-cache-dir \ 39 | # && pip install fpzip --no-binary :all: --no-cache-dir \ 40 | # setup environment variables 41 | # we have to setup first, otherwise click installation will fail 42 | && echo "export LC_ALL=C.UTF-8" >> $HOME/.bashrc \ 43 | && echo "export LANG=C.UTF-8" >> $HOME/.bashrc \ 44 | && echo "export PYTHONPATH=$HOME/workspace/chunkflow:\$PYTHONPATH" >> $HOME/.bashrc \ 45 | && pip install -r requirements.txt --no-cache-dir \ 46 | && pip install -r tests/requirements.txt --no-cache-dir \ 47 | # install the commandline chunkflow 48 | && pip install -e . \ 49 | # cleanup system libraries 50 | && apt-get remove --purge -y \ 51 | build-essential \ 52 | && apt-get clean \ 53 | && apt-get autoremove --purge -y \ 54 | && rm -rf /var/lib/apt/lists/* \ 55 | # the test will not pass due to missing of credentials. 56 | # && pytest tests \ 57 | && chunkflow 58 | 59 | WORKDIR $HOME/workspace/chunkflow/ 60 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include requirements.txt 2 | 3 | exclude tests 4 | exclude .travis.yml 5 | exclude dist 6 | exclude chunkflow.egg-info 7 | exclude docs/source/_static 8 | exclude examples 9 | exclude distributed 10 | exclude docker 11 | -------------------------------------------------------------------------------- /chunkflow/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | # Using Green Threads 5 | # import gevent.monkey 6 | #from gevent import monkey 7 | #monkey.patch_all(thread=True) 8 | 9 | from .__version__ import * 10 | -------------------------------------------------------------------------------- /chunkflow/__main__.py: -------------------------------------------------------------------------------- 1 | from chunkflow.flow.flow import main 2 | 3 | main() 4 | -------------------------------------------------------------------------------- /chunkflow/__version__.py: -------------------------------------------------------------------------------- 1 | version = "1.1.7" 2 | -------------------------------------------------------------------------------- /chunkflow/chunk/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Chunk 2 | # from .utils import load_chunk_or_volume 3 | -------------------------------------------------------------------------------- /chunkflow/chunk/affinity_map/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import AffinityMap -------------------------------------------------------------------------------- /chunkflow/chunk/affinity_map/base.py: -------------------------------------------------------------------------------- 1 | __doc__ = """Image chunk class""" 2 | 3 | import numpy as np 4 | 5 | from chunkflow.lib.cartesian_coordinate import Cartesian 6 | from chunkflow.chunk import Chunk 7 | 8 | class AffinityMap(Chunk): 9 | """ 10 | a chunk of affinity map. It has x,y,z three channels with single precision. 11 | """ 12 | def __init__(self, array: np.ndarray, 13 | voxel_offset: Cartesian=None, 14 | voxel_size: Cartesian=None, 15 | layer_type: str = None): 16 | assert isinstance(array, np.ndarray) 17 | assert array.ndim == 4 18 | assert np.issubdtype(array.dtype, np.float32) 19 | assert array.shape[0] == 3 20 | super().__init__(array, 21 | voxel_offset=voxel_offset, 22 | voxel_size=voxel_size, 23 | layer_type=layer_type) 24 | 25 | @classmethod 26 | def from_chunk(cls, chk: Chunk): 27 | assert isinstance(chk, Chunk) 28 | return cls(chk.array, 29 | voxel_offset = chk.voxel_offset, 30 | voxel_size = chk.voxel_size, 31 | layer_type = chk.layer_type) 32 | 33 | def quantize(self, mode: str='xy'): 34 | """transform affinity map to gray scale image 35 | 36 | Args: 37 | mode (str, optional): tranformation mode. Defaults to 'xy'. 38 | 39 | Raises: 40 | ValueError: only support mode of xy and z. 41 | 42 | Returns: 43 | Chunk: the gray scale image chunk 44 | """ 45 | if mode == 'z': 46 | # only use the last channel, it is the Z affinity 47 | # if this is affinitymap 48 | image = self[-1, :, :, :] 49 | elif mode == 'xy': 50 | image = (self[0,...] + self[1, ...]) / 2. 51 | else: 52 | raise ValueError(f'only support xy and z mode, but got {mode}') 53 | 54 | image = (image * 255.).astype(np.uint8) 55 | image = Chunk(image) 56 | image.set_properties(self.properties) 57 | # assert np.issubdtype(image.dtype, np.uint8) 58 | return image -------------------------------------------------------------------------------- /chunkflow/chunk/image/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import Image 2 | -------------------------------------------------------------------------------- /chunkflow/chunk/image/base.py: -------------------------------------------------------------------------------- 1 | __doc__ = """Image chunk class""" 2 | 3 | from tqdm import tqdm 4 | import numpy as np 5 | 6 | from .adjust_grey import normalize_shang 7 | from chunkflow.chunk import Chunk 8 | from chunkflow.flow.divid_conquer.inferencer import Inferencer 9 | 10 | 11 | class Image(Chunk): 12 | """ 13 | a chunk of image volume. 14 | """ 15 | def __init__(self, array: np.ndarray, **kwargs): 16 | super().__init__(array, **kwargs) 17 | 18 | @classmethod 19 | def from_chunk(cls, chk: Chunk): 20 | return cls(chk.array, **chk.properties) 21 | 22 | def inference(self, inferencer: Inferencer): 23 | """run convolutional net inference for this image chunk""" 24 | return inferencer(self) 25 | 26 | def normalize_shang(self, nominalmin, nominalmax, clipvalues): 27 | return normalize_shang(self.array, nominalmin, nominalmax, 28 | clipvalues) 29 | 30 | def _find_section_clamping_values(self, 31 | hist: np.ndarray, lower_clip_fraction: float, upper_clip_fraction: float): 32 | """compute the clamping values for each section.""" 33 | # remove the np.copy from original code since we only need this once 34 | filtered = hist 35 | 36 | # remove pure black from frequency counts as 37 | # it has no information in our images 38 | filtered[0] = 0 39 | 40 | cdf = np.zeros(shape=(len(filtered), ), dtype=np.uint64) 41 | cdf[0] = filtered[0] 42 | for i in range(1, len(filtered)): 43 | cdf[i] = cdf[i - 1] + filtered[i] 44 | 45 | total = cdf[-1] 46 | 47 | if total == 0: 48 | return (0, 0) 49 | 50 | lower = 0 51 | for i, val in enumerate(cdf): 52 | if float(val) / float(total) > lower_clip_fraction: 53 | break 54 | lower = i 55 | 56 | upper = 0 57 | for i, val in enumerate(cdf): 58 | if float(val) / float(total) > 1 - upper_clip_fraction: 59 | break 60 | upper = i 61 | 62 | return lower, upper 63 | 64 | def _hist_to_lookup_table(self, 65 | hist: np.ndarray, lower_clip_fraction: float, upper_clip_fraction: float, 66 | minval: int = 1, 67 | maxval: int = 255): 68 | """histogram to lookup table 69 | 70 | Args: 71 | hist (np.ndarray): histogram 72 | 73 | Returns: 74 | np.ndarray: lookup table 75 | """ 76 | lower, upper = self._find_section_clamping_values( 77 | hist, lower_clip_fraction, upper_clip_fraction) 78 | 79 | if lower == upper: 80 | #lookup_table = np.arange(0, 256, dtype=np.uint8) 81 | # no need to perform any transform 82 | return None 83 | else: 84 | # compute the lookup table 85 | lookup_table = np.arange(0, 256, dtype=np.float32) 86 | lookup_table = (lookup_table - float(lower)) * ( 87 | maxval / (float(upper) - float(lower))) 88 | np.clip(lookup_table, minval, maxval, out=lookup_table) 89 | lookup_table = np.round(lookup_table) 90 | lookup_table = lookup_table.astype( np.uint8 ) 91 | return lookup_table 92 | 93 | def normalize_contrast(self, 94 | lower_clip_fraction: float = 0.01, 95 | upper_clip_fraction: float = 0.01, 96 | minval: int = 1, 97 | maxval: int = 255, 98 | per_section: bool = True 99 | ): 100 | 101 | def _normalize_array(array: np.ndarray, 102 | lower_clip_fraction: float, 103 | upper_clip_fraction: float, 104 | minval: int = 1, 105 | maxval: int = 255): 106 | hist = np.bincount(array.flatten(), minlength=255) 107 | lookup_table = self._hist_to_lookup_table( 108 | hist, lower_clip_fraction, upper_clip_fraction, 109 | minval=minval, maxval=maxval) 110 | if lookup_table is not None: 111 | array = lookup_table[array] 112 | return array 113 | 114 | 115 | if per_section: 116 | for z in tqdm(range(self.bbox.start.z, self.bbox.stop.z)): 117 | slices = (slice(z, z+1), *self.slices[-2:]) 118 | section = self.cutout(slices) 119 | # section = Image.from_chunk(section) 120 | section.array = _normalize_array( 121 | section.array, 122 | lower_clip_fraction, upper_clip_fraction, 123 | minval=minval, maxval=maxval 124 | ) 125 | self.save(section) 126 | else: 127 | # chunk = Image.from_chunk(chunk) 128 | self.array = _normalize_array( 129 | self.array, 130 | lower_clip_fraction, upper_clip_fraction, 131 | minval=minval, maxval=maxval 132 | ) 133 | -------------------------------------------------------------------------------- /chunkflow/chunk/probability_map.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import numpy as np 4 | from skimage.feature import peak_local_max 5 | 6 | from chunkflow.lib.cartesian_coordinate import Cartesian 7 | from .base import Chunk 8 | 9 | 10 | class ProbabilityMap(Chunk): 11 | def __init__(self, array: np.ndarray, 12 | voxel_offset: Cartesian = None, 13 | voxel_size: Cartesian = None): 14 | super().__init__(array, voxel_offset, voxel_size) 15 | assert np.issubdtype(self.dtype, np.floating) 16 | 17 | @classmethod 18 | def from_chunk(cls, chunk: Chunk): 19 | return cls(chunk.array, chunk.voxel_offset, chunk.voxel_size) 20 | 21 | def detect_points(self, min_distance: int = 1, threshold_rel: float=0.3, 22 | exclude_border: int=True, num_peaks: int = math.inf, 23 | return_confidences: bool = True): 24 | # prob = chunk.array.copy() 25 | prob = self.array 26 | prob -= np.mean(prob) 27 | prob[prob<0.] = 0. 28 | prob /= prob.max() 29 | 30 | # since we have already normalized the map, 31 | # the absolute threshold becomes relative threshold! 32 | coords = peak_local_max( 33 | prob, 34 | min_distance=min_distance, 35 | threshold_abs=threshold_rel, 36 | exclude_border=exclude_border, 37 | num_peaks = num_peaks, 38 | ) 39 | print('number of detected points: ', coords.shape[0]) 40 | if coords.shape[0] > num_peaks: 41 | print(f'only select the first {num_peaks} points.') 42 | coords = coords[:num_peaks, :] 43 | 44 | if len(coords) == 0: 45 | coords = None 46 | if return_confidences: 47 | confidences = None 48 | else: 49 | if return_confidences: 50 | confidences = prob[coords[:, 0], coords[:, 1], coords[:, 2]] 51 | coords +=self.voxel_offset 52 | 53 | if return_confidences: 54 | return coords, confidences 55 | else: 56 | return coords 57 | -------------------------------------------------------------------------------- /chunkflow/chunk/segmentation.py: -------------------------------------------------------------------------------- 1 | __doc__ = """Image chunk class""" 2 | import os 3 | import json 4 | from typing import Union 5 | 6 | import numpy as np 7 | import fastremap 8 | 9 | from cloudfiles import CloudFiles 10 | 11 | # from ...lib.gala import evaluate 12 | from .base import Chunk 13 | from chunkflow.lib.gala import evaluate 14 | from chunkflow.lib.cartesian_coordinate import Cartesian 15 | 16 | 17 | class Segmentation(Chunk): 18 | """ 19 | a chunk of segmentation volume. 20 | """ 21 | def __init__(self, array: np.ndarray, **kwargs ): 22 | super().__init__(array, **kwargs) 23 | assert array.ndim == 3 24 | assert np.issubdtype(array.dtype, np.integer) 25 | 26 | @classmethod 27 | def from_chunk(cls, chunk): 28 | assert isinstance(chunk, Chunk) 29 | return cls( 30 | chunk.array, voxel_offset=chunk.voxel_offset, 31 | voxel_size = chunk.voxel_size) 32 | 33 | def evaluate(self, groundtruth, size_threshold: int=1000): 34 | """ 35 | Parameters: 36 | size_threshold [int]: size threshold for Edit Distance. 37 | Ignore splits or merges smaller than this number of voxels. 38 | """ 39 | if not np.issubdtype(self.dtype, np.uint64): 40 | this = self.astype(np.uint64) 41 | else: 42 | this = self 43 | 44 | if not np.issubdtype(groundtruth.dtype, np.uint64): 45 | groundtruth = groundtruth.astype(np.uint64) 46 | 47 | if isinstance(groundtruth, Chunk): 48 | groundtruth = groundtruth.array 49 | 50 | rand_index = evaluate.rand_index(this.array, groundtruth) 51 | adjusted_rand_index = evaluate.adj_rand_index(this.array, groundtruth) 52 | variation_of_information = evaluate.vi(this.array, groundtruth) 53 | fowlkes_mallows_index = evaluate.fm_index(this.array, groundtruth) 54 | edit_distance = evaluate.edit_distance(this.array, groundtruth, size_threshold=size_threshold) 55 | print(f'rand index: {rand_index: .3f}') 56 | print(f'adjusted rand index: {adjusted_rand_index: .3f}') 57 | print(f'variation of information: {variation_of_information: .3f}') 58 | print(f'edit distance: {edit_distance}') 59 | print(f'Fowlkes Mallows Index: {fowlkes_mallows_index: .3f}') 60 | 61 | ret = {} 62 | ret['rand_index'] = rand_index 63 | ret['adjusted_rand_index'] = adjusted_rand_index 64 | ret['variation_of_information'] = variation_of_information 65 | ret['fowlkes_mallows_index'] = fowlkes_mallows_index 66 | ret['edit_distance'] = edit_distance 67 | return ret 68 | 69 | def remap(self, base_id: int = 0): 70 | """renumber the object ID 71 | 72 | Args: 73 | base_id (int, optional): the maximum object ID in previous chunk. Defaults to 0. 74 | 75 | Returns: 76 | new_base_id (int): the maximum object ID in this chunk as the new base ID. 77 | """ 78 | fastremap.renumber(self.array, 79 | preserve_zero=True, in_place=True) 80 | self.array = self.array.astype(np.uint64) 81 | if base_id > 0: 82 | self.array[self.array>0] += base_id 83 | new_base_id = self.max() 84 | return new_base_id 85 | 86 | def mask_fragments(self, voxel_num_threshold: int): 87 | uniq, counts = fastremap.unique(self.array, return_counts=True) 88 | fragment_ids = uniq[counts <= voxel_num_threshold] 89 | print(f'masking out {len(fragment_ids)} fragments in {len(uniq)} with a percentage of {len(fragment_ids)/len(uniq)}') 90 | self.array = fastremap.mask(self.array, fragment_ids) 91 | 92 | def mask_except(self, selected_obj_ids: Union[str, list, set]): 93 | if selected_obj_ids is None: 94 | print('we have not selected any objects to mask out.') 95 | return 96 | 97 | if isinstance(selected_obj_ids, str) and selected_obj_ids.endswith('.json'): 98 | # assume that ids is a json file in the storage path 99 | json_storage = CloudFiles(os.path.dirname(selected_obj_ids)) 100 | ids_str = json_storage.get(os.path.basename(selected_obj_ids)) 101 | selected_obj_ids = set(json.loads(ids_str)) 102 | assert len(selected_obj_ids) > 0 103 | print(f'number of selected objects: {len(selected_obj_ids)}') 104 | elif isinstance(selected_obj_ids, str): 105 | # a simple string, like "34,45,56,23" 106 | # this is used for small object numbers 107 | selected_obj_ids = set([int(id) for id in selected_obj_ids.split(',')]) 108 | 109 | self.array = fastremap.mask_except(self.array, list(selected_obj_ids)) 110 | 111 | -------------------------------------------------------------------------------- /chunkflow/chunk/validate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from skimage.feature import match_template 4 | 5 | 6 | def validate_by_template_matching(img: np.ndarray): 7 | """ Detect 3d black boxes by template matching. 8 | 1. binarize the image. the voxels inside the black box will be false, and the outside will be true 9 | 2. The template is 7x7x2 with one section true and the other false. 10 | 3. sliding the template through the array, and detect the matching regions. 11 | 4. rotate the template to be 7x2x7 and 2x7x7, do the same detection. 12 | 5. if we can find multiple matchings in all the x,y,z direction, there is probably a black box. 13 | Note that this is always effective. If the black box is large enough to reach both sides, 14 | the detection will fail. 15 | 16 | Parameters 17 | ----------- 18 | img: 19 | 3D image volume. 20 | """ 21 | print("validation by template matching...") 22 | 23 | if np.issubdtype(img.dtype, np.floating): 24 | print( 25 | 'do not support image with floating data type, will skip the validation.' 26 | ) 27 | return True 28 | 29 | img = img.astype(dtype=bool) 30 | 31 | score_threshold = 0.9 32 | num_threshold = 100 33 | evidence_point = 0 34 | 35 | temp = np.zeros((7, 7, 2), dtype=bool) 36 | temp[:, :, 0] = True 37 | result = match_template(img, temp) 38 | if np.count_nonzero(result > score_threshold) > num_threshold: 39 | evidence_point += 1 40 | 41 | temp = np.zeros((7, 7, 2), dtype=bool) 42 | temp[:, :, 1] = True 43 | result = match_template(img, temp) 44 | if np.count_nonzero(result > score_threshold) > num_threshold: 45 | evidence_point += 1 46 | 47 | temp = np.zeros((2, 7, 7), dtype=bool) 48 | temp[0, :, :] = True 49 | result = match_template(img, temp) 50 | if np.count_nonzero(result > score_threshold) > num_threshold: 51 | evidence_point += 1 52 | 53 | temp = np.zeros((2, 7, 7), dtype=bool) 54 | temp[1, :, :] = True 55 | result = match_template(img, temp) 56 | if np.count_nonzero(result > score_threshold) > num_threshold: 57 | evidence_point += 1 58 | 59 | temp = np.zeros((7, 2, 7), dtype=bool) 60 | temp[:, 0, :] = True 61 | result = match_template(img, temp) 62 | if np.count_nonzero(result > score_threshold) > num_threshold: 63 | evidence_point += 1 64 | 65 | temp = np.zeros((7, 2, 7), dtype=bool) 66 | temp[:, 1, :] = True 67 | result = match_template(img, temp) 68 | if np.count_nonzero(result > score_threshold) > num_threshold: 69 | evidence_point += 1 70 | 71 | if evidence_point > 4: 72 | return False 73 | else: 74 | return True 75 | -------------------------------------------------------------------------------- /chunkflow/flow/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/chunkflow/flow/__init__.py -------------------------------------------------------------------------------- /chunkflow/flow/base.py: -------------------------------------------------------------------------------- 1 | class OperatorBase(object): 2 | """Real Operator should inherit from this base class.""" 3 | def __init__(self, name: str = None): 4 | assert isinstance(name, str) 5 | self.name = name 6 | 7 | def __call__(self): 8 | """The processing should happen in this function.""" 9 | pass 10 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/chunkflow/flow/divid_conquer/__init__.py -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/chunkflow/flow/divid_conquer/patch/__init__.py -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/base.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .patch_mask import PatchMask 4 | 5 | 6 | class PatchInferencerBase(object): 7 | """PatchEngine 8 | 9 | the input patch is a 10 | """ 11 | def __init__(self, input_patch_size: tuple, output_patch_size: tuple, 12 | output_patch_overlap: tuple, num_output_channels: int, 13 | dtype: str='float32'): 14 | 15 | if output_patch_size is None: 16 | output_patch_size = input_patch_size 17 | 18 | self.input_patch_size = input_patch_size 19 | self.output_patch_size = output_patch_size 20 | self.output_patch_overlap = output_patch_overlap 21 | self.num_output_channels = num_output_channels 22 | 23 | assert len(output_patch_overlap) == 3 24 | assert len(input_patch_size) == 3 25 | assert len(output_patch_size) == 3 26 | 27 | self.crop_margin = tuple( 28 | (osz-isz)//2 for osz, isz in 29 | zip(input_patch_size, output_patch_size)) 30 | 31 | self.input_patch_overlap = tuple( 32 | (opo + 2 * ocms) for opo, ocms in 33 | zip(output_patch_overlap, self.crop_margin)) 34 | 35 | self.input_patch_stride = tuple(p - o for p, o in 36 | zip(input_patch_size, self.input_patch_overlap)) 37 | self.output_patch_stride = tuple(p - o for p, o in 38 | zip(output_patch_size, self.output_patch_overlap)) 39 | 40 | # prepare patch mask 41 | self.output_patch_mask = PatchMask(output_patch_size, 42 | output_patch_overlap, 43 | dtype=dtype) 44 | # keep a version in cpu for making chunk mask 45 | self.output_patch_mask_numpy = self.output_patch_mask 46 | 47 | def __call__(self, input_patch: np.ndarray) -> np.ndarray: 48 | r"""This method should be inherited for real implementation 49 | 50 | Args: 51 | patch: a image patch with datatype of float32, 52 | The value range should be in [0,1] 53 | 54 | Returns 55 | -------- 56 | np.ndarray 57 | """ 58 | return NotImplementedError('this function should be overload by inherited class!') 59 | 60 | def _reshape_patch_to_5d(self, input_patch): 61 | """patch should be a 5d np array 62 | """ 63 | assert isinstance(input_patch, np.ndarray) 64 | if input_patch.ndim == 3: 65 | input_patch = input_patch.reshape((1, 1) + input_patch.shape) 66 | elif input_patch.ndim == 4: 67 | input_patch = input_patch.reshape((1, ) + input_patch.shape) 68 | return input_patch 69 | 70 | def _crop_output_patch(self, output_patch: np.ndarray): 71 | return output_patch[:, :self.num_output_channels, 72 | self.crop_margin[0]:output_patch.shape[-3]-self.crop_margin[0], 73 | self.crop_margin[1]:output_patch.shape[-2]-self.crop_margin[1], 74 | self.crop_margin[2]:output_patch.shape[-1]-self.crop_margin[2]] 75 | 76 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/identity.py: -------------------------------------------------------------------------------- 1 | import platform 2 | import numpy as np 3 | from .base import PatchInferencerBase 4 | 5 | 6 | class Identity(PatchInferencerBase): 7 | """ 8 | IdentityPatchInferenceEngine(PatchInferenceEngine) 9 | 10 | return the same output with the input 11 | this class was only used for tests 12 | Note that the bump parameter is not used now, only support bump function 13 | """ 14 | def __init__(self, convnet_model: str, convnet_weight_path: str, 15 | input_patch_size: tuple, output_patch_overlap: tuple, 16 | output_patch_size: tuple = None, 17 | num_output_channels: int = 1, 18 | dtype='float32', 19 | bump: str='wu'): 20 | assert bump == 'wu' 21 | super().__init__(input_patch_size, output_patch_size, 22 | output_patch_overlap, num_output_channels, 23 | dtype=dtype) 24 | 25 | @property 26 | def compute_device(self): 27 | return platform.processor() 28 | 29 | def __call__(self, input_patch): 30 | """ 31 | return the same with argument 32 | reshape the size to 5 dimension: 33 | batch, channel, z, y, x 34 | """ 35 | input_patch = self._reshape_patch_to_5d(input_patch) 36 | 37 | output_patch = input_patch.astype(np.float32) 38 | #if np.issubdtype(patch.dtype, np.integer): 39 | # # normalize to 0-1 value range 40 | # output /= np.iinfo(patch.dtype).max 41 | 42 | output_patch = self._crop_output_patch(output_patch) 43 | 44 | # mask should be done in patch engine now 45 | output_patch *= self.output_patch_mask_numpy 46 | 47 | if self.num_output_channels > 1: 48 | output_patch = np.repeat(output_patch, 49 | self.num_output_channels, axis=1) 50 | 51 | return output_patch 52 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/patch_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | from math import log10 3 | import numpy as np 4 | 5 | 6 | class PatchMask(np.ndarray): 7 | def __new__(cls, patch_size, overlap, dtype='float32'): 8 | assert len(patch_size) == 3 9 | assert len(overlap) == 3 10 | 11 | mask = make_patch_mask(patch_size, overlap, dtype=dtype) 12 | return np.asarray(mask).view(cls) 13 | 14 | 15 | def make_patch_mask(patch_size, overlap, dtype='float32'): 16 | """ 17 | _make_mask() 18 | return: 19 | an numpy array with data type of float32. The value was generated 20 | using a bump function. the overlapping borders and corners were 21 | normalized according to weight accumulation. 22 | https://en.wikipedia.org/wiki/Bump_function 23 | """ 24 | bump_map = make_bump_map(patch_size) 25 | stride = tuple(p - o for p, o in zip(patch_size, overlap)) 26 | # use 3x3x3 mask addition to figure out the normalization parameter 27 | # this is a simulation of blending 28 | base_mask = np.zeros(tuple(f + 2 * s 29 | for (f, s) in zip(patch_size, stride)), 30 | dtype='float64') 31 | for nz in range(3): 32 | for ny in range(3): 33 | for nx in range(3): 34 | base_mask[nz*stride[0]:nz*stride[0]+patch_size[0], 35 | ny*stride[1]:ny*stride[1]+patch_size[1], 36 | nx*stride[2]:nx*stride[2]+patch_size[2]] += \ 37 | bump_map 38 | 39 | bump_map /= base_mask[stride[0]:stride[0] + 40 | patch_size[0], stride[1]:stride[1] + 41 | patch_size[1], stride[2]:stride[2] + patch_size[2]] 42 | 43 | np.testing.assert_array_equal(bump_map[ 44 | overlap[0]:-overlap[0], 45 | overlap[1]:-overlap[1], 46 | overlap[2]:-overlap[2]], 1) 47 | 48 | return bump_map.astype(dtype) 49 | 50 | 51 | def make_bump_map(patch_size): 52 | x = range(patch_size[-1]) 53 | y = range(patch_size[-2]) 54 | z = range(patch_size[-3]) 55 | zv, yv, xv = np.meshgrid(z, y, x, indexing='ij') 56 | xv = (xv + 1.0) / (patch_size[-1] + 1.0) * 2.0 - 1.0 57 | yv = (yv + 1.0) / (patch_size[-2] + 1.0) * 2.0 - 1.0 58 | zv = (zv + 1.0) / (patch_size[-3] + 1.0) * 2.0 - 1.0 59 | bump_map = np.exp(-1.0 / (1.0 - xv * xv) + 60 | -1.0 / (1.0 - yv * yv) + 61 | -1.0 / (1.0 - zv * zv)) 62 | 63 | bump_map = np.interp(bump_map, (bump_map.min(), bump_map.max()), (1, 1e6)) 64 | # make the low value a little bit higher to avoid floating point error 65 | #threshold = np.max(bump_map) * 1e-8 66 | #bump_map[bump_map < threshold] = threshold 67 | 68 | return np.asarray(bump_map, dtype=np.float64) 69 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/pytorch.py: -------------------------------------------------------------------------------- 1 | # from .inference_engine import InferenceEngine 2 | # import imp 3 | import torch 4 | from .base import PatchInferencerBase 5 | from chunkflow.lib import load_source 6 | 7 | torch.backends.cudnn.benchmark = True 8 | 9 | 10 | class PyTorch(PatchInferencerBase): 11 | """perform inference for an image patch using pytorch. 12 | Parameters 13 | ---------- 14 | patch_size: size of input/output patch size. We assume that 15 | the input and output patch size is the same. 16 | patch_overlap: overlap of neighboring patches. 17 | model_file_name: file name of model 18 | weight_file_name: file name of trained weight. 19 | num_output_channels: number of output channels. 20 | 21 | You can make some customized processing in your model file. 22 | You can define `load_model` function to customize your way of 23 | loading model. This is useful for loading some models trained using 24 | old version pytorch (<=0.4.0). You can also define `pre_process` 25 | and `post_process` function to insert your own customized processing. 26 | """ 27 | def __init__(self, convnet_model: str, convnet_weight_path: str, 28 | input_patch_size: tuple, 29 | output_patch_size: tuple, 30 | output_patch_overlap: tuple, 31 | num_output_channels: int = 1, 32 | dtype: str='float32', 33 | bump: str='wu'): 34 | # To-Do: support zung function 35 | assert bump == 'wu' 36 | super().__init__(input_patch_size, output_patch_size, 37 | output_patch_overlap, num_output_channels, 38 | dtype=dtype) 39 | 40 | self.num_output_channels = num_output_channels 41 | if torch.cuda.is_available(): 42 | self.is_gpu = True 43 | # put mask to gpu 44 | self.output_patch_mask = torch.from_numpy(self.output_patch_mask).cuda() 45 | else: 46 | self.is_gpu = False 47 | 48 | net_source = load_source(convnet_model) 49 | 50 | if hasattr(net_source, "load_model"): 51 | self.model = net_source.load_model(convnet_weight_path) 52 | else: 53 | self.model = net_source.InstantiatedModel 54 | if self.is_gpu: 55 | map_location = 'cuda' 56 | else: 57 | map_location = 'cpu' 58 | chkpt = torch.load(convnet_weight_path, map_location=map_location) 59 | state_dict = chkpt['state_dict'] if 'state_dict' in chkpt else chkpt 60 | self.model.load_state_dict(state_dict) 61 | 62 | if self.is_gpu and next(self.model.parameters()).is_cuda: 63 | self.model.cuda() 64 | 65 | # data parallel do not work with old emvision net 66 | #self.model = torch.nn.DataParallel( 67 | # self.model, device_ids=range(torch.cuda.device_count())) 68 | 69 | # Print model's state_dict 70 | #print("Model's state_dict:") 71 | #for param_tensor in self.model.state_dict(): 72 | # print(param_tensor, "\t", self.model.state_dict()[param_tensor].size()) 73 | 74 | 75 | if hasattr(net_source, "pre_process"): 76 | self.pre_process = net_source.pre_process 77 | else: 78 | self.pre_process = self._pre_process 79 | 80 | if hasattr(net_source, "post_process"): 81 | self.post_process = net_source.post_process 82 | else: 83 | self.post_process = self._identity 84 | 85 | @property 86 | def compute_device(self): 87 | return torch.cuda.get_device_name(0) 88 | 89 | def _pre_process(self, input_patch): 90 | input_patch = torch.from_numpy(input_patch) 91 | if self.is_gpu: 92 | input_patch = input_patch.cuda() 93 | return input_patch 94 | 95 | def _identity(self, patch): 96 | return patch 97 | 98 | def __call__(self, input_patch): 99 | # make sure that the patch is 5d ndarray 100 | input_patch = self._reshape_patch_to_5d(input_patch) 101 | 102 | with torch.no_grad(): 103 | net_input = self.pre_process(input_patch) 104 | # the network input and output should be dict 105 | net_output = self.model(net_input) 106 | 107 | # get the required output patch from network 108 | # The processing depends on network model and application 109 | output_patch = self.post_process(net_output) 110 | 111 | # mask in gpu/cpu 112 | output_patch = self._crop_output_patch(output_patch) 113 | output_patch *= self.output_patch_mask 114 | 115 | if self.is_gpu: 116 | # transfer to cpu 117 | output_patch = output_patch.data.cpu() 118 | output_patch = output_patch.numpy() 119 | return output_patch 120 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/pznet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import sys 3 | import platform 4 | from .base import PatchInferencerBase 5 | 6 | 7 | class PZNet(PatchInferencerBase): 8 | def __init__(self, net_dir): 9 | super().__init__(dtype='float32') 10 | sys.path.append(net_dir) 11 | import pznet 12 | # self.net = pznet.znet(model_file_name, net_file_name) 13 | self.net = pznet.znet() 14 | self.net.load_net(net_dir) 15 | 16 | @property 17 | def compute_device(self): 18 | return platform.processor() 19 | 20 | def __call__(self, patch): 21 | """ 22 | args: 23 | patch (5d numpy array): input patch with dimensions \ 24 | batch/channel/z/y/x 25 | return: 26 | 5d numpy array with the same dimension arrangment. 27 | """ 28 | # make sure that the input patch is 5d ndarray 29 | patch = self._reshape_patch(patch) 30 | return self.net.forward(patch) 31 | 32 | 33 | if __name__ == "__main__": 34 | import numpy as np 35 | for i in range(10): 36 | print("patch iteration: {}".format(i)) 37 | engine = PZNet("/nets/pinky100/unet4-long/mip1/cores2") 38 | input_patch = np.random.randint(0, 39 | 255, (1, 1, 20, 256, 256), 40 | dtype='uint8') 41 | engine(input_patch) 42 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/patch/universal.py: -------------------------------------------------------------------------------- 1 | import platform 2 | 3 | import numpy as np 4 | from .base import PatchInferencerBase 5 | from chunkflow.lib import load_source 6 | 7 | 8 | class Universal(PatchInferencerBase): 9 | """perform inference for an image patch using an arbitrary backend. 10 | Parameters 11 | ---------- 12 | patch_size: size of input/output patch size. We assume that 13 | the input and output patch size is the same. 14 | patch_overlap: overlap of neighboring patches. 15 | model_file_name: file name of model 16 | weight_file_name: file name of trained weight. 17 | num_output_channels: number of output channels. 18 | 19 | You can make some customized processing in your model file. 20 | You need to define a class called `PatchInferencer`. 21 | The constructor inputs are `model_weight_file`, `patch_mask`. 22 | You should define a `__call__` function to process the `input_patch`, 23 | and the output the `output_patch` after masking. 24 | 25 | Because the masking can happen in the device, such as GPU or TPU, so 26 | it is the PatchInferencer's job to perform the masking. 27 | """ 28 | def __init__(self, convnet_model: str, convnet_weight_path: str, 29 | input_patch_size: tuple, 30 | output_patch_size: tuple, 31 | output_patch_overlap: tuple, 32 | num_output_channels: int = 1, 33 | dtype: str='float32', 34 | bump: str='wu'): 35 | # To-Do: support zung function 36 | assert bump == 'wu' 37 | super().__init__(input_patch_size, output_patch_size, 38 | output_patch_overlap, num_output_channels, 39 | dtype=dtype) 40 | 41 | self.num_output_channels = num_output_channels 42 | 43 | net_source = load_source(convnet_model) 44 | 45 | assert hasattr(net_source, "PatchInferencer") 46 | self.patch_inferencer = net_source.PatchInferencer( 47 | convnet_weight_path, 48 | self.output_patch_mask,) 49 | 50 | # this feature is not working correctly for now. 51 | #crop_output_patch_margin=self.crop_margin) 52 | 53 | @property 54 | def compute_device(self): 55 | # To-Do: make the device available to log 56 | if hasattr(self.patch_inferencer, 'compute_device'): 57 | return self.patch_inferencer.compute_device 58 | else: 59 | return platform.processor() 60 | 61 | def __call__(self, input_patch): 62 | # make sure that the patch is 5d ndarray 63 | input_patch = self._reshape_patch_to_5d(input_patch) 64 | output_patch = self.patch_inferencer( input_patch ) 65 | #output_patch = self._crop_output_patch(output_patch) 66 | #output_patch = output_patch * self.output_patch_mask 67 | #output_patch = output_patch.cpu().numpy() 68 | assert isinstance(output_patch, np.ndarray) 69 | return output_patch 70 | -------------------------------------------------------------------------------- /chunkflow/flow/divid_conquer/transform.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from itertools import product 3 | from typing import List 4 | 5 | import numpy as np 6 | 7 | 8 | class TransformBase(ABC): 9 | def __init__(self) -> None: 10 | super().__init__() 11 | 12 | @abstractmethod 13 | def forward(self, arr: np.ndarray): 14 | ... 15 | 16 | @abstractmethod 17 | def backward(self, arr: np.ndarray): 18 | pass 19 | 20 | class Lazy(TransformBase): 21 | def __init__(self) -> None: 22 | super().__init__() 23 | 24 | def forward(self, arr: np.ndarray): 25 | return arr 26 | 27 | def backward(self, arr: np.ndarray): 28 | return arr 29 | 30 | class FlipLR(TransformBase): 31 | def __init__(self) -> None: 32 | super().__init__() 33 | 34 | def _flip(self, arr: np.ndarray): 35 | for z in range(arr.shape[-3]): 36 | arr[...,z,:,:] = np.fliplr(arr[...,z,:,:]) 37 | return arr 38 | 39 | def forward(self, arr: np.ndarray): 40 | return self._flip(arr) 41 | 42 | def backward(self, arr: np.ndarray): 43 | return self._flip(arr) 44 | 45 | 46 | class FlipUD(TransformBase): 47 | def __init__(self) -> None: 48 | super().__init__() 49 | 50 | def _flip(self, arr: np.ndarray): 51 | for z in range(arr.shape[-3]): 52 | arr[...,z,:,:] = np.flipud(arr[...,z,:,:]) 53 | return arr 54 | 55 | def forward(self, arr: np.ndarray): 56 | return self._flip(arr) 57 | 58 | def backward(self, arr: np.ndarray): 59 | return self._flip(arr) 60 | 61 | class FlipZ(TransformBase): 62 | def __init__(self) -> None: 63 | super().__init__() 64 | 65 | def _flip(self, arr: np.ndarray): 66 | for x in range(arr.shape[-1]): 67 | arr[..., :, :, x] = np.flipud(arr[..., :, :, x]) 68 | return arr 69 | 70 | def forward(self, arr: np.ndarray): 71 | return self._flip(arr) 72 | 73 | def backward(self, arr: np.ndarray): 74 | return self._flip(arr) 75 | 76 | # class Rotate(TransformBase): 77 | # def __init__(self, k: int = 1) -> None: 78 | # super().__init__() 79 | # self.k = k 80 | 81 | # def _rotate(self, arr: np.ndarray, k: int): 82 | # for z in range(arr.shape[-3]): 83 | # arr[...,z,:,:] = np.rot90( 84 | # arr[...,z,:,:], 85 | # k = k, 86 | # axes=(-1, -2) 87 | # ) 88 | # return arr 89 | 90 | # def forward(self, arr: np.ndarray): 91 | # return self._rotate(arr, self.k) 92 | 93 | # def backward(self, arr: np.ndarray): 94 | # return self._rotate(arr, -self.k) 95 | 96 | class Transpose(TransformBase): 97 | def __init__(self) -> None: 98 | super().__init__() 99 | 100 | def _transpose(self, arr: np.ndarray): 101 | arr = np.swapaxes(arr, -1, -2) 102 | #for z in range(arr.shape[-3]): 103 | # #arr[..., z, :, :] = np.transpose(arr[..., z, :, :]) 104 | # arr = np.swapaxes(arr, -1, -2) 105 | return arr 106 | 107 | def forward(self, arr): 108 | return self._transpose(arr) 109 | 110 | def backward(self, arr: np.ndarray): 111 | return self._transpose(arr) 112 | 113 | 114 | class TransformSequences: 115 | def __init__(self, 116 | transpose: bool = True, 117 | fliplr: bool = True, 118 | flipud: bool = True, 119 | flipz: bool = False, 120 | ) -> None: 121 | 122 | options = [] 123 | if transpose: 124 | options.append((Lazy(), Transpose())) 125 | if fliplr: 126 | options.append((Lazy(), FlipLR())) 127 | if flipud: 128 | options.append((Lazy(), FlipUD())) 129 | if flipz: 130 | options.append((Lazy(), FlipZ())) 131 | 132 | assert len(options) > 0 133 | self.transform_sequences = [x for x in product(*options)] 134 | assert len(self.transform_sequences) == 8 135 | print(f'get {len(self.transform_sequences)} transformation sequences.') 136 | 137 | def forward(self, arr: np.ndarray): 138 | transformed_arrays = [] 139 | for transform_sequence in self.transform_sequences: 140 | tmp = np.copy(arr) 141 | for transform in transform_sequence: 142 | tmp = transform.forward(tmp) 143 | transformed_arrays.append(tmp) 144 | 145 | return transformed_arrays 146 | 147 | def backward(self, transformed_arrays: List[np.ndarray]): 148 | assert len(transformed_arrays) == len(self.transform_sequences) 149 | inversed_arrays = [] 150 | for idx, transform_sequence in enumerate(self.transform_sequences): 151 | arr = np.copy(transformed_arrays[idx]) 152 | for transform in transform_sequence: 153 | arr = transform.backward(arr) 154 | inversed_arrays.append(arr) 155 | 156 | return inversed_arrays 157 | 158 | 159 | if __name__ == '__main__': 160 | img = np.zeros(shape=(4, 128, 128), dtype=np.uint8) 161 | img[:, 60:68, :] = 255 162 | 163 | transform_sequences = TransformSequences() 164 | 165 | transformed_images = transform_sequences.forward(img) 166 | inversed_images = transform_sequences.backward(transformed_images) 167 | 168 | import matplotlib.pyplot as plt 169 | 170 | fig, axs = plt.subplots(2, 2, subplot_kw=dict(projection="polar")) 171 | axs[0, 0].imshow(img[0,...]) 172 | 173 | fig.show() 174 | 175 | breakpoint() 176 | -------------------------------------------------------------------------------- /chunkflow/flow/downsample_upload.py: -------------------------------------------------------------------------------- 1 | 2 | from cloudvolume import CloudVolume 3 | import tinybrain 4 | import numpy as np 5 | from cloudvolume.lib import Bbox 6 | from .base import OperatorBase 7 | 8 | 9 | class DownsampleUploadOperator(OperatorBase): 10 | """ 11 | Multiple mip level downsampling including image and segmenation. 12 | 13 | For image, the algorithm will be automatically choosen as average pooling. 14 | For segmentation, the algorithm will be Will Silversman's countless algorithm to perform model pooling. The most frequent segmentation ID will be choosen. 15 | 16 | The type of chunk was automatically determined from the data type. 17 | Image: uint8, floating 18 | Segmentation: uint16, uint32, uint64,... 19 | """ 20 | def __init__(self, 21 | volume_path: str, 22 | factor: tuple = (2,2,2), 23 | chunk_mip: int = 0, 24 | start_mip: int = None, 25 | stop_mip: int = 5, 26 | fill_missing: bool = True, 27 | verbose=False, 28 | name='downsample-upload'): 29 | """ 30 | volume_path: (str) path of volume 31 | factor: (tuple) downsampling factor in z,y,x order 32 | chunk_mip: (int) the mip level of input chunk 33 | start_mip: (int) the mip level for starting uploading 34 | stop_mip: (int) the mip level for stoping uploading. Note that the indexing follows python indexing, this stop mip will not be included. For example, if you would like to upload mip level 1 to 4, the start mip will be 1, and the stop mip should be 5. 35 | fill_missing: (bool) fill missing blocks with zeros or not. See same parameter in cloudvolume. 36 | """ 37 | super().__init__(name=name) 38 | 39 | if start_mip is None: 40 | start_mip = chunk_mip + 1 41 | 42 | vols = dict() 43 | for mip in range(start_mip, stop_mip): 44 | vols[mip] = CloudVolume(volume_path, 45 | fill_missing=fill_missing, 46 | bounded=False, 47 | autocrop=True, 48 | mip=mip, 49 | green_threads=True, 50 | delete_black_uploads=True, 51 | progress=verbose) 52 | 53 | self.vols = vols 54 | self.factor = factor 55 | self.chunk_mip = chunk_mip 56 | self.start_mip = start_mip 57 | self.stop_mip = stop_mip 58 | assert len(factor) == 3 59 | assert np.all([f>=1 for f in factor]) 60 | assert stop_mip > start_mip 61 | assert start_mip > chunk_mip 62 | 63 | def __call__(self, chunk): 64 | voxel_offset = chunk.voxel_offset 65 | num_mips = self.stop_mip - self.chunk_mip 66 | # tinybrain use F order and require 4D array! 67 | chunk2 = np.transpose(chunk) 68 | # chunk2 = np.reshape(chunk2, (*chunk2.shape, 1)) 69 | if chunk2.ndim == 3: 70 | chunk2 = np.expand_dims(chunk2, 3) 71 | 72 | if np.issubdtype(chunk.dtype, np.floating) or chunk.dtype == np.uint8: 73 | pyramid = tinybrain.downsample_with_averaging(chunk2, 74 | factor=self.factor[::-1], 75 | num_mips=num_mips) 76 | else: 77 | pyramid = tinybrain.downsample_segmentation(chunk2, 78 | factor=self.factor[::-1], 79 | num_mips=num_mips) 80 | 81 | #print(f'upload image pyramid...') 82 | for mip in range(self.start_mip, self.stop_mip): 83 | # the first chunk in pyramid is already downsampled! 84 | downsampled_chunk = pyramid[mip - self.chunk_mip - 1] 85 | # compute new offset, only downsample the y,x dimensions 86 | offset = np.divide(voxel_offset, np.asarray([ 87 | self.factor[0]**(mip - self.chunk_mip), 88 | self.factor[1]**(mip - self.chunk_mip), 89 | self.factor[2]**(mip - self.chunk_mip)])) 90 | bbox = Bbox.from_delta(offset, downsampled_chunk.shape[0:3][::-1]) 91 | # upload downsampled chunk, note that we should use F order in the indexing 92 | self.vols[mip][bbox.to_slices()[::-1]] = downsampled_chunk 93 | -------------------------------------------------------------------------------- /chunkflow/flow/load_pngs.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | 4 | import numpy as np 5 | 6 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 7 | from chunkflow.chunk import Chunk 8 | 9 | from tqdm import tqdm 10 | import pyspng 11 | 12 | 13 | def load_png_image(file_name: str): 14 | with open(file_name, "rb") as f: 15 | arr = pyspng.load(f.read()) 16 | if np.ndim(arr) == 3: 17 | arr = arr[:, :, 0] 18 | return arr 19 | 20 | 21 | def load_png_images( 22 | path_prefix: str, 23 | bbox: BoundingBox = None, 24 | voxel_offset: Cartesian = Cartesian(0, 0, 0), 25 | voxel_size: Cartesian = Cartesian(1, 1, 1), 26 | digit_num: int = 5, 27 | dtype: np.dtype = None): 28 | if isinstance(dtype, str): 29 | dtype = np.dtype(dtype) 30 | 31 | file_names = [] 32 | 33 | if bbox is None: 34 | if os.path.isfile(path_prefix): 35 | file_names.append(path_prefix) 36 | else: 37 | if os.path.isdir(path_prefix): 38 | dir_path = path_prefix 39 | else: 40 | dir_path = os.path.dirname(path_prefix) 41 | fname = os.path.expanduser(dir_path) 42 | for fname in sorted(os.listdir(dir_path)): 43 | if fname.endswith('.png'): 44 | fname = os.path.join(dir_path, fname) 45 | file_names.append(fname) 46 | arr = load_png_image(file_names[0]) 47 | shape = Cartesian(len(file_names), arr.shape[0], arr.shape[1]) 48 | if dtype is None: 49 | dtype = arr.dtype 50 | bbox = BoundingBox.from_delta(voxel_offset, shape) 51 | else: 52 | for z in range(bbox.start[0], bbox.stop[0]): 53 | file_name = f'{path_prefix}{z:0>{digit_num}d}.png' 54 | file_name = os.path.expanduser(file_name) 55 | file_names.append(file_name) 56 | 57 | chunk = Chunk.from_bbox( 58 | bbox, dtype=dtype, 59 | pattern='zero', 60 | voxel_size=voxel_size 61 | ) 62 | 63 | for z_offset, file_name in tqdm(enumerate(file_names)): 64 | if os.path.exists(file_name): 65 | if z_offset > 0: 66 | arr = load_png_image(file_name) 67 | 68 | if arr.dtype != dtype: 69 | arr = arr.astype(dtype) 70 | # breakpoint() 71 | chunk.array[z_offset, :, :] = arr[ 72 | bbox.start[1]:bbox.stop[1], 73 | bbox.start[2]:bbox.stop[2]] 74 | else: 75 | print(f'image file do not exist: {file_name}') 76 | 77 | return chunk -------------------------------------------------------------------------------- /chunkflow/flow/log_summary.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # coding: utf-8 3 | # Usage: 4 | # download log folder in cloud storage 5 | # run this script to compute summary 6 | 7 | import os 8 | import json 9 | import numpy as np 10 | import pandas as pd 11 | from tqdm import tqdm 12 | 13 | pd.set_option('display.precision', 0) 14 | 15 | 16 | def load_log(log_dir: str) -> pd.DataFrame: 17 | """load log from executed chunkflow tasks 18 | 19 | Args: 20 | log_dir (str): the directory path of the log files 21 | 22 | Returns: 23 | dataframe: the collected items of log 24 | """ 25 | 26 | collected = dict() 27 | 28 | for file_name in tqdm(os.listdir(log_dir), desc='loading log files'): 29 | complete_file_name = os.path.join(log_dir, file_name) 30 | # print('file name: {}'.format(complete_file_name)) 31 | with open(complete_file_name) as f: 32 | d = json.load(f) 33 | 34 | timing = d['timer'] 35 | 36 | if not collected: 37 | # this is the first time, construct the keys 38 | collected['compute_device'] = [] 39 | collected['complete_task'] = [] 40 | for k in timing.keys(): 41 | collected[k] = [] 42 | 43 | time_per_task = 0.0 44 | for t in d['timer'].values(): 45 | time_per_task += t 46 | collected['complete_task'].append( time_per_task ) 47 | 48 | collected['compute_device'].append( d['compute_device'] ) 49 | 50 | for k, v in timing.items(): 51 | collected[k].append( v ) 52 | 53 | df = pd.DataFrame.from_dict(collected, orient='columns') 54 | return df 55 | 56 | 57 | def print_log_statistics(df: pd.DataFrame, output_size: tuple=None): 58 | grouped_df = df.groupby('compute_device') 59 | 60 | print('\n\nmean time (sec):') 61 | print(grouped_df.mean()) 62 | 63 | print('\n\nmax time (sec):') 64 | print(grouped_df.max()) 65 | 66 | print('\n\nmin time (sec):') 67 | print(grouped_df.min()) 68 | 69 | if output_size: 70 | print('\n\nmean speed (mv/s)') 71 | print(np.prod(output_size) / grouped_df.mean() / 1e3) 72 | 73 | print('\n\nsummation of time (hour):') 74 | print(grouped_df.sum() / 3600) 75 | -------------------------------------------------------------------------------- /chunkflow/flow/mask.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | from cloudvolume import CloudVolume 5 | 6 | from chunkflow.chunk import Chunk 7 | from chunkflow.lib.cartesian_coordinate import Cartesian 8 | from .base import OperatorBase 9 | 10 | 11 | class MaskOperator(OperatorBase): 12 | def __init__(self, 13 | volume_path: str, 14 | mask_mip: int, 15 | chunk_mip: int, 16 | inverse: bool = False, 17 | fill_missing: bool = True, 18 | name: str = 'mask'): 19 | super().__init__(name=name) 20 | 21 | self.mask_mip = mask_mip 22 | self.chunk_mip = chunk_mip 23 | self.inverse = inverse 24 | 25 | self.mask_vol = CloudVolume(volume_path, 26 | bounded=False, 27 | fill_missing=fill_missing, 28 | progress=False, 29 | parallel=1, 30 | mip=mask_mip) 31 | 32 | print(f'build mask operator based on {volume_path} at mip {mask_mip}') 33 | 34 | def __call__(self, chunks: list): 35 | """ Make part of chunk to be black according to a mask chunk. 36 | Note that the operation is inplace and the data in the input chunk is changed. 37 | """ 38 | if isinstance(chunks, Chunk): 39 | chunks = [chunks] 40 | elif len(chunks)>1: 41 | for chunk in chunks: 42 | assert isinstance(chunk, Chunk) 43 | assert chunk.voxel_offset == chunks[0].voxel_offset 44 | assert chunk.voxel_size == chunks[0].voxel_size 45 | 46 | assert isinstance(chunks, list) 47 | assert len(chunks)>0 48 | 49 | voxel_size = chunks[0].voxel_size 50 | 51 | mask_voxel_size = Cartesian.from_collection( 52 | self.mask_vol.resolution[::-1] 53 | ) 54 | factor = mask_voxel_size // voxel_size 55 | # factor = tuple(m//c for m, c in zip(self.mask_vol.resolution[::-1], chunk.voxel_size)) 56 | for m, c in zip(mask_voxel_size, voxel_size): 57 | assert m >= c 58 | assert m % c == 0 59 | 60 | mask_in_high_mip = self._read_mask_in_high_mip(chunks[0].bbox, factor) 61 | 62 | if np.all(mask_in_high_mip == 0): 63 | print('the mask is all black, mask all the voxels directly') 64 | for chunk in chunks: 65 | np.multiply(chunk, 0, out=chunk) 66 | return chunks 67 | if np.all(mask_in_high_mip): 68 | print("mask elements are all positive, return directly") 69 | return chunks 70 | 71 | assert np.any(mask_in_high_mip) 72 | 73 | 74 | for chunk in chunks: 75 | # make it the same type with input 76 | mask_in_high_mip = mask_in_high_mip.astype(chunk.dtype) 77 | for offset in np.ndindex(factor): 78 | chunk.array[..., 79 | np.s_[offset[0]::factor[0]], 80 | np.s_[offset[1]::factor[1]], 81 | np.s_[offset[2]::factor[2]]] *= mask_in_high_mip 82 | 83 | return chunks 84 | 85 | 86 | def _read_mask_in_high_mip(self, chunk_bbox, factor): 87 | """ 88 | chunk_bbox: the bounding box of the chunk in lower mip level 89 | """ 90 | # print("download mask chunk...") 91 | # make sure that the slices only contains zyx without channel 92 | chunk_slices = chunk_bbox.slices[-3:] 93 | 94 | # only scale the indices in XY plane 95 | # only scale the indices in XY plane 96 | mask_slices = tuple( 97 | slice(a.start // f, a.stop // f) 98 | for a, f in zip(chunk_slices[-3:], factor)) 99 | 100 | # the slices did not contain the channel dimension 101 | mask = self.mask_vol[mask_slices[::-1]] 102 | # this is a cloudvolume VolumeCutout rather than a normal numpy array 103 | # which will make np.all(mask_in_high_mip == 0) to be 104 | # VolumeCutout(False) rather than False, so we need to transform it 105 | # to numpy 106 | mask = mask.astype(bool) 107 | mask = np.asarray(mask) 108 | mask = np.transpose(mask) 109 | mask = np.squeeze(mask, axis=0) 110 | 111 | if self.inverse: 112 | mask = (mask == 0) 113 | return mask 114 | -------------------------------------------------------------------------------- /chunkflow/flow/mesh_manifest.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import re 4 | from collections import defaultdict 5 | from time import sleep 6 | from typing import Union 7 | 8 | from tqdm import tqdm 9 | 10 | from cloudvolume import CloudVolume 11 | from cloudfiles import CloudFiles 12 | 13 | 14 | from .base import OperatorBase 15 | 16 | 17 | class MeshManifestOperator(OperatorBase): 18 | """Create mesh manifest files for Neuroglancer visualization.""" 19 | def __init__(self, 20 | volume_path: str, 21 | lod: int = 0, 22 | name: str = 'mesh-manifest'): 23 | """ 24 | Parameters 25 | ------------ 26 | volume_path: 27 | path to store mesh manifest files 28 | lod: 29 | level of detail. we always use 0! 30 | """ 31 | super().__init__(name=name) 32 | self.lod = lod 33 | vol = CloudVolume(volume_path) 34 | info = vol.info 35 | assert 'mesh' in info 36 | self.mesh_path = os.path.join(volume_path, info['mesh']) 37 | self.storage = CloudFiles(self.mesh_path) 38 | 39 | def __call__(self, prefix: Union[int, str], digits: int) -> None: 40 | assert int(prefix) < 10**digits 41 | prefix = str(prefix).zfill(digits) 42 | print(f'manifesting meshes with prefix: {prefix}') 43 | 44 | id2filenames = defaultdict(list) 45 | for filename in tqdm( 46 | self.storage.list(prefix=prefix), 47 | desc='list mesh files' 48 | ): 49 | 50 | filename = os.path.basename(filename) 51 | # `match` implies the beginning (^). `search` matches whole string 52 | matches = re.search(r'(\d+):(\d+):', filename) 53 | 54 | if not matches: 55 | continue 56 | 57 | seg_id, lod = matches.groups() 58 | seg_id, lod = int(seg_id), int(lod) 59 | # currently we are not using `level of detail`, it is always 0 60 | # will need to adjust code if we start using variants 61 | assert lod == self.lod 62 | id2filenames[seg_id].append(filename) 63 | 64 | for seg_id, frags in tqdm( 65 | id2filenames.items(), 66 | desc='upload aggregated manifest file' 67 | ): 68 | 69 | # print(f'segment id: {seg_id}') 70 | # print(f'fragments: {frags}') 71 | self.storage.put_json( 72 | path=f'{seg_id}:{self.lod}', 73 | content={"fragments": frags}, 74 | ) 75 | # the last few hundred files will not be uploaded without sleeping! 76 | sleep(0.01) 77 | -------------------------------------------------------------------------------- /chunkflow/flow/napari.py: -------------------------------------------------------------------------------- 1 | try: 2 | import napari 3 | except: 4 | print('napari not installed and related operators could not be used. This does not impact other operators.') 5 | 6 | import numpy as np 7 | from .base import OperatorBase 8 | 9 | 10 | class NapariOperator(OperatorBase): 11 | def __init__(self, 12 | name: str = 'neuroglancer', 13 | voxel_size: tuple = None): 14 | super().__init__(name=name) 15 | self.voxel_size = voxel_size 16 | 17 | def __call__(self, datas: dict, selected: str=None): 18 | if selected is None: 19 | selected = datas.keys() 20 | elif isinstance(selected, str): 21 | selected = selected.split(',') 22 | 23 | for name in selected: 24 | data = datas[name] 25 | 26 | if np.issubdtype(data.dtype, np.uint8): 27 | viewer = napari.view_image(data.array, rgb=False) 28 | napari.run() # start the event loop and show viewer -------------------------------------------------------------------------------- /chunkflow/flow/plugin.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os 4 | import os.path as path 5 | from typing import Union 6 | 7 | import numpy as np 8 | 9 | from chunkflow.lib.cartesian_coordinate import Cartesian 10 | from chunkflow.lib.utils import str_to_dict 11 | 12 | from .base import OperatorBase 13 | 14 | from chunkflow.lib import load_source 15 | from chunkflow.chunk import Chunk 16 | from chunkflow.point_cloud import PointCloud 17 | 18 | 19 | def array_to_chunk(arr: Union[np.ndarray, Chunk], voxel_offset: Cartesian, 20 | voxel_size: Cartesian, shape: tuple): 21 | if isinstance(arr, np.ndarray): 22 | # in case the plugin did some symmetric cropping 23 | offset = tuple(vo + (ins - outs)//2 for vo, ins, outs in zip(voxel_offset, shape[-3:], arr.shape[-3:]) ) 24 | return Chunk(arr, voxel_offset=offset, voxel_size=voxel_size) 25 | else: 26 | return arr 27 | 28 | 29 | class Plugin(OperatorBase): 30 | r""" 31 | Chunk operation using a custom python file. 32 | """ 33 | def __init__(self, plugin_file_name: str, 34 | name: str = 'plugin-1'): 35 | r""" 36 | Loads a custom python file specified in `opprogram`, which 37 | should contain a callable named "exec" such that 38 | a call of `op_call(chunk, args)` operates on the chunk. 39 | """ 40 | super().__init__(name=name) 41 | 42 | if not plugin_file_name.endswith('.py'): 43 | plugin_file_name += '.py' 44 | 45 | plugin_dir = path.join(path.dirname(path.realpath(__file__)), '../plugins') 46 | plugin_dir1 =path.join(plugin_dir, 'chunkflow-plugins/chunkflowplugins') 47 | 48 | plugin_dirs = ['./', plugin_dir, plugin_dir1] 49 | if 'CHUNKFLOW_PLUGIN_DIR' in os.environ: 50 | plugin_dirs.append(os.environ['CHUNKFLOW_PLUGIN_DIR']) 51 | 52 | for plugin_dir in plugin_dirs: 53 | fname = path.join(plugin_dir, plugin_file_name) 54 | if path.exists(fname): 55 | print(f'loading plugin {fname}') 56 | program = load_source(fname) 57 | # assuming this is a func / static functor for now, maybe make it a class? 58 | self.execute = program.execute 59 | break 60 | 61 | assert os.path.exists(fname), f'did not find plugin: {fname}' 62 | assert hasattr(self, 'execute') 63 | 64 | def __call__(self, inputs: list, args: str = None): 65 | voxel_offset = None 66 | voxel_size = None 67 | shape = None 68 | for inp in inputs: 69 | if isinstance(inp, Chunk): 70 | voxel_offset = inp.voxel_offset 71 | voxel_size = inp.voxel_size 72 | shape = inp.shape 73 | break 74 | if args is not None and '=' in args: 75 | args = str_to_dict(args) 76 | 77 | if len(inputs) == 0 and args is None: 78 | outputs = self.execute() 79 | elif len(inputs) == 0 and args is not None: 80 | if isinstance(args, str): 81 | outputs = self.execute(args=args) 82 | elif isinstance(args, dict): 83 | outputs = self.execute(**args) 84 | else: 85 | raise ValueError(f'unsupported argument: {args}') 86 | elif len(inputs) > 0 and args is None: 87 | outputs = self.execute(*inputs) 88 | else: 89 | if isinstance(args, str): 90 | outputs = self.execute(*inputs, args=args) 91 | elif isinstance(args, dict): 92 | outputs = self.execute(*inputs, **args) 93 | else: 94 | raise ValueError(f'unsupported argument: {args}') 95 | 96 | # assert isinstance(outputs, list) or isinstance(outputs, tuple) or outputs is None 97 | if isinstance(outputs, tuple): 98 | outputs = [*outputs] 99 | # automatically convert the ndarrays to Chunks 100 | if voxel_offset is not None and outputs is not None: 101 | assert shape is not None 102 | if isinstance(outputs, list) or isinstance(outputs, tuple): 103 | for idx, output in enumerate(outputs): 104 | outputs[idx] = array_to_chunk(output, voxel_offset, voxel_size, shape) 105 | elif isinstance(outputs, PointCloud): 106 | pass 107 | elif isinstance(outputs, np.ndarray): 108 | outputs = array_to_chunk(outputs, voxel_offset, voxel_size, shape) 109 | elif isinstance(outputs, Chunk): 110 | # this is good 111 | pass 112 | else: 113 | # will return outputs as is 114 | pass 115 | 116 | return outputs 117 | -------------------------------------------------------------------------------- /chunkflow/flow/save_pngs.py: -------------------------------------------------------------------------------- 1 | from .base import OperatorBase 2 | 3 | from warnings import warn 4 | import os 5 | import numpy as np 6 | 7 | from tqdm import tqdm 8 | import pyspng 9 | from PIL import Image 10 | 11 | from chunkflow.chunk import Chunk 12 | 13 | 14 | class SavePNGsOperator(OperatorBase): 15 | def __init__(self, 16 | output_path: str = './pngs/', 17 | dtype: str = 'uint8', 18 | name: str = 'save-pngs'): 19 | super().__init__(name=name) 20 | 21 | if not os.path.isdir(output_path): 22 | warn(f'output path do not exist, will create %s {output_path}') 23 | os.makedirs(output_path) 24 | 25 | self.dtype = np.dtype(dtype) 26 | self.output_path = output_path 27 | 28 | def __call__(self, chunk: Chunk): 29 | assert isinstance(chunk, Chunk) 30 | if not np.issubdtype(chunk.dtype, self.dtype): 31 | chunk = chunk.astype(self.dtype) 32 | 33 | if chunk.is_affinity_map: 34 | properties = chunk.properties 35 | chunk = (chunk[1,...] + chunk[2,...]) / 2. * 255. 36 | chunk = chunk.astype(np.uint8) 37 | chunk = Chunk(chunk) 38 | chunk.set_properties(properties) 39 | 40 | assert chunk.ndim == 3 41 | for z in tqdm(range(chunk.voxel_offset[0], chunk.bbox.maxpt[0])): 42 | img = chunk.cutout((slice(z,z+1), chunk.slices[1], chunk.slices[2])) 43 | img = img.array[0,:,:] 44 | filename = os.path.join(self.output_path, f"{z:05d}.png") 45 | if np.issubdtype(img.dtype, np.uint16): 46 | # pyspng do not support 16 bit image well. 47 | img = Image.fromarray(img) 48 | img.save(filename) 49 | else: 50 | binary = pyspng.encode( 51 | img, 52 | progressive=pyspng.ProgressiveMode.PROGRESSIVE, 53 | compress_level=6, 54 | ) 55 | with open(filename, "wb") as fout: 56 | fout.write(binary) -------------------------------------------------------------------------------- /chunkflow/flow/save_precomputed.py: -------------------------------------------------------------------------------- 1 | 2 | import time 3 | import os 4 | import json 5 | import numpy as np 6 | 7 | from cloudvolume import CloudVolume 8 | from cloudvolume.lib import Vec, yellow 9 | from cloudfiles import CloudFiles 10 | 11 | from chunkflow.lib.cartesian_coordinate import BoundingBox 12 | from chunkflow.lib.igneous.tasks import downsample_and_upload 13 | from chunkflow.chunk import Chunk 14 | 15 | from .base import OperatorBase 16 | #from .downsample_upload import DownsampleUploadOperator 17 | 18 | 19 | class SavePrecomputedOperator(OperatorBase): 20 | def __init__(self, 21 | volume_path: str, 22 | mip: int, 23 | upload_log: bool = True, 24 | create_thumbnail: bool = False, 25 | fill_missing: bool = False, 26 | parallel: int = 1, 27 | name: str = 'save-precomputed'): 28 | super().__init__(name=name) 29 | 30 | self.upload_log = upload_log 31 | self.create_thumbnail = create_thumbnail 32 | self.mip = mip 33 | 34 | # if not volume_path.startswith('precomputed://'): 35 | # volume_path = 'precomputed://' + volume_path 36 | self.volume_path = volume_path 37 | 38 | # gevent.monkey.patch_all(thread=False) 39 | self.volume = CloudVolume( 40 | self.volume_path, 41 | fill_missing=fill_missing, 42 | bounded=False, 43 | autocrop=True, 44 | mip=self.mip, 45 | cache=False, 46 | green_threads=True, 47 | delete_black_uploads=True, 48 | parallel=parallel, 49 | progress=True) 50 | 51 | if upload_log: 52 | log_path = os.path.join(volume_path, 'log') 53 | self.log_storage = CloudFiles(log_path) 54 | 55 | def create_chunk_with_zeros(self, bbox, num_channels, dtype): 56 | """Create a fake all zero chunk. 57 | this is used in skip some operation based on mask.""" 58 | shape = (num_channels, *bbox.size3()) 59 | arr = np.zeros(shape, dtype=dtype) 60 | chunk = Chunk(arr, voxel_offset=(0, *bbox.minpt)) 61 | return chunk 62 | 63 | def __call__(self, chunk: Chunk, log=None): 64 | assert isinstance(chunk, Chunk) 65 | print(f'save chunk {chunk.bbox.string} to {self.volume_path}') 66 | 67 | start = time.time() 68 | chunk = self._auto_convert_dtype(chunk, self.volume) 69 | 70 | # transpose czyx to xyzc order 71 | arr = np.transpose(chunk.array) 72 | self.volume[chunk.slices[::-1]] = arr 73 | 74 | if self.create_thumbnail: 75 | self._create_thumbnail(chunk) 76 | 77 | # add timer for save operation itself 78 | if log: 79 | log['timer'][self.name] = time.time() - start 80 | 81 | if self.upload_log: 82 | self._upload_log(log, chunk.bbox) 83 | 84 | def _auto_convert_dtype(self, chunk, volume): 85 | """convert the data type to fit volume datatype""" 86 | if np.issubdtype(volume.dtype, np.floating) and np.issubdtype(chunk.dtype, np.uint8): 87 | chunk = chunk.astype(volume.dtype) 88 | chunk /= 255. 89 | # chunk = chunk / chunk.array.max() * np.iinfo(volume.dtype).max 90 | elif np.issubdtype(volume.dtype, np.uint8) and np.issubdtype(chunk.dtype, np.floating): 91 | chunk.max() <= 1. 92 | chunk *= 255 93 | 94 | if volume.dtype != chunk.dtype: 95 | print(yellow(f'converting chunk data type {chunk.dtype} ' + 96 | f'to volume data type: {volume.dtype}')) 97 | # float_chunk = chunk.astype(np.float64) 98 | # chunk = float_chunk / np.iinfo(chunk.dtype).max * np.iinfo(self.volume.dtype).max 99 | # chunk = chunk / chunk.array.max() * np.iinfo(volume.dtype).max 100 | return chunk.astype(volume.dtype) 101 | else: 102 | return chunk 103 | 104 | def _create_thumbnail(self, chunk): 105 | print('creating thumbnail...') 106 | 107 | thumbnail_volume_path = os.path.join(self.volume_path, 'thumbnail') 108 | thumbnail_volume = CloudVolume( 109 | thumbnail_volume_path, 110 | compress='gzip', 111 | fill_missing=True, 112 | bounded=False, 113 | autocrop=True, 114 | mip=self.mip, 115 | cache=False, 116 | green_threads=True, 117 | delete_black_uploads=True, 118 | progress=False) 119 | 120 | # only use the last channel, it is the Z affinity 121 | # if this is affinitymap 122 | image = chunk[-1, :, :, :] 123 | if np.issubdtype(image.dtype, np.floating): 124 | image = (image * 255).astype(np.uint8) 125 | 126 | #self.thumbnail_operator(image) 127 | # transpose to xyzc 128 | image = np.transpose(image) 129 | image_bbox = BoundingBox.from_slices(chunk.slices[::-1][:3]) 130 | 131 | downsample_and_upload(image, 132 | image_bbox, 133 | thumbnail_volume, 134 | Vec(*(image.shape)), 135 | mip=self.mip, 136 | max_mip=6, 137 | axis='z', 138 | skip_first=True, 139 | only_last_mip=True) 140 | 141 | def _upload_log(self, log, output_bbox): 142 | assert log 143 | assert isinstance(output_bbox, BoundingBox) 144 | 145 | print(f'uploaded log: {log}') 146 | 147 | # write to google cloud storage 148 | self.log_storage.put_json(output_bbox.string + 149 | '.json', 150 | content=json.dumps(log)) 151 | -------------------------------------------------------------------------------- /chunkflow/flow/view.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from microviewer import view, hyperview 4 | from .base import OperatorBase 5 | 6 | 7 | class ViewOperator(OperatorBase): 8 | def __init__(self, name: str = 'view'): 9 | super().__init__(name=name) 10 | 11 | def __call__(self, chunk, seg=None): 12 | """view chunk using cloudvolume view""" 13 | # cloudvolume use fortran order 14 | chunk = chunk.transpose() 15 | if seg: 16 | seg = seg.transpose() 17 | hyperview(chunk, seg) 18 | elif np.issubdtype(chunk.dtype, 19 | np.floating) or chunk.dtype == np.uint8: 20 | # this is an image 21 | view(chunk) 22 | else: 23 | view(chunk, seg=True) 24 | -------------------------------------------------------------------------------- /chunkflow/lib/__init__.py: -------------------------------------------------------------------------------- 1 | import types 2 | from importlib.machinery import SourceFileLoader 3 | 4 | 5 | def load_source(fname: str): 6 | """ Imports a module from source. 7 | 8 | Parameters 9 | ----------- 10 | fname: 11 | file path of the python code. 12 | """ 13 | loader = SourceFileLoader("Model", fname) 14 | mod = types.ModuleType(loader.name) 15 | loader.exec_module(mod) 16 | return mod 17 | -------------------------------------------------------------------------------- /chunkflow/lib/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/chunkflow/lib/aws/__init__.py -------------------------------------------------------------------------------- /chunkflow/lib/flow.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from typing import Union 3 | 4 | from functools import update_wrapper, wraps 5 | 6 | import click 7 | from .cartesian_coordinate import Cartesian 8 | 9 | 10 | class CartesianParamType(click.ParamType): 11 | name = 'Cartesian' 12 | 13 | def convert(self, value: Union[list, tuple], param, ctx): 14 | assert len(value) == 3 15 | return Cartesian.from_collection(value) 16 | 17 | CartesianParam = CartesianParamType() 18 | 19 | # global dict to hold the operators and parameters 20 | state = {'operators': {}} 21 | DEFAULT_CHUNK_NAME = 'chunk' 22 | DEFAULT_SYNAPSES_NAME = 'syns' 23 | DEFAULT_SKELETON_NAME = 'skels' 24 | 25 | 26 | def get_initial_task(): 27 | return {'log': {'timer': {}}} 28 | 29 | 30 | def default_none(ctx, _, value): 31 | """ 32 | click currently can not use None with tuple type 33 | it will return an empty tuple if the default=None details: 34 | https://github.com/pallets/click/issues/789 35 | """ 36 | if not value: 37 | return None 38 | else: 39 | return value 40 | 41 | 42 | # the code design is based on: 43 | # https://github.com/pallets/click/blob/master/examples/imagepipe/imagepipe.py 44 | @click.group(chain=True) 45 | @click.option('--mip', '-m', 46 | type=click.INT, default=0, 47 | help='default mip level of chunks.') 48 | @click.option('--dry-run/--real-run', default=False, 49 | help='dry run or real run. default is real run.') 50 | @click.option('--verbose/--quiet', default=False, 51 | help='show more information or not. default is False.') 52 | def main(mip, dry_run, verbose): 53 | """Compose operators and create your own pipeline.""" 54 | 55 | state['mip'] = mip 56 | state['dry_run'] = dry_run 57 | state['verbose'] = verbose 58 | if dry_run: 59 | print('\nYou are using dry-run mode, will not do the work!') 60 | 61 | 62 | @main.result_callback() 63 | def process_commands(operators, mip, dry_run, verbose): 64 | """This result callback is invoked with an iterable of all 65 | the chained subcommands. As in this example each subcommand 66 | returns a function we can chain them together to feed one 67 | into the other, similar to how a pipe on unix works. 68 | """ 69 | # It turns out that a tuple will not work correctly! 70 | stream = [get_initial_task(), ] 71 | 72 | # Pipe it through all stream operators. 73 | for operator in operators: 74 | stream = operator(stream) 75 | # task = next(stream) 76 | 77 | # Evaluate the stream and throw away the items. 78 | for _ in stream: 79 | pass 80 | 81 | 82 | def operator(func): 83 | """ 84 | Help decorator to rewrite a function so that 85 | it returns another function from it. 86 | """ 87 | @wraps(func) 88 | def wrapper(*args, **kwargs): 89 | def operator(stream): 90 | return func(stream, *args, **kwargs) 91 | return operator 92 | 93 | return wrapper 94 | 95 | 96 | def generator(func): 97 | """Similar to the :func:`operator` but passes through old values unchanged 98 | and does not pass through the values as parameter. 99 | """ 100 | @operator 101 | def new_func(stream, *args, **kwargs): 102 | for item in func(*args, **kwargs): 103 | yield item 104 | 105 | return update_wrapper(new_func, func) -------------------------------------------------------------------------------- /chunkflow/lib/gala/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/chunkflow/lib/gala/__init__.py -------------------------------------------------------------------------------- /chunkflow/lib/igneous/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/chunkflow/lib/igneous/__init__.py -------------------------------------------------------------------------------- /chunkflow/lib/igneous/downsample_scales.py: -------------------------------------------------------------------------------- 1 | # @license 2 | # Copyright 2016,2017 The Neuroglancer Authors 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from __future__ import division 16 | from builtins import range 17 | 18 | import numpy as np 19 | import click 20 | 21 | DEFAULT_MAX_DOWNSAMPLING = 64 # maximum factor to downsample by 22 | DEFAULT_MAX_DOWNSAMPLED_SIZE = 128 # minimum length of a side after downsampling 23 | DEFAULT_MAX_DOWNSAMPLING_SCALES = float('inf') 24 | 25 | 26 | def compute_near_isotropic_downsampling_scales( 27 | size, 28 | voxel_size, 29 | dimensions_to_downsample, 30 | max_scales=DEFAULT_MAX_DOWNSAMPLING_SCALES, 31 | max_downsampling=DEFAULT_MAX_DOWNSAMPLING, 32 | max_downsampled_size=DEFAULT_MAX_DOWNSAMPLED_SIZE): 33 | """Compute a list of successive downsampling factors.""" 34 | 35 | num_dims = len(voxel_size) 36 | cur_scale = np.ones((num_dims, ), dtype=click.INT) 37 | scales = [tuple(cur_scale)] 38 | while (len(scales) < max_scales and (np.prod(cur_scale) < max_downsampling) 39 | and np.all(size / cur_scale > max_downsampled_size) 40 | and np.all(np.mod(size, cur_scale) == 0) 41 | ): # the number of chunks should be integer 42 | # Find dimension with smallest voxelsize. 43 | cur_voxel_size = cur_scale * voxel_size 44 | smallest_cur_voxel_size_dim = dimensions_to_downsample[np.argmin( 45 | cur_voxel_size[dimensions_to_downsample])] 46 | cur_scale[smallest_cur_voxel_size_dim] *= 2 47 | target_voxel_size = cur_voxel_size[smallest_cur_voxel_size_dim] * 2 48 | for d in dimensions_to_downsample: 49 | if d == smallest_cur_voxel_size_dim: 50 | continue 51 | d_voxel_size = cur_voxel_size[d] 52 | if abs(d_voxel_size - target_voxel_size) > abs(d_voxel_size * 2 - 53 | target_voxel_size): 54 | cur_scale[d] *= 2 55 | scales.append(tuple(cur_scale)) 56 | return scales 57 | 58 | 59 | def compute_two_dimensional_near_isotropic_downsampling_scales( 60 | size, 61 | voxel_size, 62 | max_scales=float('inf'), 63 | max_downsampling=DEFAULT_MAX_DOWNSAMPLING, 64 | max_downsampled_size=DEFAULT_MAX_DOWNSAMPLED_SIZE): 65 | """Compute a list of successive downsampling factors for 2-d tiles.""" 66 | 67 | max_scales = min(max_scales, 10) 68 | 69 | # First compute a set of 2-d downsamplings for XY, XZ, and YZ with a high 70 | # number of max_scales, and ignoring other criteria. 71 | scales_transpose = [ 72 | compute_near_isotropic_downsampling_scales( 73 | size=size, 74 | voxel_size=voxel_size, 75 | dimensions_to_downsample=dimensions_to_downsample, 76 | max_scales=max_scales, 77 | max_downsampling=float('inf'), 78 | max_downsampled_size=0, 79 | ) for dimensions_to_downsample in [[0, 1], [0, 2], [1, 2]] 80 | ] 81 | 82 | # Truncate all list of scales to the same length, once the stopping criteria 83 | # is reached for all values of dimensions_to_downsample. 84 | scales = [((1, ) * 3, ) * 3] 85 | size = np.array(size) 86 | 87 | def scale_satisfies_criteria(scale): 88 | return np.prod(scale) < max_downsampling and ( 89 | size / scale).max() > max_downsampled_size 90 | 91 | for i in range(1, max_scales): 92 | cur_scales = tuple(scales_transpose[d][i] for d in range(3)) 93 | if all(not scale_satisfies_criteria(scale) for scale in cur_scales): 94 | break 95 | scales.append(cur_scales) 96 | return scales 97 | 98 | 99 | def compute_plane_downsampling_scales( 100 | size, 101 | preserve_axis='z', 102 | max_scales=DEFAULT_MAX_DOWNSAMPLING_SCALES, 103 | max_downsampling=DEFAULT_MAX_DOWNSAMPLING, 104 | max_downsampled_size=DEFAULT_MAX_DOWNSAMPLED_SIZE): 105 | 106 | axis_map = {'x': 0, 'y': 1, 'z': 2} 107 | preserve_axis = axis_map[preserve_axis] 108 | 109 | size = np.array(size) 110 | 111 | if np.any(size <= 0): 112 | return [(1, 1, 1)] 113 | 114 | size[preserve_axis] = size[(preserve_axis + 1) % 3] 115 | 116 | dimension = min(*size) 117 | num_downsamples = int(np.log2(dimension / max_downsampled_size)) 118 | num_downsamples = min(num_downsamples, max_scales) 119 | 120 | factor = 2 121 | scales = [(1, 1, 1)] 122 | for i in range(num_downsamples): 123 | if factor > max_downsampling: 124 | break 125 | elif dimension / factor < max_downsampled_size: 126 | break 127 | 128 | scale = [factor, factor, factor] 129 | scale[preserve_axis] = 1 130 | 131 | scales.append(tuple(scale)) 132 | factor *= 2 133 | 134 | return scales 135 | -------------------------------------------------------------------------------- /chunkflow/lib/igneous/tasks.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 4 | from cloudvolume.lib import min2, Vec 5 | 6 | from . import downsample, downsample_scales 7 | 8 | 9 | def downsample_and_upload(image, 10 | bounds, 11 | vol, 12 | ds_shape, 13 | mip=0, 14 | max_mip: int=None, 15 | axis='z', 16 | skip_first=False, 17 | sparse=False, 18 | only_last_mip=True): 19 | """ 20 | mip:int, the current mip level of image 21 | only_last_mip::bool, only save the last mip level or not. 22 | In default as False, we'll save all the intermediate mip level 23 | """ 24 | ds_shape = min2(vol.volume_size, ds_shape[:3]) 25 | 26 | # sometimes we downsample a base layer of 512x512 27 | # into underlying chunks of 64x64 which permits more scales 28 | underlying_mip = (mip + 1) if (mip + 1) in vol.available_mips else mip 29 | underlying_shape = vol.mip_underlying(underlying_mip).astype(np.float32) 30 | toidx = {'x': 0, 'y': 1, 'z': 2} 31 | preserved_idx = toidx[axis] 32 | underlying_shape[preserved_idx] = float('inf') 33 | 34 | # Need to use ds_shape here. Using image bounds means truncated 35 | # edges won't generate as many mip levels 36 | fullscales = downsample_scales.compute_plane_downsampling_scales( 37 | size=ds_shape, 38 | preserve_axis=axis, 39 | max_downsampled_size=int(min(*underlying_shape)), 40 | ) 41 | if max_mip: 42 | fullscales = fullscales[:(max_mip-mip)] 43 | factors = downsample.scale_series_to_downsample_factors(fullscales) 44 | 45 | if len(factors) == 0: 46 | print( 47 | "No factors generated. Image Shape: {}, Downsample Shape: {}, Volume Shape: {}, Bounds: {}" 48 | .format(image.shape, ds_shape, vol.volume_size, bounds)) 49 | 50 | downsamplefn = downsample.method(vol.layer_type, sparse=sparse) 51 | 52 | vol.mip = mip 53 | if not skip_first: 54 | vol[bounds.slices] = image 55 | 56 | new_bounds = bounds.clone() 57 | 58 | for factor3 in factors: 59 | vol.mip += 1 60 | image = downsamplefn(image, factor3) 61 | new_bounds //= factor3 62 | delta = Cartesian.from_collection(image.shape[:3]) 63 | new_bounds = BoundingBox.from_delta(new_bounds.start, delta) 64 | if factor3 is factors[-1]: 65 | # this is the last mip level 66 | vol[new_bounds.slices] = image 67 | else: 68 | # this is not the last mip level 69 | if not only_last_mip: 70 | vol[new_bounds.slices] = image 71 | -------------------------------------------------------------------------------- /chunkflow/lib/region_of_interest.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import Union 4 | 5 | import numpy as np 6 | 7 | from .cartesian_coordinate import Cartesian, BoundingBox 8 | 9 | 10 | class RegionOfInterest(BoundingBox): 11 | """Physical Region of Interest. 12 | Basically Bounding Box with voxel size. 13 | 14 | Args: 15 | BoundingBox (BoundingBox): _description_ 16 | """ 17 | def __init__(self, 18 | minpt: Union[list, Cartesian], maxpt: Union[list, Cartesian], 19 | voxel_size: Cartesian, 20 | dtype=None): 21 | """Physical region of interest. 22 | Basically BoundingBox with voxel size. 23 | 24 | Args: 25 | minpt (Union[list, Cartesian]): the start of bounding box 26 | maxpt (Union[list, Cartesian]): the stop of bounding box 27 | voxel_size (Cartesian): the physical size of each voxel/unit. We normally use nanometer here. 28 | dtype (_type_, optional): integer datatype of bounding box. Defaults to None. 29 | """ 30 | super().__init__(minpt, maxpt, dtype) 31 | 32 | self.voxel_size = voxel_size 33 | 34 | @classmethod 35 | def from_bbox(cls, bbox: BoundingBox, voxel_size: Cartesian): 36 | return cls(bbox.minpt, bbox.maxpt, voxel_size) 37 | 38 | def __repr__(self): 39 | return f'from {self.minpt} to {self.maxpt}, dtype={self.dtype}, voxel_size={self.voxel_size})' 40 | 41 | @property 42 | def bounding_box(self): 43 | return BoundingBox(self.minpt, self.maxpt, dtype=self.dtype) 44 | 45 | @property 46 | def physical_size(self): 47 | return self.voxel_size * self.shape 48 | 49 | def clone(self): 50 | bbox = self.bounding_box.clone() 51 | # the Cartesian is immutable, so we do not need to clone it. 52 | return RegionOfInterest(bbox.minpt, bbox.maxpt, self.voxel_size, dtype=bbox.dtype) 53 | 54 | def slices_in_scale(self, voxel_size: Cartesian) -> tuple: 55 | """slices with specific voxel size volume 56 | 57 | Args: 58 | voxel_size (Cartesian): the target volume voxel size 59 | 60 | Returns: 61 | tuple: tuple of slices 62 | """ 63 | minpt = tuple( p * s1 // s2 for p, s1, s2 in zip( 64 | self.minpt, self.voxel_size, voxel_size 65 | )) 66 | maxpt = tuple( p * s1 // s2 for p, s1, s2 in zip( 67 | self.maxpt, self.voxel_size, voxel_size 68 | )) 69 | bbox = BoundingBox(minpt, maxpt) 70 | return bbox.slices 71 | 72 | 73 | class ROITree: 74 | def __init__(self, roi: RegionOfInterest, axis: int, 75 | left: ROITree, right: ROITree): 76 | """Split the volume hierarchically using a modified aligned K-D tree 77 | 78 | Args: 79 | roi (RegionOfInterest): physical region covered 80 | axis (int): the splitting dimension of left and right roi 81 | left (ROIKDTree): left roi with lower coordinate values in axis dim 82 | right (ROIKDTree): right roi with higher coordinate values in axis dim 83 | """ 84 | assert axis >= 0 and axis < 3 85 | self.roi = roi 86 | self.axis = axis 87 | self.left = left 88 | self.right = right 89 | 90 | @classmethod 91 | def from_roi(cls, roi: RegionOfInterest, factor: Cartesian, 92 | atomic_block_size: Cartesian, atomic_voxel_size: Cartesian): 93 | """Generate the ROITree from a single roi. 94 | This roi is not required to be aligned with the atomic block size. 95 | If it is not aligned, a roi will partially cover the volume. 96 | 97 | Args: 98 | roi (RegionOfInterest): the total roi covered by this tree. 99 | factor (Cartesian): downsampling factor in each level. 100 | atomic_block_size (Cartesian): the size of the leaf node/block 101 | atomic_voxel_size (Cartesian): the voxel size of leaf block 102 | """ 103 | pass 104 | 105 | # assert roi.voxel_size % atomic_voxel_size == Cartesian(0, 0, 0) 106 | # assert roi.voxel_size // atomic_voxel_size % factor == Cartesian(0, 0, 0) 107 | 108 | # if roi.voxel_size == atomic_voxel_size: 109 | # # this is the leaf roi/block 110 | # return cls(roi, None, None, None) 111 | 112 | # # find the relatively longest axis to split 113 | # children_voxel_size = roi.voxel_size // factor 114 | # block_nums = roi.physical_size / (children_voxel_size * ) 115 | # block_nums = np.ceil(block_nums) 116 | # axis = np.argmax(block_nums) 117 | 118 | # # split along axis 119 | # left_start = roi.start * factor 120 | # left_block_nums = 121 | # left_stop = left_start + 122 | # left_roi = RegionOfInterest() 123 | # left = cls.from_roi(left_roi, factor, atomic_block_size, atomic_voxel_size) 124 | 125 | 126 | 127 | @property 128 | def is_leaf(self): 129 | return self.axis is None -------------------------------------------------------------------------------- /chunkflow/lib/utils.py: -------------------------------------------------------------------------------- 1 | from ast import literal_eval 2 | 3 | import numpy as np 4 | 5 | # from chunkflow.chunk import Chunk 6 | # from chunkflow.volume import PrecomputedVolume 7 | # from chunkflow.flow.plugin import str_to_dict 8 | 9 | 10 | 11 | def simplest_type(s: str): 12 | try: 13 | return literal_eval(s) 14 | except: 15 | return s 16 | 17 | def str_to_dict(string: str): 18 | keywords = {} 19 | for item in string.split(';'): 20 | assert '=' in item 21 | item = item.split('=') 22 | keywords[item[0]] = simplest_type(item[1]) 23 | return keywords 24 | 25 | -------------------------------------------------------------------------------- /chunkflow/plugins/agglomerate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.chunk import Chunk 4 | 5 | from waterz import agglomerate 6 | 7 | 8 | def execute(affs: Chunk, 9 | fragments: np.ndarray = None, 10 | threshold: float = 0.7, 11 | aff_threshold_low: float = 0.001, 12 | aff_threshold_high: float = 0.9999, 13 | scoring_function: str = 'OneMinus>', 14 | flip_channel: bool = True, 15 | ): 16 | """ 17 | Mean/max agglomeration of affinity map including watershed step. 18 | 19 | Parameters: 20 | ----------- 21 | affs: affinity map with 4 dimensions: channel, z, y, x 22 | """ 23 | properties = affs.properties 24 | 25 | # our affinity map channel order is x,y,z! 26 | # meaning the first channel is x, the second is y, 27 | # the third is z. We have to reverse the channel for waterz. 28 | if flip_channel: 29 | affs = np.flip(affs, axis=0) 30 | 31 | # waterz need datatype float32 and 32 | # the array memory layout is contiguous 33 | affs = np.ascontiguousarray(affs, dtype=np.float32) 34 | 35 | # the output is a generator, and the computation is delayed 36 | seg_generator = agglomerate( 37 | affs, [threshold], fragments=fragments, 38 | aff_threshold_low=aff_threshold_low, 39 | aff_threshold_high=aff_threshold_high, 40 | scoring_function=scoring_function, 41 | force_rebuild=False) 42 | # there is only one threshold, so there is also only one result 43 | # in generator 44 | seg = next(seg_generator) 45 | seg = Chunk(seg) 46 | seg.set_properties(properties) 47 | 48 | return [seg] 49 | -------------------------------------------------------------------------------- /chunkflow/plugins/aggregate_skeleton_fragments.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import re 4 | from collections import defaultdict 5 | from time import sleep 6 | 7 | from cloudfiles import CloudFiles 8 | from cloudvolume import PrecomputedSkeleton 9 | 10 | import kimimaro 11 | 12 | from .base import OperatorBase 13 | 14 | 15 | class AggregateSkeletonFragmentsOperator(OperatorBase): 16 | """Merge skeleton fragments for Neuroglancer visualization.""" 17 | def __init__(self, 18 | fragments_path: str, 19 | output_path: str, 20 | name: str = 'aggregate-skeleton-fragments'): 21 | """ 22 | Parameters 23 | ------------ 24 | fragments_path: 25 | path to store fragment files 26 | output_path: 27 | save the merged skeleton file here. 28 | """ 29 | super().__init__(name=name) 30 | self.fragments_storage = CloudFiles(fragments_path) 31 | self.output_storage = CloudFiles(output_path) 32 | 33 | def __call__(self, prefix: str): 34 | """To-do: convert it to a plugin 35 | 36 | Args: 37 | prefix (str): _description_ 38 | """ 39 | print(f'aggregate skeletons with prefix of {prefix}') 40 | 41 | id2filenames = defaultdict(list) 42 | for filename in self.fragments_storage.list_files(prefix=prefix): 43 | filename = os.path.basename(filename) 44 | # `match` implies the beginning (^). `search` matches whole string 45 | matches = re.search(r'(\d+):', filename) 46 | 47 | if not matches: 48 | continue 49 | 50 | # skeleton ID 51 | skl_id = int(matches.group(0)[:-1]) 52 | id2filenames[skl_id].append(filename) 53 | 54 | for skl_id, filenames in id2filenames.items(): 55 | print(f'skeleton id: {skl_id}') 56 | frags = self.fragments_storage.get(filenames) 57 | frags = [PrecomputedSkeleton.from_precomputed(x['content']) for x in frags] 58 | skel = PrecomputedSkeleton.simple_merge(frags).consolidate() 59 | skel = kimimaro.postprocess( 60 | skel, 61 | dust_threshold=1000, 62 | tick_threshold=3500 63 | ) 64 | self.output_storage.put( 65 | file_path=str(skl_id), 66 | content=skel.to_precomputed(), 67 | ) 68 | # the last few hundred files will not be uploaded without sleeping! 69 | sleep(0.01) 70 | -------------------------------------------------------------------------------- /chunkflow/plugins/aws/cloud_watch.py: -------------------------------------------------------------------------------- 1 | from warnings import warn 2 | import boto3 3 | from cloudvolume.secrets import aws_credentials 4 | 5 | 6 | class CloudWatch: 7 | """monitor time elapsed of each operator using AWS CloudWatch.""" 8 | def __init__(self, log_name: str, credentials: dict = None): 9 | """ 10 | Parameters 11 | ---------- 12 | log_name: 13 | the name of the log. 14 | credentials: 15 | aws credential. If it is None, we'll try to construct it automatically. 16 | """ 17 | self.log_name = log_name 18 | if not credentials: 19 | credentials = aws_credentials() 20 | self.client = boto3.client( 21 | 'cloudwatch', 22 | region_name=credentials['AWS_DEFAULT_REGION'], 23 | aws_secret_access_key=credentials['AWS_SECRET_ACCESS_KEY'], 24 | aws_access_key_id=credentials['AWS_ACCESS_KEY_ID']) 25 | 26 | def put_metric_data(self, log: dict): 27 | """ 28 | Parameters 29 | ----------- 30 | log: 31 | the log containing the operator time elapsed in the `timer` key. 32 | """ 33 | assert isinstance(log, dict) 34 | 35 | if 'compute_device' in log: 36 | compute_device = log['compute_device'] 37 | else: 38 | warn( 39 | 'did not find compute device in log, will create one based on CPU.' 40 | ) 41 | import platform 42 | compute_device = platform.processor() 43 | 44 | # create metric data 45 | metric_data = [] 46 | for key, value in log['timer'].items(): 47 | metric_data.append({ 48 | 'MetricName': 49 | key, 50 | 'Dimensions': [{ 51 | 'Name': 'compute_device', 52 | 'Value': compute_device 53 | }], 54 | 'Value': 55 | value, 56 | 'Unit': 57 | 'Seconds' 58 | }) 59 | 60 | # submit the metric data 61 | self.client.put_metric_data(Namespace=self.log_name, 62 | MetricData=metric_data) 63 | 64 | 65 | def execute(log: dict, log_name: str): 66 | cw = CloudWatch(log_name=log_name) 67 | cw.put_metric_data(log) 68 | print(f'log submitted to cloud watch: {log}') 69 | -------------------------------------------------------------------------------- /chunkflow/plugins/aws/test_cloud_watch.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import pytest 3 | 4 | from chunkflow.lib.aws.cloud_watch import CloudWatch 5 | 6 | 7 | @pytest.mark.skip(reason="need credential") 8 | class TestCloudWatch(unittest.TestCase): 9 | def setUp(self): 10 | print('test cloud watch...') 11 | log_name = 'chunkflow-test' 12 | self.cloud_watch = CloudWatch(log_name) 13 | 14 | def test_put_metric_data(self): 15 | log = { 16 | 'compute_device': 'X86-64', 17 | 'timer': { 18 | 'cutout': 24, 19 | 'inference': 200, 20 | 'save': 50 21 | } 22 | } 23 | self.cloud_watch.put_metric_data(log) 24 | 25 | 26 | if __name__ == '__main__': 27 | unittest.main() 28 | -------------------------------------------------------------------------------- /chunkflow/plugins/clahe.py: -------------------------------------------------------------------------------- 1 | from tqdm import tqdm 2 | import cv2 3 | 4 | from chunkflow.chunk import Chunk 5 | 6 | clipLimit = 20.0 7 | tileGridSize = (512, 512) 8 | 9 | #clahe = cv2.createCLAHE(clipLimit=5., tileGridSize = (18, 26)) 10 | clahe = cv2.createCLAHE(clipLimit=clipLimit, tileGridSize = tileGridSize) 11 | 12 | def execute(img: Chunk): 13 | assert img.is_image 14 | for z in tqdm(range(img.shape[-3])): 15 | img.array[...,z,:,:] = clahe.apply(img.array[...,z,:,:]) 16 | return img 17 | -------------------------------------------------------------------------------- /chunkflow/plugins/cutout_dvid_label.py: -------------------------------------------------------------------------------- 1 | import json 2 | from cloudvolume.lib import Bbox 3 | import pydvid 4 | 5 | from chunkflow.chunk import Chunk 6 | 7 | # these example parameters are from the example of [neuclease](https://github.com/janelia-flyem/neuclease) 8 | supervoxels = False 9 | server = "http://hemibrain-dvid.janelia.org:80" 10 | server = server.strip('/') 11 | if ':' not in server[-5:]: 12 | # add a default port 13 | server += ':8000' 14 | 15 | instance = "segmentation" 16 | # uuid = dvid.find_master(server) 17 | uuid = "20631f" 18 | print('server: ', server) 19 | print('instance: ', instance) 20 | print('uuid: ', uuid) 21 | 22 | config = {} 23 | config['supervoxels'] = False 24 | config['server'] = server 25 | config['instance'] = 'segmentation' 26 | config['uuid'] = uuid 27 | 28 | config_str = json.dumps(config) 29 | 30 | def execute(bbox: Bbox, args: str= config_str ): 31 | if args is not None: 32 | config = json.loads(args) 33 | if 'server' in config: 34 | server = config['server'] 35 | if 'supervoxels' in config: 36 | supervoxels = config['supervoxels'] 37 | if 'instance' in config: 38 | instance = config['instance'] 39 | if 'uuid' in config: 40 | uuid = config['uuid'] 41 | 42 | if ":main" in uuid or ":master" in uuid: 43 | uuid, _ = uuid.split(':') 44 | #uuid = dvid.find_master(uuid) 45 | 46 | print('bounding box: ', bbox) 47 | box = [tuple(bbox.minpt), tuple(bbox.maxpt)] 48 | subvol = pydvid.fetch_labelmap_voxels( 49 | server, uuid, instance, box, 50 | scale=0, 51 | supervoxels=supervoxels) 52 | # print('cutout volume: \n', subvol) 53 | chunk = Chunk(subvol, voxel_offset=bbox.minpt) 54 | 55 | return [chunk] -------------------------------------------------------------------------------- /chunkflow/plugins/czann_inference.py: -------------------------------------------------------------------------------- 1 | #%% load model 2 | 3 | -------------------------------------------------------------------------------- /chunkflow/plugins/detect_points.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.chunk import Chunk 4 | from chunkflow.chunk.probability_map import ProbabilityMap 5 | from chunkflow.point_cloud import PointCloud 6 | 7 | 8 | def execute(prob: Chunk, 9 | min_distance: float = 15., 10 | threshold_rel: float = 0.3): 11 | if prob is None: 12 | print('get None probability map!') 13 | return None 14 | 15 | assert threshold_rel > 0. 16 | assert threshold_rel < 1. 17 | 18 | if np.issubdtype(prob.dtype, np.uint8): 19 | prob = prob.astype(np.float32) 20 | prob /= 255. 21 | prob = ProbabilityMap.from_chunk(prob) 22 | 23 | # drop confidence for now 24 | points, _= prob.detect_points( 25 | min_distance=min_distance, 26 | threshold_rel=threshold_rel 27 | ) 28 | 29 | points = PointCloud(points, prob.voxel_size) 30 | return points -------------------------------------------------------------------------------- /chunkflow/plugins/fill_segmentation_holes.py: -------------------------------------------------------------------------------- 1 | from chunkflow.chunk import Chunk 2 | import fill_voids 3 | 4 | def execute(seg: Chunk): 5 | properties = seg.properties 6 | array = fill_voids.fill(seg.array) 7 | seg2 = Chunk(array) 8 | seg2.set_properties(properties) 9 | return [seg2,] -------------------------------------------------------------------------------- /chunkflow/plugins/gaussian_filter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from chunkflow.chunk import Chunk 3 | 4 | from scipy.ndimage import gaussian_filter 5 | 6 | 7 | def execute(chunk: Chunk, sigma: float=1., inplace=False): 8 | if not inplace: 9 | chunk = chunk.clone() 10 | for z in range(chunk.shape[-3]): 11 | if chunk.ndim == 4: 12 | for channel in range(chunk.shape[0]): 13 | chunk.array[channel, z,:,:] = gaussian_filter( 14 | chunk.array[channel,z,:,:], sigma=sigma) 15 | elif chunk.ndim == 3: 16 | chunk.array[z,:,:] = gaussian_filter( 17 | chunk.array[z,:,:], sigma=sigma) 18 | else: 19 | raise ValueError(f'only support 4 or 3d, but got {chunk.ndim}') 20 | 21 | # 2d processing in xy plane 22 | # chunk.array = gaussian_filter(chunk.array, sigma=sigma) 23 | return chunk 24 | -------------------------------------------------------------------------------- /chunkflow/plugins/inverse.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def execute(chunk: np.ndarray): 4 | if np.issubdtype(chunk.dtype, np.uint8): 5 | chunk = 255 - chunk 6 | elif np.issubdtype(chunk.dtype, np.float32) and chunk.max() <= 1 and chunk.min()>=0 : 7 | chunk = 1.0 - chunk 8 | else: 9 | raise TypeError("unsupported chunk data type.") 10 | return [chunk] 11 | -------------------------------------------------------------------------------- /chunkflow/plugins/load_mrc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 4 | from chunkflow.chunk import Chunk 5 | 6 | import mrcfile 7 | 8 | def execute(bbox: BoundingBox, fname: str = None): 9 | 10 | with mrcfile.mmap(fname) as mrc: 11 | # print(mrc.header) 12 | # print(f'volume shape: {mrc.data.shape}') 13 | # print(f'voxel size: {mrc.voxel_size}') 14 | voxel_size = mrc.voxel_size 15 | # voxel_size = [int(x/10.) for x in mrc.voxel_size] 16 | # voxel_size = Cartesian.from_collection(voxel_size) 17 | # print(f'voxel size: {voxel_size}') 18 | arr = mrc.data[bbox.slices] 19 | 20 | # somehow, mrcfile pick it up with int8! 21 | arr = arr.view(np.uint8) 22 | z = int(voxel_size.z / 10) 23 | y = int(voxel_size.y / 10) 24 | x = int(voxel_size.x / 10) 25 | voxel_size = Cartesian(z, y, x) 26 | chunk = Chunk(arr, voxel_offset=bbox.start, voxel_size=voxel_size) 27 | return chunk 28 | -------------------------------------------------------------------------------- /chunkflow/plugins/load_n5.py: -------------------------------------------------------------------------------- 1 | 2 | import zarr 3 | 4 | import boto3 5 | from botocore import UNSIGNED 6 | from botocore.client import Config 7 | s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED)) 8 | 9 | from chunkflow.chunk import Chunk 10 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 11 | 12 | def execute(bbox: BoundingBox, 13 | n5_dir: str = None, 14 | group_path: str = None, 15 | voxel_size: tuple = None, 16 | type: str = None, 17 | driver: str = 'n5', ): 18 | # container_url = os.path.join(n5_dir) 19 | # only n5 driver is supported for now 20 | assert driver == 'n5' 21 | if isinstance(voxel_size, tuple): 22 | voxel_size = Cartesian.from_collection(voxel_size) 23 | 24 | fsstore = zarr.N5FSStore(n5_dir, anon=True) 25 | img_zarr = zarr.open(fsstore, path=group_path) 26 | img_arr = img_zarr[bbox.slices] 27 | 28 | img_chk = Chunk(img_arr, 29 | voxel_offset = bbox.start, 30 | voxel_size = voxel_size, 31 | type = type 32 | ) 33 | 34 | return img_chk 35 | 36 | 37 | -------------------------------------------------------------------------------- /chunkflow/plugins/load_nrrd.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.chunk import Chunk 4 | import nrrd 5 | 6 | 7 | def execute(file_name: str=None, voxel_offset: tuple=None, 8 | voxel_size: tuple=None, dtype: str=None, layer_type: str=None): 9 | """Read NRRD file.""" 10 | arr, _ = nrrd.read(file_name) 11 | 12 | if dtype: 13 | arr = arr.astype(dtype) 14 | 15 | chunk = Chunk(arr, voxel_offset=voxel_offset, voxel_size=voxel_size) 16 | breakpoint() 17 | return chunk 18 | -------------------------------------------------------------------------------- /chunkflow/plugins/load_tensorstore.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from chunkflow.chunk import Chunk 3 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 4 | import tensorstore as ts 5 | 6 | 7 | def execute(bbox: BoundingBox, driver: str=None, kvstore: str=None, cache: int=None, 8 | voxel_size: tuple=None, output_name: str='chunk'): 9 | """Load chunk from dataset using tensorstore""" 10 | if driver == 'n5': 11 | kv_driver, path = kvstore.split('://') 12 | kvstore = { 13 | 'driver': kv_driver, 14 | 'path': path 15 | } 16 | dataset_future = ts.open({ 17 | 'driver': driver, 18 | 'kvstore': kvstore, 19 | 'context': { 20 | 'cache_pool': { 21 | 'total_bytes_limit': cache, 22 | } 23 | }, 24 | 'recheck_cached_data': 'open', 25 | }) 26 | dataset = dataset_future.result() 27 | 28 | slices = bbox.slices 29 | arr = dataset[slices[0], slices[1], slices[2]].read().result() 30 | assert arr.ndim == 4 31 | arr = arr.transpose() 32 | if arr.shape[0] == 1: 33 | arr = np.squeeze(arr, axis=0) 34 | chunk = Chunk(arr, 35 | voxel_offset = bbox.start, 36 | voxel_size = Cartesian.from_collection(voxel_size), 37 | ) 38 | return chunk 39 | -------------------------------------------------------------------------------- /chunkflow/plugins/mapto01.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def execute(chunk: np.ndarray): 4 | if np.issubdtype(chunk.dtype, np.uint8): 5 | chunk = chunk.astype(np.float32) 6 | chunk = chunk / 255.0 7 | assert np.issubdtype(chunk.dtype, np.float32) 8 | assert chunk.min() >= -1 9 | assert chunk.max() <= 1 10 | chunk = (chunk + 1.0) / 2.0 11 | return [chunk] 12 | -------------------------------------------------------------------------------- /chunkflow/plugins/median_filter.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from scipy.ndimage import median_filter 4 | 5 | def execute(chunk: np.ndarray, size: tuple=(3,1,1), mode: str='reflect'): 6 | print('median filtering of chunk...') 7 | chunk = median_filter(chunk, size=size, mode=mode) 8 | return [chunk,] -------------------------------------------------------------------------------- /chunkflow/plugins/napari_pixel_classifier.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | from chunkflow.chunk import Chunk 5 | 6 | import apoc 7 | 8 | 9 | def execute(img: Chunk, opencl_filename: str='./PixelClassifier.cl', 10 | obj_id: int=None, reverse: bool = False): 11 | """ 12 | Note that we need to set the following command to use the opencl backends in cluster nodes. 13 | module load modules/2.1.1-20230405 gcc/11.3.0 boost/1.80.0 14 | export OCL_ICD_VENDORS="$CONDA_PREFIX/etc/OpenCL/vendors" 15 | export XDG_CACHE_HOME="/tmp/opencl" 16 | """ 17 | assert os.path.exists(opencl_filename), \ 18 | f'traied classifier file not found: {opencl_filename}' 19 | clf = apoc.PixelClassifier(opencl_filename=opencl_filename) 20 | breakpoint() 21 | pred = clf.predict(img.array) 22 | pred = np.asarray(pred) 23 | 24 | if obj_id is not None: 25 | pred = (pred == obj_id) 26 | if reverse: 27 | pred = np.logical_not(pred) 28 | else: 29 | pred = pred.astype(np.uint32) 30 | 31 | pred = Chunk(pred) 32 | pred.set_properties(img.properties) 33 | pred.layer_type = 'segmentation' 34 | 35 | return pred -------------------------------------------------------------------------------- /chunkflow/plugins/print_max_id.py: -------------------------------------------------------------------------------- 1 | from chunkflow.chunk import Chunk 2 | 3 | def execute(seg: Chunk): 4 | print(f'maximum object ID: {seg.max()}') 5 | -------------------------------------------------------------------------------- /chunkflow/plugins/requirements.txt: -------------------------------------------------------------------------------- 1 | scipy 2 | mrcfile -------------------------------------------------------------------------------- /chunkflow/plugins/rescale_intensity.py: -------------------------------------------------------------------------------- 1 | 2 | from skimage.exposure import rescale_intensity 3 | 4 | from chunkflow.chunk import Chunk 5 | 6 | def execute(img: Chunk, low: int = None, high: int = None): 7 | img.array = rescale_intensity(img.array, in_range=(low, high), out_range=(0, 255)) 8 | return img -------------------------------------------------------------------------------- /chunkflow/plugins/save_nrrd.py: -------------------------------------------------------------------------------- 1 | import nrrd 2 | import numpy as np 3 | from chunkflow.chunk import Chunk 4 | 5 | def execute(chk: Chunk, file_name: str=None): 6 | breakpoint() 7 | chk.array = np.squeeze(chk.array, axis=0) 8 | 9 | if file_name is None: 10 | file_name = f'{chk.bbox.string}.nrrd' 11 | elif not file_name.endswith('.nrrd'): 12 | file_name += f'_{chk.bbox.string}.nrrd' 13 | 14 | print(f'write chunk to file: {file_name}') 15 | nrrd.write(file_name, chk.array) 16 | 17 | -------------------------------------------------------------------------------- /chunkflow/plugins/skeletonize.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import multiprocessing 4 | 5 | import kimimaro 6 | from chunkflow.chunk import Chunk 7 | 8 | from cloudfiles import CloudFiles 9 | 10 | def skeletonize(seg: Chunk, voxel_size): 11 | skels = kimimaro.skeletonize( 12 | seg.array, 13 | anisotropy=voxel_size, 14 | parallel=multiprocessing.cpu_count() // 2 15 | ) 16 | return skels 17 | 18 | def execute(seg: Chunk, voxel_size: tuple, output_path: str=None): 19 | """Create mesh files from segmentation.""" 20 | storage = CloudFiles(output_path) 21 | 22 | if seg is None: 23 | print('no segmentation, skip computation.') 24 | return None 25 | 26 | print('skeletonize segmentation...') 27 | 28 | skels = skeletonize(seg, voxel_size) 29 | bbox_str = seg.bbox.string 30 | for neuron_id, skel in skels.items(): 31 | file_name = f'{neuron_id}:{bbox_str}' 32 | storage.put(file_name, skel.to_precomputed()) 33 | return skels 34 | -------------------------------------------------------------------------------- /chunkflow/plugins/statistics.py: -------------------------------------------------------------------------------- 1 | from chunkflow.chunk import Chunk 2 | import fastremap 3 | 4 | 5 | def execute(chunk: Chunk): 6 | if chunk.is_segmentation: 7 | uniq = fastremap.unique(chunk.array, return_counts=False) 8 | print(f'{len(uniq)} objects with min id {uniq.min()} and max id {uniq.max()}') -------------------------------------------------------------------------------- /chunkflow/plugins/stretch_intensity.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | from chunkflow.chunk import Chunk 5 | 6 | 7 | def execute(chk: Chunk, min_threshold: int = 0, max_threshold: int = 256): 8 | assert np.issubdtype(chk.dtype, np.uint8) 9 | assert min_threshold>=0 10 | assert max_threshold <= 256 11 | assert min_threshold < max_threshold 12 | 13 | lookup_table = np.arange(256, dtype=np.float32) 14 | lookup_table -= min_threshold 15 | lookup_table /= (max_threshold - min_threshold) 16 | lookup_table[lookup_table<0.] = 0. 17 | lookup_table[lookup_table>1.] = 1. 18 | lookup_table *= 255 19 | lookup_table = lookup_table.astype(np.uint8) 20 | 21 | chk.array = lookup_table[chk.array] 22 | return chk 23 | 24 | -------------------------------------------------------------------------------- /chunkflow/plugins/synapse/adjust_pre.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | import numpy as np 3 | 4 | from chunkflow.chunk import Chunk 5 | 6 | from chunkflow.synapses import Synapses 7 | 8 | from edt import edt 9 | 10 | 11 | def execute(syns: Synapses, img: Chunk, img_threshold: int=60, radius: int=4) -> Synapses: 12 | syns2 = deepcopy(syns) 13 | 14 | mask = img.array < img_threshold 15 | dt = edt(mask) 16 | 17 | for idx in range(syns2.pre_num): 18 | coord = syns2.pre[idx, :] 19 | 20 | bbox = img.bbox.clone() 21 | shrinked_bbox = bbox.adjust(-radius) 22 | if not shrinked_bbox.contains(coord): 23 | print(f'skip this tbar: {coord}') 24 | continue 25 | coord = coord - img.voxel_offset 26 | 27 | box = dt[ 28 | coord[0]-radius : coord[0]+radius, 29 | coord[1]-radius : coord[1]+radius, 30 | coord[2]-radius : coord[2]+radius, 31 | ] 32 | if box.shape != (2*radius, 2*radius, 2*radius): 33 | breakpoint() 34 | 35 | ind = np.unravel_index(np.argmax(box), box.shape) 36 | coord[0] += ind[0] - radius + img.voxel_offset[0] 37 | coord[1] += ind[1] - radius + img.voxel_offset[1] 38 | coord[2] += ind[2] - radius + img.voxel_offset[2] 39 | 40 | syns2.pre[idx, :] = coord 41 | 42 | return syns2 -------------------------------------------------------------------------------- /chunkflow/plugins/synapse/detect_duplicate_post.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.chunk import Chunk 4 | from chunkflow.synapses import Synapses 5 | from chunkflow.lib.cartesian_coordinate import BoundingBox 6 | 7 | OUT_DIR = "/mnt/ceph/users/neuro/wasp_em/jwu/sample1/75_duplicated_post" 8 | 9 | def execute(bbox: BoundingBox, syns: Synapses, seg: Chunk, args: str = None): 10 | if syns is None: 11 | return 12 | 13 | MAX_POST_NUM = int(args) 14 | 15 | to_be_removed = syns.find_redundent_post(num_threshold=MAX_POST_NUM, distance_threshold=50.) 16 | 17 | # find the post synapses on the same neuron within a distance 18 | same_neuron_set = syns.find_duplicate_post_on_same_neuron(seg) 19 | to_be_removed = to_be_removed.union(same_neuron_set) 20 | 21 | # keep the proofread ones 22 | predicted = syns.post_indices_from_user('jwu') 23 | to_be_removed = to_be_removed.intersection(predicted) 24 | 25 | # transform to array of points 26 | to_be_removed = list(to_be_removed) 27 | posts = syns.post[to_be_removed, :] 28 | pres = syns.pre[posts[:, 0], :] 29 | # save the result 30 | # to_be_removed = np.asarray(to_be_removed, dtype=np.int32) 31 | fname = f'{OUT_DIR}/zyx_{bbox.string}.npz' 32 | posts = posts[:, 1:] 33 | np.savez(fname, pres, posts) 34 | 35 | -------------------------------------------------------------------------------- /chunkflow/plugins/synapse/detect_post_synapses.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import torch 4 | from neutorch.model.IsoRSUNet import Model 5 | 6 | import numpy as np 7 | from chunkflow.chunk import Chunk 8 | from chunkflow.chunk.probability_map import ProbabilityMap 9 | from chunkflow.synapses import Synapses 10 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 11 | 12 | # import patch_inferencer 13 | # from .patch_inferencer import PatchInferencer 14 | 15 | 16 | model_weight_file = '/mnt/ceph/users/neuro/wasp_em/jwu/14_post_synapse_net/12_net/model_570000.chkpt' 17 | #model_weight_file = '/mnt/ceph/users/neuro/wasp_em/jwu/14_post_synapse_net/32/model_1000000.chkpt' 18 | #model_weight_file = '/mnt/ceph/users/neuro/wasp_em/jwu/14_post_synapse_net/32/model_500000.chkpt' 19 | patch_size = Cartesian(192, 192, 192) 20 | is_data_parallel = True 21 | CROP_MARGIN = 28 22 | 23 | # model_weight_file = '/mnt/ceph/users/neuro/wasp_em/jwu/14_post_synapse_net/11_net/model_129000.chkpt' 24 | # patch_size = Cartesian(256, 256, 256) 25 | # is_data_parallel = False 26 | 27 | resolution = Cartesian(8, 8, 8) 28 | 29 | for ele in patch_size: 30 | assert ele == patch_size[0] 31 | 32 | 33 | class PatchInferencer: 34 | def __init__(self, 35 | model_weight_file: str, 36 | output_patch_mask: np.ndarray): 37 | 38 | if output_patch_mask is not None: 39 | output_patch_mask = torch.tensor(output_patch_mask) 40 | self.output_patch_mask = output_patch_mask 41 | 42 | model = Model(1,1) 43 | if is_data_parallel: 44 | model = torch.nn.DataParallel(model) 45 | #model.load(model_weight_file) 46 | if torch.cuda.is_available(): 47 | device = torch.device('cuda') 48 | else: 49 | device = torch.device('cpu') 50 | checkpoint = torch.load(model_weight_file, map_location=device) 51 | model.load_state_dict(checkpoint['state_dict']) 52 | model.eval() 53 | if torch.cuda.is_available(): 54 | model.cuda() 55 | if self.output_patch_mask is not None: 56 | self.output_patch_mask = self.output_patch_mask.cuda() 57 | self.model = model 58 | 59 | @property 60 | def compute_device(self): 61 | return torch.cuda.get_device_name() 62 | 63 | def __call__(self, input_patch): 64 | with torch.no_grad(): 65 | input_patch = torch.from_numpy(input_patch) 66 | if torch.cuda.is_available(): 67 | input_patch = input_patch.cuda() 68 | output_patch = self.model( input_patch ).sigmoid() 69 | if self.output_patch_mask is not None: 70 | output_patch = output_patch * self.output_patch_mask 71 | output_patch = output_patch.cpu().numpy() 72 | return output_patch 73 | 74 | inferencer = PatchInferencer(model_weight_file=model_weight_file, output_patch_mask=None) 75 | 76 | def detect_post(pre_idx: int, synapses: Synapses, img: Chunk): 77 | pre = synapses.pre[pre_idx, :] 78 | # fetch the image patch 79 | bbox = BoundingBox.from_center(pre, extent=patch_size[0]//2) 80 | np.testing.assert_array_equal(bbox.size3(), patch_size) 81 | patch = img.cutout(bbox) 82 | patch = patch.array 83 | patch = patch.astype(np.float32) 84 | patch /= 255. 85 | # inference requires 5D array 86 | patch = patch.reshape((1,1) + patch.shape) 87 | 88 | post_map = inferencer(patch) 89 | post_map = np.squeeze(post_map) 90 | # print('post map statistics: ', describe(post_map.flatten())) 91 | np.testing.assert_array_equal(post_map.shape, bbox.size3()) 92 | post_map = ProbabilityMap.from_array( 93 | post_map, bbox, 94 | voxel_size=img.voxel_size 95 | ) 96 | 97 | points, confidences = post_map.detect_points( 98 | min_distance=6, 99 | threshold_rel=0.3, 100 | exclude_border=CROP_MARGIN, 101 | num_peaks=15 102 | ) 103 | 104 | if points is None: 105 | posts = None 106 | confidences = None 107 | else: 108 | posts = [] 109 | for idx in range(points.shape[0]): 110 | post = (pre_idx, *points[idx, :]) 111 | posts.append(post) 112 | 113 | return posts, confidences 114 | 115 | def execute(img: Chunk, synapses: Synapses, aligned_bbox: BoundingBox): 116 | if synapses is None or img is None : 117 | return (None,) 118 | 119 | # this synapses should be downloaded from DVID with latest update 120 | # we also have added some unpredicted synapses in it 121 | # we only consider some synapses inside our bounding box 122 | synapses.remove_synapses_outside_bounding_box(aligned_bbox) 123 | # synapses.remove_pre_duplicates() 124 | if synapses.pre_num == 0: 125 | return(None, ) 126 | 127 | # this is for debug since we want to run it quickly 128 | # synapses.pre = synapses.pre[:10, :] 129 | 130 | # pre_index2post_indices = synapses.pre_index2post_indices 131 | posts = [] 132 | predicted_connectivities = [] 133 | connectivity_confidence = [] 134 | post_confidences = [] 135 | for pre_idx, post_indices in enumerate(synapses.pre_index2post_indices): 136 | if post_indices is not None: 137 | # skip the inference, use existing manually annotated post synapses 138 | for post_idx in post_indices: 139 | post = tuple(synapses.post[post_idx, :]) 140 | posts.append(post) 141 | 142 | if synapses.post_confidence is None: 143 | confidence = 1. 144 | else: 145 | confidence = synapses.post_confidence[post_idx] 146 | post_confidences.append(confidence) 147 | else: 148 | predicted_posts, confidences = detect_post(pre_idx, synapses, img) 149 | if predicted_posts is not None: 150 | posts.extend(predicted_posts) 151 | post_confidences.extend(confidences) 152 | connectivity_confidence.extend(confidences) 153 | 154 | for predicted_post in predicted_posts: 155 | predicted_connectivity = (*synapses.pre[pre_idx, :], *predicted_post[1:]) 156 | predicted_connectivities.append(predicted_connectivity) 157 | 158 | synapses.post = np.asarray(posts, dtype=np.int32) 159 | synapses.post_confidence = np.asarray(post_confidences, dtype=np.float32) 160 | 161 | return (synapses,) 162 | -------------------------------------------------------------------------------- /chunkflow/plugins/synapse/detect_pre_synapses.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from chunkflow.chunk import Chunk 4 | from chunkflow.chunk.probability_map import ProbabilityMap 5 | from chunkflow.synapses import Synapses 6 | from chunkflow.lib.cartesian_coordinate import BoundingBox 7 | 8 | 9 | def execute(prob: Chunk, bbox: BoundingBox): 10 | if prob is None: 11 | print('get None probability map!') 12 | return None 13 | 14 | if np.issubdtype(prob.dtype, np.uint8): 15 | prob = prob.astype(np.float32) 16 | prob /= 255. 17 | prob = ProbabilityMap.from_chunk(prob) 18 | 19 | pre, pre_confidences = prob.detect_points(min_distance=15, threshold_rel=0.3) 20 | if pre is None: 21 | syns = None 22 | else: 23 | syns = Synapses(pre, pre_confidence=pre_confidences, resolution=prob.voxel_size) 24 | syns.remove_synapses_outside_bounding_box(bbox) 25 | 26 | return syns -------------------------------------------------------------------------------- /chunkflow/plugins/synapse/find_tbar_object.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import h5py 5 | 6 | from chunkflow.synapses import Synapses 7 | from chunkflow.chunk import Chunk 8 | 9 | file_path = '/mnt/ceph/users/neuro/wasp_em/jwu/sample1/72_tbar/02_tbar_points' 10 | 11 | def execute(synapses: Synapses, seg: Chunk): 12 | tbars = synapses.tbars 13 | if len(tbars) == 0: 14 | print('tbars with 0 element: ', seg.bbox.string) 15 | return 16 | 17 | if seg.bbox.string == '0-1024_3072-4096_4096-5120': 18 | return 19 | 20 | # if seg.bbox.string == '0-1024_5120-6144_3072-4096': 21 | # breakpoint() 22 | 23 | tbars = tbars.astype(np.int64) 24 | voxel_offset = np.asarray(seg.voxel_offset, np.int64) 25 | tbars -= voxel_offset 26 | 27 | assert np.all(tbars >= 0) 28 | 29 | # eliminate the tbars 30 | mask = ( tbars[:, 0]!=1024 ) 31 | mask = np.logical_and(mask, tbars[:, 1] != 1024) 32 | mask = np.logical_and(mask, tbars[:, 2] != 1024) 33 | tbars = tbars[mask, :] 34 | 35 | objects = seg.array[tbars[:, 0], tbars[:, 1], tbars[:, 2]] 36 | assert len(objects) == tbars.shape[0] 37 | 38 | fname = os.path.join(file_path, f'{seg.bbox.string}.h5') 39 | with h5py.File(fname, mode='r+') as hf: 40 | hf['object_ids'] = objects 41 | hf['tbars_3'] = tbars + voxel_offset -------------------------------------------------------------------------------- /chunkflow/plugins/transpose.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | from chunkflow.chunk import Chunk 5 | 6 | 7 | def execute(chunk: Chunk, only_array: bool=True): 8 | chunk = chunk.transpose(only_array=only_array) 9 | return [chunk,] 10 | -------------------------------------------------------------------------------- /chunkflow/point_cloud.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | from functools import cached_property 3 | 4 | import h5py 5 | import numpy as np 6 | from .lib.cartesian_coordinate import Cartesian, BoundingBox 7 | 8 | class PointCloud: 9 | def __init__(self, points: np.ndarray, 10 | voxel_size: Cartesian) -> None: 11 | assert points.ndim == 2 12 | assert points.shape[1] == 3 13 | 14 | if not isinstance(voxel_size, Cartesian): 15 | assert len(voxel_size) == 3 16 | voxel_size = Cartesian.from_collection(voxel_size) 17 | 18 | self.points = points 19 | self.voxel_size = voxel_size 20 | 21 | @classmethod 22 | def from_h5(cls, file_path: str) -> PointCloud: 23 | assert file_path.endswith('.h5') 24 | # assert h5py.is_hdf5(file_path) 25 | with h5py.File(file_path) as hf: 26 | points = np.asarray(hf['points']) 27 | voxel_size = Cartesian.from_collection(hf['voxel_size']) 28 | return cls(points, voxel_size) 29 | 30 | def to_h5(self, file_path: str): 31 | assert file_path.endswith('.h5') 32 | with h5py.File(file_path, 'w') as hf: 33 | hf['points'] = self.points 34 | hf['voxel_size'] = self.voxel_size 35 | 36 | @classmethod 37 | def from_swc(cls, file_path: str) -> PointCloud: 38 | assert file_path.endswith('.swc') 39 | 40 | 41 | @cached_property 42 | def bounding_box(self): 43 | bbox = BoundingBox.from_points(self.points) 44 | return bbox 45 | 46 | @cached_property 47 | def point_num(self): 48 | return self.points.shape[0] -------------------------------------------------------------------------------- /chunkflow/pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "chunkflow" 3 | version = "1.1.4" 4 | description = "Composable image chunk operators to create pipeline for distributed computation." 5 | readme = "README.md" 6 | requires-python = ">=3.8" 7 | license = {file = "LICENSE.txt"} 8 | keywords = ["distributed computing", "image processing", "cloud computing", "deep learning", "connectomics"] 9 | authors = [ 10 | {email = "jingpeng.wu@gmail.com"}, 11 | {name = "Jingpeng Wu"} 12 | ] 13 | maintainers = [ 14 | {name = "Jingpeng Wu", email = "jingpeng.wu@gmail.com"} 15 | ] 16 | classifiers = [ 17 | "Development Status :: 2 - Beta", 18 | "Programming Language :: Python", 19 | 'Environment :: Console', 20 | 'Intended Audience :: End Users/Desktop', 21 | 'Intended Audience :: Developers', 22 | "License :: OSI Approved :: Apache Software License", 23 | "Operating System :: OS Independent", 24 | "Programming Language :: Python :: 3", 25 | "Programming Language :: Python :: 3.8", 26 | "Programming Language :: Python :: 3.9", 27 | "Programming Language :: Python :: 3.10", 28 | ] 29 | 30 | dependencies = [ 31 | "numpy>=1.20.0", 32 | "boto3", 33 | "cloud-volume>=0.47.0", 34 | "click>=8.1", 35 | "tqdm>=4.36.1", 36 | "connected-components-3d", 37 | "scikit-image>=0.14.2", 38 | "scikit-learn", 39 | "h5py", 40 | "neuroglancer>=2.28", 41 | "tinybrain", 42 | "tifffile", 43 | "pandas>=1.3.5", 44 | "zmesh", 45 | "fastremap", 46 | "matplotlib", 47 | "kimimaro", 48 | "pyspng-seunglab~=1.0.0", 49 | "pyparsing>=2.4.2", 50 | "urllib3>=1.25.11", 51 | "docutils>=0.10", 52 | "pynrrd", 53 | "psutil", 54 | "gputil", 55 | "traitlets", 56 | ] 57 | 58 | [project.optional-dependencies] 59 | test = [ 60 | "pytest", 61 | ] 62 | 63 | [project.urls] 64 | homepage = "https://github.com/seung-lab/chunkflow" 65 | documentation = "https://pychunkflow.readthedocs.io" 66 | repository = "https://github.com/seung-lab/chunkflow" 67 | changelog = "https://github.com/seung-lab/chunkflow/blob/master/CHANGELOG.md" 68 | 69 | [project.scripts] 70 | #chunkflow = "chunkflow.flow.flow:main" 71 | 72 | [project.gui-scripts] 73 | 74 | [project.entry-points."spam.magical"] 75 | chunkflow = "chunkflow.flow.flow:main" 76 | -------------------------------------------------------------------------------- /conftest.py: -------------------------------------------------------------------------------- 1 | pytest_plugins = ['chunkflow'] 2 | -------------------------------------------------------------------------------- /count_operators.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | printf "\n number of generators: " 3 | grep -o '@generator' chunkflow/flow/flow.py| wc -l 4 | 5 | printf "\nnumber of operators:" 6 | grep -o '@operator' chunkflow/flow/flow.py| wc -l 7 | 8 | printf "\nnumber of plugins in chunkflow: " 9 | \ls chunkflow/plugins/*.py | grep -e __init__.py --invert-match | wc -w 10 | 11 | printf "\nnumber of external plugins: " 12 | find chunkflow/plugins/chunkflow-plugins/chunkflowplugins/ -type f | wc -l 13 | -------------------------------------------------------------------------------- /dev-environments.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | neovim 3 | pynvim 4 | twine 5 | flake8 6 | -------------------------------------------------------------------------------- /distributed/kubernetes/.dockerignore: -------------------------------------------------------------------------------- 1 | .cache 2 | .git 3 | .ropeproject 4 | tests 5 | docker 6 | .env 7 | jwu 8 | dist 9 | build 10 | *.egg-info 11 | __pycache__ 12 | -------------------------------------------------------------------------------- /distributed/kubernetes/README.md: -------------------------------------------------------------------------------- 1 | # Usage of kubernetes 2 | 3 | ## [Install NVIDIA drivers](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus#installing_drivers) 4 | ``` 5 | kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/stable/nvidia-driver-installer/cos/daemonset-preloaded.yaml 6 | ``` 7 | 8 | ## build secret pod for mounting 9 | ``` 10 | kubectl create secret generic secrets \ 11 | --from-file=/secrets/google-secret.json \ 12 | --from-file=/secrets/aws-secret.json \ 13 | --from-file=/secrets/boss-secret.json 14 | ``` 15 | 16 | delete secret 17 | kubectl delete secret --all 18 | 19 | ## reconfig cluster 20 | ``` 21 | gcloud container clusters resize my-cluster --region us-central1 --node-pool gpu-pool-1 --size 1 22 | 23 | kubectl edit configmap kube-dns-autoscaler --namespace=kube-system 24 | ``` 25 | ## reconnect 26 | `gcloud container clusters get-credentials my-cluster` 27 | 28 | ## monitor and watch 29 | get the pod id 30 | kubectl get pods 31 | 32 | watch the logs 33 | watch kubectl logs pod-id 34 | 35 | We can monitor the traffic operations per second in Kubernetes deployment monitor interface. 36 | There is no bandwidth usage in Kubernetes, we can only monitor it in `instance group` monitor interface. 37 | 38 | ## deployment 39 | - create: `kubectl apply -f deploy.yml` 40 | - check: `kubectl get deployments` 41 | - delete: `kubectl delete deployment inference` 42 | - sclae: `kubectl scale --replicas=85 -f deploy.yml` 43 | -------------------------------------------------------------------------------- /distributed/kubernetes/connect_to_cluster.sh: -------------------------------------------------------------------------------- 1 | #/bin/sh 2 | gcloud container clusters get-credentials $1 --region us-central1 3 | -------------------------------------------------------------------------------- /distributed/kubernetes/deploy.yml: -------------------------------------------------------------------------------- 1 | apiVersion: apps/v1beta1 2 | kind: Deployment 3 | metadata: 4 | name: inference 5 | targets: 6 | app: chunkflow 7 | spec: 8 | replicas: 1 9 | selector: 10 | matchtargets: 11 | app: chunkflow 12 | strategy: 13 | rollingUpdate: 14 | maxSurge: 100% 15 | maxUnavailable: 100% 16 | type: RollingUpdate 17 | template: 18 | metadata: 19 | targets: 20 | app: chunkflow 21 | spec: 22 | hostNetwork: true 23 | dnsPolicy: "Default" 24 | containers: 25 | - name: chunkflow 26 | image: seunglab/chunkflow:latest 27 | imagePullPolicy: Always 28 | command: ["/bin/bash", "-c"] 29 | args: 30 | - source /root/.bashrc; 31 | mkdir /nets; cd /nets; wget -nc -nv "${CONVNET_PATH}${CONVNET_FILE_NAME}"; 32 | export PYTHONPATH=$HOME/workspace/chunkflow:$PYTHONPATH; 33 | export PYTHONPATH=$HOME/workspace/cloud-volume:$PYTHONPATH; 34 | export PYTHONPATH=$HOME/workspace/chunkflow/docker/inference/pytorch-model:$PYTHONPATH; 35 | export PYTHONPATH=$HOME/workspace/chunkflow/docker/inference/pytorch-emvision:$PYTHONPATH; 36 | echo "started all processes..."; 37 | seq "$PROCESS_NUM" | parallel -j "$PROCESS_NUM" --delay 300 --ungroup echo Starting worker {}\; chunkflow --mip 2 fetch-task --queue-name="$QUEUE_NAME" --visibility-timeout=$VISIBILITY_TIMEOUT mask --name='check-all-zero-and-skip-to-save' --check-all-zero --volume-path="$OUTPUT_MASK_volume_path" --mip 6 --skip-to='save' --fill-missing --inverse cutout --volume-path="$IMAGE_volume_path" --expand-margin-size 10 128 128 --fill-missing mask --name='mask-image' --volume-path="$IMAGE_MASK_volume_path" --mip 6 --fill-missing --inverse inference --convnet-model='rsunet' --convnet-weight-path="/nets/${CONVNET_FILE_NAME}" --patch-size 20 256 256 --patch-overlap 10 128 128 --output-key='affinity' --original-num-output-channels 3 --num-output-channels 3 --framework='pytorch-multitask' --batch-size="$BATCH_SIZE" --bump="$BUMP" crop-margin mask --name='mask-aff' --volume-path="$OUTPUT_MASK_volume_path" --mip 6 --fill-missing --inverse save --volume-path="$OUTPUT_volume_path" --upload-log --nproc 0 --create-thumbnail cloud-watch delete-task-in-queue; 38 | env: 39 | - name: PROCESS_NUM 40 | value: 2 41 | - name: BUMP 42 | value: "zung" 43 | - name: BATCH_SIZE 44 | value: "1" 45 | - name: CONVNET_PATH 46 | value: "path/of/convnet" 47 | - name: CONVNET_FILE_NAME 48 | value: "convnet.chkpt" 49 | - name: IMAGE_volume_path 50 | value: "image/layer/path" 51 | - name: IMAGE_MASK_volume_path 52 | value: "image/mask/layer/path" 53 | - name: OUTPUT_volume_path 54 | value: "output/layer/path" 55 | - name: OUTPUT_MASK_volume_path 56 | value: "output/mask/layer/path" 57 | - name: VISIBILITY_TIMEOUT 58 | value: "1800" 59 | - name: QUEUE_NAME 60 | value: "aws-sqs-queue-name" 61 | - name: LC_ALL 62 | value: C.UTF-8 63 | - name: LANG 64 | value: C.UTF-8 65 | volumeMounts: 66 | - name: secrets 67 | mountPath: /root/.cloudvolume/secrets 68 | readOnly: true 69 | resources: 70 | requests: 71 | cpu: "3.2" 72 | memory: 23.6G 73 | limits: 74 | nvidia.com/gpu: 1 75 | volumes: 76 | - name: secrets 77 | secret: 78 | secretName: secrets 79 | -------------------------------------------------------------------------------- /distributed/restapi/server.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | from typing import Union 5 | from queue import Queue 6 | 7 | from fastapi import FastAPI 8 | from pydantic import BaseModel 9 | 10 | from chunkflow.lib.cartesian_coordinate import BoundingBox 11 | 12 | app = FastAPI() 13 | 14 | maxid = 115015938 15 | 16 | @app.get("/objids/{id_num}") 17 | def get_base_id(id_num: int): 18 | global maxid 19 | base_id = maxid 20 | 21 | maxid += id_num 22 | print(f'updated maxid to {maxid}') 23 | return base_id 24 | 25 | -------------------------------------------------------------------------------- /distributed/restapi/task.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import json 4 | import math 5 | 6 | from functools import cached_property 7 | import multiprocessing as mp 8 | 9 | import psutil 10 | import GPUtil 11 | 12 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 13 | 14 | 15 | class ComputationalResource: 16 | def __init__(self, cores: int, ram: int, 17 | gpus: int = None, gpu_ram: int = None) -> None: 18 | """the required computational resources 19 | 20 | Args: 21 | cores (int): the number of CPU cores 22 | ram (int): the minimum size of RAM (GB) 23 | gpus (int): the number of GPUs. 24 | gpu_ram (int): the minimum size of GPU RAM (GB) 25 | """ 26 | self.cores = cores 27 | self.ram = ram 28 | self.gpus = gpus 29 | self.gpu_ram = gpu_ram 30 | 31 | @cached_property 32 | def is_satisfied(self, max_availability: float=1.0) -> bool: 33 | gpus = GPUtil.getGPUs( 34 | maxMemory=max_availability, maxLoad=max_availability) 35 | 36 | return self.cores <= mp.cpu_count() and \ 37 | self.ram <= psutil.virtual_memory().available // 1e9 and \ 38 | len(gpus)>=self.gpus if ( 39 | self.gpus is not None and self.gpus>0) else True 40 | 41 | @property 42 | def dict(self): 43 | return { 44 | 'cores': self.cores, 45 | 'ram': self.ram, 46 | 'gpus': self.gpus, 47 | 'gpu_ram': self.gpu_ram 48 | } 49 | 50 | 51 | class Task: 52 | def __init__(self, bbox: BoundingBox, resource: ComputationalResource = None) -> None: 53 | assert isinstance(bbox, BoundingBox) 54 | 55 | self.bbox = bbox 56 | self.resource = resource 57 | 58 | 59 | 60 | 61 | 62 | 63 | class SpatialTaskTree: 64 | def __init__(self, bbox: BoundingBox, block_size: Cartesian, parent: SpatialTaskTree = None) -> None: 65 | self.bbox = bbox 66 | self.parent = parent 67 | self.block_size = block_size 68 | 69 | # find the axis with maximum number of blocks 70 | block_num = [math.ceil(x) for x in (bbox.shape / block_size)] 71 | max_block_num = max(block_num) 72 | axis = block_num.index(max_block_num) 73 | 74 | 75 | if block_num[axis] == 1: 76 | # this is already the leaf chunk, no spliting any more 77 | # leaf chunk do not have any dependent children, so it 78 | # is naturally ready to execute 79 | self.state = 'ready' 80 | self.left = None 81 | self.right = None 82 | return 83 | else: 84 | self.state = 'not ready' 85 | left_size = (block_num[axis] // 2) * block_size[axis] 86 | left_shape = [*bbox.shape] 87 | left_shape[axis] = left_size 88 | left_bbox = BoundingBox.from_delta(bbox.start, left_shape) 89 | self.left = SpatialTaskTree(left_bbox, block_size, parent=self) 90 | 91 | right_start = [*bbox.start] 92 | right_start[axis] = left_bbox.stop[axis] 93 | right_bbox = BoundingBox(right_start, bbox.stop) 94 | self.right = SpatialTaskTree(right_bbox, block_size, parent=self) 95 | return 96 | 97 | @classmethod 98 | def from_json(cls, data: str): 99 | pass 100 | 101 | @property 102 | def json(self): 103 | if self.left is None: 104 | left = 'None' 105 | else: 106 | left = self.left.bbox.string 107 | 108 | if self.right is None: 109 | right = 'None' 110 | else: 111 | right = self.right.bbox.string 112 | 113 | if self.parent is None: 114 | parent = 'None' 115 | else: 116 | parent = self.parent.bbox.string 117 | 118 | return json.dumps({ 119 | 'state': self.state, 120 | 'bbox': self.bbox.string, 121 | 'block_size': self.block_size, 122 | 'left': left, 123 | 'right': right, 124 | 'parent': parent, 125 | }) 126 | 127 | @property 128 | def is_leaf(self): 129 | return self.left is None and self.right is None 130 | 131 | def set_state_working_on(self): 132 | self.state = 'working on' 133 | 134 | def try_to_set_parent_done(self): 135 | if self.parent.left.is_done and self.parent.right.is_done: 136 | self.parent.set_state_done() 137 | 138 | def set_state_done(self): 139 | self.state = 'done' 140 | self.try_to_set_parent_done() 141 | 142 | @property 143 | def is_done(self): 144 | return self.state == 'done' 145 | 146 | @property 147 | def leaf_list(self): 148 | """the list of leaf bounding boxes 149 | 150 | Returns: 151 | list: the bounding box list 152 | """ 153 | lst = [] 154 | if self.left.is_leaf: 155 | lst.append(self.left) 156 | else: 157 | lst.extend(self.left.leaf_list) 158 | 159 | if self.right.is_leaf: 160 | lst.append(self.right) 161 | else: 162 | lst.extend(self.right.leaf_list) 163 | return lst 164 | -------------------------------------------------------------------------------- /distributed/restapi/worker.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/distributed/restapi/worker.py -------------------------------------------------------------------------------- /docker/base/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:20.04 2 | 3 | RUN apt-get update && apt-get install -y -qq --no-install-recommends \ 4 | apt-utils \ 5 | python3 \ 6 | python3-pip \ 7 | python3-dev \ 8 | && pip3 install -U pip \ 9 | # test whether pip is working 10 | # there is an issue of pip: 11 | # https://github.com/laradock/laradock/issues/1496 12 | # we need this hash to solve this issue 13 | && hash -r pip3 \ 14 | && pip3 \ 15 | && apt-get clean \ 16 | && apt-get autoremove --purge -y \ 17 | && rm -rf /var/lib/apt/lists/* \ 18 | -------------------------------------------------------------------------------- /docker/inference/Dockerfile.pytorch: -------------------------------------------------------------------------------- 1 | #FROM pytorch/pytorch:0.4_cuda9_cudnn7 2 | #FROM pytorch/pytorch:0.4.1-cuda9-cudnn7-runtime 3 | #FROM pytorch/pytorch:1.0-cuda10.0-cudnn7-runtime 4 | FROM pytorch/pytorch:1.4-cuda10.1-cudnn7-runtime 5 | #FROM gcr.io/tpu-pytorch/xla:r0.5 6 | 7 | ENV HOME /root 8 | 9 | # make pip3 command available 10 | RUN apt update && \ 11 | apt install -y -qq --no-install-recommends \ 12 | python3-dev \ 13 | python3-pip 14 | 15 | # ln -s /opt/conda/bin/pip /usr/bin/pip3 16 | 17 | COPY pytorch-emvision $HOME/workspace/chunkflow/docker/inference/pytorch-emvision/ 18 | COPY pytorch-model $HOME/workspace/chunkflow/docker/inference/pytorch-model/ 19 | 20 | RUN echo "export PYTHONPATH=$HOME/workspace/chunkflow/docker/inference/pytorch-model:$HOME/workspace/chunkflow/docker/inference/pytorch-emvision:\$PYTHONPATH" >> $HOME/.bashrc; 21 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- 1 | # get help for all the options 2 | `make help` 3 | 4 | # build html 5 | 6 | make html 7 | 8 | open it in browser for Mac OS: 9 | 10 | open build/html/index.html 11 | 12 | # build pdf 13 | 14 | make latexpdf 15 | -------------------------------------------------------------------------------- /docs/logo/CMYK_print/Chunkflow_logo_CMYK.eps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/CMYK_print/Chunkflow_logo_CMYK.eps -------------------------------------------------------------------------------- /docs/logo/CMYK_print/Chunkflow_logo_CMYK.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/CMYK_print/Chunkflow_logo_CMYK.jpg -------------------------------------------------------------------------------- /docs/logo/CMYK_print/Chunkflow_logo_CMYK.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/CMYK_print/Chunkflow_logo_CMYK.pdf -------------------------------------------------------------------------------- /docs/logo/CMYK_print/Chunkflow_logo_CMYK.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/CMYK_print/Chunkflow_logo_CMYK.png -------------------------------------------------------------------------------- /docs/logo/RGB_web/Chunkflow_logo_RBG.eps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/RGB_web/Chunkflow_logo_RBG.eps -------------------------------------------------------------------------------- /docs/logo/RGB_web/Chunkflow_logo_RBG.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/RGB_web/Chunkflow_logo_RBG.jpg -------------------------------------------------------------------------------- /docs/logo/RGB_web/Chunkflow_logo_RBG.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/logo/RGB_web/Chunkflow_logo_RBG.png -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | sphinx<7 2 | sphinx-autodoc-typehints 3 | -------------------------------------------------------------------------------- /docs/source/_static/custom.css: -------------------------------------------------------------------------------- 1 | pre { 2 | white-space: pre-wrap; 3 | } 4 | -------------------------------------------------------------------------------- /docs/source/_static/image/cleft.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/cleft.png -------------------------------------------------------------------------------- /docs/source/_static/image/cleft_label.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/cleft_label.png -------------------------------------------------------------------------------- /docs/source/_static/image/cremi_image.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/cremi_image.png -------------------------------------------------------------------------------- /docs/source/_static/image/image_aff.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/image_aff.png -------------------------------------------------------------------------------- /docs/source/_static/image/image_seg.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/image_seg.png -------------------------------------------------------------------------------- /docs/source/_static/image/log_summary.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/log_summary.png -------------------------------------------------------------------------------- /docs/source/_static/image/operator_list.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/operator_list.png -------------------------------------------------------------------------------- /docs/source/_static/image/random_image_in_cloudvolume_viewer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/random_image_in_cloudvolume_viewer.png -------------------------------------------------------------------------------- /docs/source/_static/image/random_image_in_neuroglancer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/random_image_in_neuroglancer.png -------------------------------------------------------------------------------- /docs/source/_static/image/tasks_in_sqs.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/docs/source/_static/image/tasks_in_sqs.png -------------------------------------------------------------------------------- /docs/source/api.rst: -------------------------------------------------------------------------------- 1 | .. _api: 2 | 3 | API Reference 4 | *************** 5 | 6 | Unless otherwise stated below functions are available at top level import, 7 | i.e. after ``import chunkflow``. 8 | 9 | Chunk 10 | ======= 11 | .. autoclass:: chunkflow.chunk.Chunk 12 | :members: 13 | 14 | Image 15 | --------- 16 | .. automodule:: chunkflow.chunk.image 17 | :members: 18 | 19 | .. automodule:: chunkflow.chunk.validate 20 | :members: 21 | 22 | Flow 23 | =========== 24 | .. automodule:: chunkflow.flow.flow 25 | :members: 26 | 27 | Operators 28 | ------------ 29 | .. automodule:: chunkflow.flow 30 | :members: 31 | 32 | .. autoclass:: chunkflow.flow.cloud_watch.CloudWatchOperator 33 | :members: 34 | 35 | .. autoclass:: chunkflow.flow.load_precomputed.LoadPrecomputedOperator 36 | :members: 37 | 38 | .. autoclass:: chunkflow.flow.downsample_upload.DownsampleUploadOperator 39 | :members: 40 | 41 | .. autoclass:: chunkflow.flow.mask.MaskOperator 42 | :members: 43 | 44 | .. autoclass:: chunkflow.flow.mesh.MeshOperator 45 | :members: 46 | 47 | .. autoclass:: chunkflow.flow.neuroglancer.NeuroglancerOperator 48 | :members: 49 | 50 | .. autoclass:: chunkflow.flow.save_precomputed.SavePrecomputedOperator 51 | :members: 52 | 53 | .. autoclass:: chunkflow.flow.save_pngs.SavePNGsOperator 54 | :members: 55 | 56 | .. autoclass:: chunkflow.flow.view.ViewOperator 57 | :members: 58 | 59 | 60 | Lib 61 | ========= 62 | 63 | .. automodule:: chunkflow.lib 64 | :members: 65 | 66 | AWS 67 | ------------ 68 | .. automodule:: chunkflow.lib.aws.cloud_watch 69 | :members: 70 | 71 | .. automodule:: chunkflow.lib.aws.sqs_queue 72 | :members: 73 | -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # http://www.sphinx-doc.org/en/master/config 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | from os import path 14 | import re 15 | import sys 16 | sys.path.insert(0, path.abspath('../..')) 17 | 18 | # mock some modules 19 | from unittest.mock import MagicMock 20 | 21 | class Mock(MagicMock): 22 | @classmethod 23 | def __getattr__(cls, name): 24 | return MagicMock() 25 | 26 | MOCK_MODULES = ['waterz', 'libchunkflow'] 27 | sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES) 28 | 29 | 30 | # -- Project information ----------------------------------------------------- 31 | 32 | project = 'chunkflow' 33 | copyright = '2019, Jingpeng Wu' 34 | author = 'Jingpeng Wu' 35 | 36 | # The version info for the project you're documenting, acts as replacement for 37 | # |version| and |release|, also used in various other places throughout the 38 | # built documents. 39 | # 40 | # The short X.Y version. 41 | VERSIONFILE = path.join("../../chunkflow/__version__.py") 42 | verstrline = open(VERSIONFILE, "rt").read() 43 | VSRE = r"^__version__ = ['\"]([^'\"]*)['\"]" 44 | mo = re.search(VSRE, verstrline, re.M) 45 | if mo: 46 | verstr = mo.group(1) 47 | else: 48 | verstr = '0.0.0' 49 | 50 | # The version info for the project you're documenting, acts as replacement for 51 | # |version| and |release|, also used in various other places throughout the 52 | # built documents. 53 | # 54 | # The short X.Y version. 55 | version = verstr 56 | # The full version, including alpha/beta/rc tags. 57 | release = verstr 58 | 59 | # The full version, including alpha/beta/rc tags. 60 | #release = chunkflow.__release__ 61 | 62 | # -- General configuration --------------------------------------------------- 63 | autodoc_mock_imports = [] 64 | 65 | # Add any Sphinx extension module names here, as strings. They can be 66 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 67 | # ones. 68 | extensions = [ 69 | #'sphinx.ext.autosummary', 70 | 'sphinx.ext.autodoc', 71 | 'sphinx.ext.viewcode', # This will add links to source code to autodoc 72 | 'sphinx.ext.githubpages', 73 | #'sphinx.ext.linkcode', # This is similar to viewcode but links to external source -> need to define a function for this 74 | 'sphinx.ext.napoleon', 75 | #'sphinx.ext.mathjax', # mathjax is interactive and configurable but can also misbehave when rendering - switched to imgmath instead 76 | #'sphinx.ext.imgmath', 77 | #'matplotlib.sphinxext.plot_directive', 78 | 'sphinx_autodoc_typehints', 79 | ] 80 | 81 | # Add any paths that contain templates here, relative to this directory. 82 | templates_path = ['_templates'] 83 | 84 | # List of patterns, relative to source directory, that match files and 85 | # directories to ignore when looking for source files. 86 | # This pattern also affects html_static_path and html_extra_path. 87 | exclude_patterns = [] 88 | 89 | # -- Options for HTML output ------------------------------------------------- 90 | 91 | # The theme to use for HTML and HTML Help pages. See the documentation for 92 | # a list of builtin themes. 93 | # 94 | html_theme = 'default' 95 | 96 | # Add any paths that contain custom static files (such as style sheets) here, 97 | # relative to this directory. They are copied after the builtin static files, 98 | # so a file named "default.css" will overwrite the builtin "default.css". 99 | html_static_path = ['_static'] 100 | 101 | # These paths are either relative to html_static_path 102 | # or fully qualified paths (eg. https://...) 103 | html_css_files = ['custom.css',] 104 | 105 | master_doc = 'index' 106 | 107 | # enable the documentation of __init__ function for classes 108 | autoclass_content = 'both' 109 | 110 | # enable the typehints for Numpy style documentation 111 | napoleon_use_param = True 112 | -------------------------------------------------------------------------------- /docs/source/development.rst: -------------------------------------------------------------------------------- 1 | .. _development: 2 | 3 | Development 4 | ############ 5 | 6 | Install 7 | ********** 8 | Install with development mode (preferably in a python virtual environment):: 9 | 10 | git clone https://github.com/seung-lab/chunkflow.git 11 | cd chunkflow 12 | pip3 install -r requirements.txt 13 | python3 setup.py develop 14 | 15 | .. note:: 16 | 17 | we support python version >=3.5 18 | 19 | Release 20 | *********** 21 | 22 | #. Update version number. The version number is defined in `chunkflow/__version__.py`, increase the version number before releasing. 23 | 24 | #. Create a new release in PyPi. 25 | 26 | Our travis build/test system can automatically release/update a version if you tag the commit. It is recommended to use travis system for new version release. 27 | 28 | .. code-block:: console 29 | 30 | git tag v1.0.0 31 | git push origin v1.0.0 32 | 33 | Of course, you can also do it manually. 34 | 35 | .. code-block:: console 36 | 37 | python setup.py sdist 38 | twine upload dist/chunkflow-version.tar.gz 39 | 40 | .. note:: 41 | 42 | If you would like to include/exclude some files/folders, please edit the MANIFEST.in file. 43 | 44 | Documentation 45 | *************** 46 | We use `Sphinx`_ with `reStructuredText`_ to make documentation. You can make it locally for tests:: 47 | 48 | cd docs 49 | pip3 install -r requirements.txt 50 | make html 51 | 52 | .. _Sphinx: https://www.sphinx-doc.org 53 | .. _reStructuredText: http://docutils.sourceforge.net/docs/ref/rst/restructuredtext.html 54 | 55 | You can also make other formats, such as pdf and epub. Checkout it out with `make help`. 56 | 57 | Add a New Operator 58 | ******************* 59 | #. Take a look of the existing simple operators, such as :file:`chunkflow/flow/create_chunk.py`. 60 | 61 | #. Create your own operator file in the `flow` folder. 62 | 63 | #. Create your own operator class such as `MyOperator`, and this class should inherit from the `.base.OperatorBase`. 64 | 65 | #. Define the job of you operator in `__call__(self, chunk, *other_args, **kwargs)`. 66 | 67 | #. add the import operator code in :file:`chunkflow/flow/operators/__init__.py`. 68 | 69 | #. Add unit test in :file:`tests/flow/`. 70 | 71 | #. Run your unit test by :code:`pytest tests/flow/my_operator.py` 72 | 73 | #. Add documentation in :ref:`api`. -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | .. chunkflow documentation master file, created by 2 | sphinx-quickstart on Thu Jul 25 17:29:21 2019. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Welcome to chunkflow's documentation! 7 | ###################################### 8 | 9 | .. toctree:: 10 | :maxdepth: 2 11 | 12 | introduction 13 | install 14 | tutorial 15 | api 16 | development 17 | 18 | Indices and tables 19 | ################### 20 | 21 | * :ref:`genindex` 22 | * :ref:`modindex` 23 | * :ref:`search` 24 | -------------------------------------------------------------------------------- /docs/source/install.rst: -------------------------------------------------------------------------------- 1 | .. _install: 2 | 3 | Install 4 | ######## 5 | Install from Pypi 6 | ================== 7 | .. note:: 8 | 9 | We support python version >=3.6 in Ubuntu >= 16.04. It should work with other Linux version too. If you need support of other python version or OS, please create an issue in github. 10 | 11 | Install the version released in pypi:: 12 | 13 | pip3 install chunkflow 14 | 15 | .. note:: 16 | 17 | You need g++>=4.8 since some dependent packages need to compile c++ code. 18 | 19 | Manual Installation 20 | =================== 21 | Install the latest from repo:: 22 | 23 | git clone https://github.com/seung-lab/chunkflow.git 24 | cd chunkflow 25 | pip install -r requirements.txt 26 | python setup.py install 27 | 28 | .. note:: 29 | 30 | You need g++>=4.8 since some dependent packages need to compile c++ code. 31 | 32 | Use Docker Image without Installation 33 | ====================================== 34 | We have built `docker 35 | `_ images ready to use. You can use docker images directly without installation and configuration. The docker images are in `Docker Hub 36 | `_. 37 | 38 | You can get the image by: 39 | 40 | .. code-block:: 41 | 42 | docker pull seunglab/chunkflow 43 | 44 | 45 | -------------------------------------------------------------------------------- /docs/source/introduction.rst: -------------------------------------------------------------------------------- 1 | .. _introduction: 2 | 3 | Introduction 4 | ############# 5 | 6 | Problem 7 | =========== 8 | Benefited from the rapid development of microscopy technologies, we can acquire large scale 3D volumetric datasets with both high resolution and large field of view. These 3D image datasets are too big to be processed in a single computer, and distributed processing is required. In most cases, the image dataset could be choped to chunks with/without overlap and distributed to computers for processing. 9 | 10 | Inside each single computer, we perform some operations of the image chunk. The type or order of operations varies according to the image type, quality and application, it is hard to make a universal pipeline for general usage or create specialized pipeline for each usage case. Composing and reusing operators to create customized pipeline easily will facilitate usage. 11 | 12 | Solution 13 | ========= 14 | 15 | Chunkflow provides a framework to perform distributed chunk processing for large scale 3D image dataset. For each task in a single computer, you can compose operators to create pipeline easily for each use case. 16 | 17 | Terminology 18 | ----------- 19 | - patch: ndarray as input to ConvNet. Normally it is pretty small due to the limited memory capacity of GPU. 20 | - chunk: ndarray with global offset and arbitrary shape. 21 | - block: the array with a shape and global offset aligned with storage backend. The block could be saved directly to storage backend. The alignment with storage files ensures that there is no writting conflict when saved parallelly. 22 | 23 | -------------------------------------------------------------------------------- /examples/downsample/create_tasks.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # downsample from mip 0 to mip 5 4 | chunkflow generate-tasks -l gs://my/dataset/path -m 0 -o 0 0 0 -c 112 2048 2048 -q my-queue-name 5 | 6 | # downsample from mip 5 to mip 9 7 | #chunkflow generate-tasks -l gs://my/dataset/path -m 5 -o 0 0 0 -c 112 2048 2048 -q my-queue-name 8 | -------------------------------------------------------------------------------- /examples/downsample/downsample.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export PROCESS_NUM=16 3 | export MIP=0 4 | export STOP_MIP=6 5 | export DATASET_PATH="gs://my/dataset/path" 6 | 7 | # perform downsample 8 | seq "$PROCESS_NUM" | parallel -j "$PROCESS_NUM" --delay 1 --ungroup echo Starting worker {}\; chunkflow --quiet --mip "$MIP" fetch-task -q my-queue-name -v 60 cutout -v "$DATASET_PATH" --fill-missing downsample-upload -v "$DATASET_PATH" --stop-mip "$STOP_MIP" delete-task-in-queue 9 | -------------------------------------------------------------------------------- /examples/inference/01_download.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | wget -nc -nv https://github.com/seung-lab/DeepEM/releases/download/S1/ac3ac4_model790000.chkpt 4 | -------------------------------------------------------------------------------- /examples/inference/batchnorm3d_to_instancenorm3d.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import torch 5 | from torch import nn 6 | from copy import deepcopy 7 | 8 | def batchnorm3d_to_instancenorm3d(model): 9 | conversion_count = 0 10 | for name, module in reversed(model._modules.items()): 11 | if len(list(module.children())) > 0: 12 | # recurse 13 | model._modules[name], num_converted = batchnorm3d_to_instancenorm3d(module) 14 | conversion_count += num_converted 15 | 16 | if isinstance(module, nn.BatchNorm3d): 17 | layer_new = nn.InstanceNorm3d(module.num_features, eps=module.eps, 18 | momentum=module.momentum, 19 | affine=module.affine, 20 | track_running_stats=False) 21 | layer_new.weight=module.weight 22 | layer_new.bias = module.bias 23 | model._modules[name] = layer_new 24 | conversion_count += 1 25 | 26 | return model, conversion_count 27 | 28 | if __name__ == '__main__': 29 | model = torch.identity() 30 | model2 = deepcopy(model).cpu() 31 | new_model, count = batchnorm3d_to_instancenorm3d(model2, nn.BatchNorm3d, nn.InstanceNorm3d) 32 | -------------------------------------------------------------------------------- /examples/inference/inference.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export QUEUE_NAME="chunkflow" 4 | export VISIBILITY_TIMEOUT="3600" 5 | export IMAGE_LAYER_PATH="gs://my/image/path" 6 | export NORMALIZED_IMAGE_LAYER_PATH="gs://my/normalized/image/path" 7 | export IMAGE_MASK_LAYER_PATH="gs://my/image/mask/path" 8 | export AFF_CONVNET_MODEL_FILE="./aff_net.py" 9 | export AFF_CONVNET_WEIGHT_FILE="./aff_net.chkpt" 10 | export SEM_CONVNET_MODEL_FILE="./sem_net.py" 11 | export SEM_CONVNET_WEIGHT_FILE="./sem_net.chkpt" 12 | export AFF_LAYER_PATH="gs://my/affinity/map/path" 13 | export AFF_MASK_LAYER_PATH="gs://my/affinity/map/path" 14 | export SEM_LAYER_PATH="gs://my/semantic/map/path" 15 | 16 | export IMAGE_HISTOGRAM_PATH="gs://my/image/histogram/path" 17 | 18 | # perform boundary detection and semantic map inference together 19 | #chunkflow --verbose 2 --mip 1 \ 20 | # fetch-task -r 20 --queue-name="$QUEUE_NAME" --visibility-timeout=$VISIBILITY_TIMEOUT \ 21 | # cutout --volume-path="$IMAGE_LAYER_PATH" --expand-margin-size 4 64 64 --fill-missing -o "image" \ 22 | # normalize-section-contrast -i "image" -o "image" -p $IMAGE_HISTOGRAM_PATH -l 0.01 -u 0.01 \ 23 | # crop-margin -i "image" -o "cropped-image" -m 4 64 64 \ 24 | # save --name "save-img" -i "cropped-image" --volume-path="$NORMALIZED_IMAGE_LAYER_PATH" --nproc 1 \ 25 | # delete-chunk --name "delete-cropped-image" -c "cropped-image" \ 26 | # inference --name "sem-inference" -i "image" --dtype float32 --convnet-model="$SEM_CONVNET_MODEL_FILE" --convnet-weight-path="${SEM_CONVNET_WEIGHT_FILE}" \ 27 | # --input-patch-size 20 256 256 --output-patch-overlap 4 64 64 --num-output-channels 5 --framework='pytorch' --batch-size 1 --patch-num 7 11 11 \ 28 | # channel-voting \ 29 | # save --name "save-sem" --volume-path="$SEM_LAYER_PATH" --nproc 1 \ 30 | # delete-chunk --name "delete-sem" -c "chunk" \ 31 | # mask --name='mask-image' -i "image" -o "image" --volume-path="$IMAGE_MASK_LAYER_PATH" --mip 4 --fill-missing --inverse \ 32 | # inference --name "aff-inference" -i "image" --mask-myelin-threshold 0.3 --dtype float32 --convnet-model="$AFF_CONVNET_MODEL_FILE" \ 33 | # --convnet-weight-path="${AFF_CONVNET_WEIGHT_FILE}" --input-patch-size 20 256 256 --output-patch-size 20 256 256 \ 34 | # --output-patch-overlap 4 64 64 --output-crop-margin 4 64 64 --num-output-channels 4 --framework='pytorch' --batch-size 1 --patch-num 7 11 11 \ 35 | # delete-chunk --name "delete-img" -c "image" \ 36 | # mask --name='mask-aff' --volume-path="$AFF_MASK_LAYER_PATH" --mip 4 --fill-missing --inverse \ 37 | # save --name "save-aff" --volume-path="$AFF_LAYER_PATH" --upload-log --nproc 1 --create-thumbnail \ 38 | # delete-chunk --name "delete-aff" -c "chunk" \ 39 | # cloud-watch 40 | # # delete-task-in-queue 41 | 42 | # perform boundary detection only 43 | chunkflow --verbose 2 --mip 0 \ 44 | generate-tasks -m 0 --roi-start 17850 165000 175000 \ 45 | --chunk-size 100 1024 1024 -g 1 1 1 \ 46 | cutout --volume-path="$IMAGE_LAYER_PATH" --expand-margin-size 4 64 64 \ 47 | --fill-missing -o "image" \ 48 | normalize-section-contrast -i "image" -o "image" -p $IMAGE_HISTOGRAM_PATH \ 49 | -l 0.01 -u 0.01 \ 50 | inference --name "aff-inference" -i "image" --mask-myelin-threshold 0.3 \ 51 | --dtype float32 --convnet-model="$AFF_CONVNET_MODEL_FILE" \ 52 | --convnet-weight-path="${AFF_CONVNET_WEIGHT_FILE}" \ 53 | --input-patch-size 20 256 256 --output-patch-size 20 256 256 \ 54 | --output-patch-overlap 4 64 64 --output-crop-margin 4 64 64 \ 55 | --num-output-channels 4 --framework='pytorch' --batch-size 1 \ 56 | --mask-output-chunk \ 57 | delete-chunk --name "delete-img" -c "image" \ 58 | write-h5 -f "affinitymap_175000_165000_17850.h5" 59 | -------------------------------------------------------------------------------- /examples/inference/universal_identity.py: -------------------------------------------------------------------------------- 1 | import platform 2 | 3 | import numpy as np 4 | 5 | 6 | class PatchInferencer: 7 | """ 8 | IdentityPatchInferenceEngine(PatchInferenceEngine) 9 | 10 | return the same output with the input 11 | this class was only used for tests 12 | Note that the bump parameter is not used now, only support bump function 13 | """ 14 | def __init__(self, model_weight_file, output_patch_mask): 15 | self.output_patch_mask = output_patch_mask 16 | 17 | @property 18 | def compute_device(self): 19 | return platform.processor() 20 | 21 | def __call__(self, input_patch): 22 | """ 23 | return the same with argument 24 | reshape the size to 5 dimension: 25 | batch, channel, z, y, x 26 | """ 27 | if np.issubdtype(input_patch.dtype, np.integer): 28 | # normalize to 0-1 value range 29 | input_patch /= np.iinfo(input_patch.dtype).max 30 | output_patch = input_patch.astype(np.float32) 31 | 32 | # mask should be done in patch engine now 33 | output_patch *= self.output_patch_mask 34 | 35 | return output_patch 36 | -------------------------------------------------------------------------------- /examples/inference/universal_pytorch.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import torch 4 | from pytorch_model.model import load_model 5 | from types import SimpleNamespace 6 | 7 | 8 | class PatchInferencer: 9 | def __init__(self, model_weight_file, output_patch_mask): 10 | self.output_patch_mask = torch.tensor(output_patch_mask) 11 | 12 | patch_size = (20, 256, 256) 13 | patch_overlap = (4, 64, 64) 14 | output_key = 'affinity' 15 | num_output_channels = 3 16 | 17 | d = dict() 18 | d['model'] = 'rsunet' 19 | d['width'] = [16, 32, 64, 128] 20 | d['in_spec'] = {'input': (1, *patch_size)} 21 | d['out_spec'] = {output_key: (16, *patch_size)} 22 | d['scan_spec'] = {output_key: (num_output_channels, *patch_size)} 23 | d['pretrain'] = True 24 | d['precomputed'] = False 25 | d['edges'] = [(0, 0, 1), (0, 1, 0), (1, 0, 0)] 26 | d['overlap'] = tuple(patch_overlap) 27 | d['bump'] = None 28 | d['cropsz'] = None 29 | d['gpu_ids'] = None 30 | 31 | self.opt = SimpleNamespace(**d) 32 | self.net = load_model(self.opt, model_weight_file) 33 | if torch.cuda.is_available(): 34 | self.net.cuda() 35 | self.output_patch_mask = self.output_patch_mask.cuda() 36 | assert len(self.opt.in_spec) == 1 37 | 38 | @property 39 | def compute_device(self): 40 | return torch.cuda.get_device_name() 41 | 42 | def __call__(self, input_patch): 43 | with torch.no_grad(): 44 | input_patch = torch.from_numpy(input_patch).cuda() 45 | output_patch = self.net( input_patch ).sigmoid() 46 | output_patch = output_patch * self.output_patch_mask 47 | output_patch = output_patch.cpu().numpy() 48 | return output_patch 49 | -------------------------------------------------------------------------------- /examples/inference/universal_pznet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import sys 3 | import platform 4 | 5 | 6 | class PatchInferencer: 7 | def __init__(self, model_weight_file, output_patch_mask): 8 | self.output_patch_mask = output_patch_mask 9 | 10 | sys.path.append(model_weight_file) 11 | from pznet.pznet import PZNet 12 | 13 | self.net = PZNet(model_weight_file) 14 | 15 | @property 16 | def compute_device(self): 17 | return platform.processor() 18 | 19 | def __call__(self, input_patch): 20 | """ 21 | args: 22 | patch (5d numpy array): input patch with dimensions \ 23 | batch/channel/z/y/x 24 | return: 25 | 5d numpy array with the same dimension arrangment. 26 | """ 27 | # make sure that the input patch is 5d ndarray 28 | output_patch = self.net.forward(input_patch) 29 | output_patch *= self.output_patch_mask 30 | return output_patch 31 | -------------------------------------------------------------------------------- /examples/mesh/create_tasks.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | chunkflow generate-tasks -l gs://my/dataset/path -m 3 -o 1 1 1 -c 226 2050 2050 -q my-queue 4 | -------------------------------------------------------------------------------- /examples/mesh/manifest.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export PROCESS_NUM=10 3 | 4 | 5 | seq "$PROCESS_NUM" | parallel -j "$PROCESS_NUM" --delay 20 --ungroup echo Starting worker {}\; chunkflow mesh-manifest -v gs://my/dataset/path --prefix {} 6 | -------------------------------------------------------------------------------- /examples/mesh/mesh.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export PROCESS_NUM=20 3 | export DATASET_PATH="gs://my/dataset/path" 4 | 5 | seq "$PROCESS_NUM" | parallel -j "$PROCESS_NUM" --delay 120 --ungroup echo Starting worker {}\; chunkflow --mip 3 fetch-task -q my-queue -v 1200 cutout -v "$DATASET_PATH" --fill-missing mesh -o "$DATASET_PATH" --dust-threshold 100 --ids 76181,76182 delete-task-in-queue 6 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy>=1.20.0 2 | boto3 3 | cloud-volume>=0.47.0 4 | click>=8.1 5 | tqdm>=4.36.1 6 | connected-components-3d 7 | scikit-image>=0.14.2 8 | scikit-learn 9 | h5py 10 | neuroglancer>=2.28 11 | tinybrain 12 | tifffile 13 | pandas>=1.3.5 14 | zmesh 15 | fastremap 16 | pyparsing>=2.4.2 17 | urllib3>=1.25.11 18 | docutils>=0.10 19 | pynrrd 20 | psutil 21 | gputil 22 | traitlets 23 | zarr 24 | scipy 25 | microviewer 26 | pyspng-seunglab~=1.0.0 27 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from setuptools import setup, find_packages 4 | 5 | PACKAGE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | 7 | with open(os.path.join(PACKAGE_DIR, 'requirements.txt')) as f: 8 | install_requires = f.read().splitlines() 9 | install_requires = [l for l in install_requires if not l.startswith('#')] 10 | 11 | tests_requirements_file = os.path.join(PACKAGE_DIR, 'tests/requirements.txt') 12 | # in the release tar.gz file, there will be no tests_requirements_file 13 | if os.path.isfile(tests_requirements_file): 14 | with open(tests_requirements_file) as f: 15 | tests_require = f.read().splitlines() 16 | tests_require = [l for l in tests_require if not l.startswith('#')] 17 | else: 18 | tests_require = [] 19 | 20 | with open("README.md", "r") as fh: 21 | LONG_DESCRIPTION = fh.read() 22 | 23 | VERSIONFILE = os.path.join(PACKAGE_DIR, "chunkflow/__version__.py") 24 | verstrline = open(VERSIONFILE, "rt").read() 25 | VSRE = r"^version = ['\"]([^'\"]*)['\"]" 26 | mo = re.search(VSRE, verstrline, re.M) 27 | if mo: 28 | version = mo.group(1) 29 | else: 30 | raise RuntimeError("Unable to find version string in %s." % 31 | (VERSIONFILE, )) 32 | 33 | 34 | setup( 35 | name='chunkflow', 36 | description='Composable image chunk operators to create pipeline' + 37 | ' for distributed computation.', 38 | long_description=LONG_DESCRIPTION, 39 | long_description_content_type="text/markdown", 40 | license='Apache License 2.0', 41 | version=version, 42 | author='Jingpeng Wu', 43 | author_email='jingpeng.wu@gmail.com', 44 | packages=find_packages(exclude=['bin', 'docker', 'kubernetes']), 45 | url='https://github.com/seung-lab/chunkflow', 46 | install_requires=install_requires, 47 | tests_require=tests_require, 48 | entry_points=''' 49 | [console_scripts] 50 | chunkflow=chunkflow.flow.flow:main 51 | ''', 52 | classifiers=[ 53 | 'Environment :: Console', 54 | 'Intended Audience :: End Users/Desktop', 55 | 'Intended Audience :: Developers', 56 | "License :: OSI Approved :: Apache Software License", 57 | "Operating System :: OS Independent", 58 | "Programming Language :: Python :: 3", 59 | "Programming Language :: Python :: 3.8", 60 | "Programming Language :: Python :: 3.9", 61 | "Programming Language :: Python :: 3.10", 62 | ], 63 | python_requires='>=3', 64 | zip_safe=False, 65 | ) 66 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | # Using Green Threads 5 | # # import gevent.monkey 6 | # gevent.monkey.patch_all(thread=False) 7 | -------------------------------------------------------------------------------- /tests/chunk/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/chunk/__init__.py -------------------------------------------------------------------------------- /tests/chunk/affinity_map/test_affinity_map.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import numpy as np 5 | 6 | from chunkflow.chunk.affinity_map import AffinityMap 7 | 8 | def test_affinity_map_construction(): 9 | arr = np.random.rand(3,3,4,5).astype(np.float32) 10 | aff = AffinityMap(arr, voxel_offset=(0, -1,-1,-1)) 11 | 12 | -------------------------------------------------------------------------------- /tests/chunk/image/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/chunk/image/__init__.py -------------------------------------------------------------------------------- /tests/chunk/image/test_adjust_grey.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /tests/chunk/image/test_image.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import numpy as np 5 | from chunkflow.chunk.image import Image 6 | 7 | def test_image(): 8 | arr = np.random.randint(256, size=(3,4,5), dtype=np.uint8) 9 | image = Image(arr, voxel_offset=(-1,-1,-1)) 10 | 11 | # test normalization 12 | image.normalize_contrast() 13 | -------------------------------------------------------------------------------- /tests/chunk/segmentation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/chunk/segmentation/__init__.py -------------------------------------------------------------------------------- /tests/chunk/segmentation/test_segmentation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from chunkflow.chunk.segmentation import Segmentation 3 | from math import isclose 4 | 5 | 6 | def test_segmentation(): 7 | arr = np.random.randint(256, size=(30,40,50), dtype=np.uint32) 8 | seg = Segmentation(arr, voxel_offset=(-1,-1,-1)) 9 | 10 | arr2 = arr.astype(np.uint64) 11 | seg2 = Segmentation(arr2, voxel_offset=(-1,-1,-1)) 12 | 13 | 14 | print('compare two segments.') 15 | scores = seg.evaluate(seg2) 16 | print('evaluate scores: \n', scores) 17 | assert scores['variation_of_information'] == 0 18 | assert scores['rand_index'] == 1 19 | assert isclose(scores['adjusted_rand_index'], 1.) 20 | assert isclose(scores['fowlkes_mallows_index'], 1.) 21 | -------------------------------------------------------------------------------- /tests/chunk/test_validate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import unittest 3 | 4 | from chunkflow.chunk.validate import validate_by_template_matching 5 | 6 | 7 | class TestValidateByTemplateMatching(unittest.TestCase): 8 | def test_validate_by_template_matching(self): 9 | print('test validate by template matching...') 10 | image = np.random.randint(0, 256, size=(64, 64, 64), dtype=np.uint8) 11 | assert validate_by_template_matching(image) 12 | 13 | # make a black box 14 | image[16:-16, 16:-16, 16:-16] = 0 15 | assert not validate_by_template_matching(image) 16 | -------------------------------------------------------------------------------- /tests/command_line.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | echo -e "\ngenerate tasks..." 4 | chunkflow generate-tasks -c 0 0 0 -s 0 0 0 -g 1 1 1 5 | 6 | echo -e "\ndownsample image..." 7 | chunkflow create-chunk --dtype uint8 --pattern sin downsample 8 | echo -e "\ndownsample segmentation..." 9 | chunkflow create-chunk --dtype uint32 --pattern sin downsample 10 | # we do not support float32 yet, but will be add later. 11 | # chunkflow create-chunk --dtype float32 --pattern sin downsample 12 | 13 | 14 | # this example is from [nuclease](https://github.com/janelia-flyem/neuclease) 15 | # echo -e "test nuclease..." 16 | # chunkflow generate-tasks --roi-start 20789 21341 17019 --chunk-size 16 128 128 plugin -f cutout_dvid_label -i bbox -o chunk 17 | 18 | #echo -e "\nread chunks using tensorstore." 19 | #chunkflow \ 20 | # generate-tasks --roi-start 20000 20000 20000 --chunk-size 32 32 32 \ 21 | # load-tensorstore -s gs://neuroglancer-janelia-flyem-hemibrain/v1.1/segmentation/ --voxel-size 8 8 8 \ 22 | 23 | echo -e "\ncreate a hdf5 file, then test the skip-task operator." 24 | chunkflow create-chunk --pattern zero skip-all-zero save-h5 -f /tmp/ 25 | # the file already exist, so we'll skip that task 26 | # if the skip is not successful, there is no chunk to write, we'll get an error. 27 | chunkflow create-chunk save-h5 -f /tmp/ 28 | 29 | echo -e "\ntest skip tasks..." 30 | chunkflow generate-tasks --roi-start 0 0 0 --chunk-size 64 64 64 skip-task --prefix /tmp/ --suffix .h5 31 | rm /tmp/0-64_0-64_0-64.h5 32 | 33 | echo -e "\ntest plugin with arguments..." 34 | chunkflow create-chunk plugin -f median_filter -i chunk -o chunk --args "size=(3,1,1);mode=reflect" 35 | 36 | echo -e "\nevaluate segmentation quality..." 37 | chunkflow create-chunk -o seg create-chunk -o gt evaluate-segmentation -s seg -g gt 38 | 39 | echo -e "\ngenerate summary of log..." 40 | chunkflow log-summary -l tests/data/log 41 | 42 | echo -e "\ntest write png files..." 43 | chunkflow create-chunk save-pngs -o /tmp/pngs; rm -rf /tmp/pngs 44 | 45 | echo -e "\ntest normalize intensity..." 46 | chunkflow create-chunk normalize-intensity 47 | 48 | echo -e "\nconnected component analysis..." 49 | chunkflow create-chunk save-h5 --file-name="/tmp/img.h5" connected-components --threshold 128 save-h5 --file-name=/tmp/seg.h5 50 | if test -f /tmp/img.h5 ; then echo -e "File found"; else exit 1; fi 51 | chunkflow load-h5 --file-name=/tmp/img.h5 52 | 53 | #echo -e "\nwrite and read nrrd file." 54 | #chunkflow create-chunk save-nrrd -f "/tmp/img.nrrd"; 55 | #chunkflow load-nrrd -f "/tmp/img.nrrd"; 56 | #rm /tmp/img.nrrd; 57 | 58 | echo -e "\nwrite and read tif file." 59 | chunkflow create-chunk save-tif -f "/tmp/img.tif"; 60 | chunkflow load-tif -f "/tmp/img.tif"; 61 | rm /tmp/img.tif; 62 | chunkflow create-chunk --dtype uint16 --pattern sin save-tif -f /tmp/seg.tif 63 | chunkflow load-tif -f /tmp/seg.tif 64 | rm /tmp/seg.tif 65 | 66 | echo -e "\ncreate the info file of Neuroglancer Precomputed format." 67 | mkdir /tmp/seg 68 | chunkflow \ 69 | create-chunk --size 128 128 128 --dtype uint32 --pattern sin \ 70 | create-info --voxel-size 8 8 8 --block-size 64 64 64 --output-layer-path file:///tmp/seg/ 71 | 72 | echo -e "\nwrite image to precomputed volume." 73 | # somehow, we have to separate creation of info and writing out precomputed operation! 74 | chunkflow \ 75 | create-chunk --size 128 128 128 --dtype uint32 --pattern sin \ 76 | save-precomputed --volume-path file:///tmp/seg \ 77 | mesh --voxel-size 8 8 8 --output-format precomputed --output-path file:///tmp/seg 78 | rm -rf /tmp/seg 79 | 80 | echo -e "\ndry run of setup-environment in the cloud." 81 | chunkflow --dry-run --log-level info \ 82 | setup-env -l "gs://my/path" --volume-start 2002 25616 12304 \ 83 | --volume-stop 2068 26128 12816 --max-ram-size 14 \ 84 | --input-patch-size 20 128 128 \ 85 | --output-patch-size 16 96 96 --output-patch-overlap 6 32 32 \ 86 | --channel-num 3 --dtype float32 -m 0 --encoding raw \ 87 | --voxel-size 45 16 16 --max-mip 5 88 | 89 | echo -e "\nconvolutional net inference." 90 | chunkflow --log-level debug \ 91 | create-chunk --size 36 448 448 \ 92 | inference --input-patch-size 20 256 256 --patch-num 2 2 2 \ 93 | --framework "universal" \ 94 | --convnet-model "./examples/inference/universal_identity.py" \ 95 | --batch-size 3 #cloud-watch --log-name chunkflow-test 96 | 97 | chunkflow \ 98 | create-chunk --size 36 448 448 \ 99 | inference --input-patch-size 20 256 256 \ 100 | --patch-num 2 2 2 --framework identity --batch-size 3 101 | 102 | chunkflow create-chunk --pattern zero --size 36 448 448 \ 103 | inference --input-patch-size 20 256 256 --patch-num 2 2 2 \ 104 | --framework identity --batch-size 3 105 | 106 | #echo -e "\nmask out objects and skeletonize." 107 | #chunkflow \ 108 | # create-chunk --size 36 448 448 --dtype "uint32" \ 109 | # connected-components mask-out-objects -d 50 -s "2,3,4"\ 110 | # skeletonize --voxel-size 1 1 1 --output-path "file:///tmp/test/skeleton" 111 | -------------------------------------------------------------------------------- /tests/data/levels/1/26774: -------------------------------------------------------------------------------- 1 | {"levels": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 4, 29, 73, 169, 307, 561, 957, 1426, 2157, 2876, 3531, 4138, 4852, 5521, 5700, 5670, 5787, 5796, 5911, 6031, 6602, 7206, 7595, 8114, 8840, 9376, 10113, 10727, 11268, 11743, 12581, 13112, 13480, 14229, 14656, 15046, 15353, 15559, 16135, 16362, 16451, 16725, 16802, 16969, 16780, 16880, 16799, 17040, 17239, 17927, 19292, 21522, 25054, 31882, 42071, 56831, 80877, 116645, 168336, 240698, 342168, 477998, 653351, 884384, 1193180, 1618320, 2219800, 3074908, 4279748, 5941999, 8112931, 10814858, 13975996, 17410623, 20886074, 24177223, 27085345, 29507047, 31431476, 32911423, 34045615, 34886045, 35480405, 35827734, 35928221, 35795168, 35386690, 34720617, 33880232, 32954614, 32029977, 31102770, 29870584, 28152115, 25819320, 23029243, 19976930, 16806513, 13622274, 10539506, 7713779, 5317132, 3465759, 2141745, 1263925, 714780, 390816, 208012, 107627, 55262, 27765, 14312, 7363, 4060, 2341, 1369, 885, 667, 451, 362, 261, 192, 159, 126, 81, 63, 63, 34, 29, 24, 9, 9, 3, 3, 1, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "patch_size": [2048, 2048, 1], "num_patches": 202, "coverage_ratio": 0.09935684165680474} -------------------------------------------------------------------------------- /tests/data/levels/normalized.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/data/levels/normalized.npz -------------------------------------------------------------------------------- /tests/data/log/0-3_16384-16492_86294-88342_121142-123190.json: -------------------------------------------------------------------------------- 1 | {"timer": {"cutout": 69.0063271522522, "normalize-section-contrast": 21.265912294387817, "sem-inference": 535.2609465122223, "save-sem": 7.987687587738037, "mask-image": 4.825762033462524, "aff-inference": 273.0924451351166, "mask-aff": 13.666078090667725, "save-aff": 71.74230360984802}, "bbox": "16384-16492_86294-88342_121142-123190", "compute_device": "TITAN X (Pascal)"} -------------------------------------------------------------------------------- /tests/data/log/0-3_16384-16492_88342-90390_110902-112950.json: -------------------------------------------------------------------------------- 1 | {"timer": {"cutout": 76.88676476478577, "normalize-section-contrast": 2.3699328899383545, "sem-inference": 537.3570463657379, "save-sem": 8.875961303710938, "mask-image": 5.057928085327148, "aff-inference": 271.429256439209, "mask-aff": 3.7008891105651855, "save-aff": 89.11688446998596}, "bbox": "16384-16492_88342-90390_110902-112950", "compute_device": "TITAN X (Pascal)"} -------------------------------------------------------------------------------- /tests/flow/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/flow/__init__.py -------------------------------------------------------------------------------- /tests/flow/divid_conquer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/flow/divid_conquer/__init__.py -------------------------------------------------------------------------------- /tests/flow/divid_conquer/test_patch_mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from scipy.stats import describe 4 | 5 | from chunkflow.flow.divid_conquer.patch.patch_mask \ 6 | import make_bump_map, PatchMask 7 | 8 | 9 | def test_patch_mask(): 10 | patch_size = (20, 256, 256) 11 | bump_map = make_bump_map(patch_size) 12 | 13 | print('describe bump map:') 14 | print(describe(bump_map, axis=None)) 15 | 16 | patch_mask = PatchMask(patch_size, (4, 64, 64)) 17 | print('shape of mask: {}'.format(patch_mask.shape)) 18 | print('describe patch mask:') 19 | print(describe(patch_mask, axis=None)) 20 | 21 | #import os 22 | #import h5py 23 | #file_name = '/tmp/patch_mask.h5' 24 | #if os.path.exists(file_name): 25 | # os.remove(file_name) 26 | #with h5py.File(file_name, 'w') as f: 27 | # f['/main'] = patch_mask 28 | -------------------------------------------------------------------------------- /tests/flow/divid_conquer/visualize_transform.py: -------------------------------------------------------------------------------- 1 | #%% 2 | from itertools import product 3 | 4 | # import numpy as np 5 | import matplotlib.pyplot as plt 6 | 7 | from chunkflow.chunk import Chunk 8 | from chunkflow.flow.divid_conquer.transform import TransformSequences 9 | 10 | # img = np.zeros(shape=(4, 128, 128), dtype=np.uint8) 11 | # img[:, 60:68, :] = 255 12 | # img = Chunk.create() 13 | img = Chunk.from_h5('/Users/jwu/dropbox/40_gt/13_wasp_sample3/vol_03700/img_zyx_3600-4056_4900-5356_4150-4606.h5') 14 | # img = img.array 15 | 16 | transform_sequences = TransformSequences() 17 | 18 | transformed_images = transform_sequences.forward(img) 19 | inversed_images = transform_sequences.backward(transformed_images) 20 | 21 | 22 | fig, axs = plt.subplots(4, 4, figsize=(15,15)) 23 | 24 | for x, y in product(range(4), range(4)): 25 | idx = y*4 + x 26 | arr = transformed_images[idx] 27 | axs[x, y].imshow(arr[0,...], cmap='gray') 28 | 29 | 30 | fig.show() 31 | 32 | 33 | # %% 34 | -------------------------------------------------------------------------------- /tests/flow/test_downsample_upload.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | import tempfile 3 | import numpy as np 4 | from chunkflow.chunk import Chunk 5 | 6 | from cloudvolume import CloudVolume 7 | 8 | from chunkflow.flow.downsample_upload import DownsampleUploadOperator 9 | 10 | 11 | def hierarchical_downsample(chunk, layer_type='segmentation'): 12 | # save the input to disk 13 | tempdir = tempfile.mkdtemp() 14 | volume_path = 'file://' + tempdir 15 | CloudVolume.from_numpy(chunk.transpose(), 16 | vol_path=volume_path, 17 | voxel_offset=(32, 32, 2), 18 | chunk_size=(32, 32, 4), 19 | max_mip=4, 20 | layer_type=layer_type) 21 | 22 | operator = DownsampleUploadOperator( 23 | volume_path, 24 | factor=(1, 2, 2), 25 | chunk_mip=0, 26 | start_mip=1, 27 | stop_mip=4) 28 | 29 | operator(chunk) 30 | shutil.rmtree(tempdir) 31 | 32 | def test_segmentation(): 33 | print('test downsample and upload...') 34 | # compute parameters 35 | size = (16, 512, 512) 36 | 37 | # create image dataset using cloud-volume 38 | img = np.random.randint(np.iinfo(np.uint32).max, 39 | size=size, dtype=np.uint32) 40 | chunk = Chunk(img, voxel_offset=[2, 32, 32]) 41 | hierarchical_downsample(chunk) 42 | 43 | def test_image(): 44 | print('test downsample and upload...') 45 | # compute parameters 46 | size = (16, 512, 512) 47 | 48 | # create image dataset using cloud-volume 49 | img = np.random.randint(np.iinfo(np.uint8).max, 50 | size=size, dtype=np.uint8) 51 | chunk = Chunk(img, voxel_offset=[2, 32, 32]) 52 | hierarchical_downsample(chunk, layer_type='image') 53 | 54 | def test_psd_map(): 55 | print('test downsample and upload...') 56 | # compute parameters 57 | size = (16, 512, 512) 58 | 59 | # create image dataset using cloud-volume 60 | img = np.random.rand(*size).astype(np.float32) 61 | chunk = Chunk(img, voxel_offset=[2, 32, 32]) 62 | hierarchical_downsample(chunk, layer_type='image') 63 | -------------------------------------------------------------------------------- /tests/flow/test_load_precomputed.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import shutil 3 | import numpy as np 4 | 5 | from cloudvolume import CloudVolume 6 | from cloudvolume.lib import generate_random_string 7 | from cloudfiles import CloudFiles 8 | 9 | from chunkflow.lib.cartesian_coordinate import BoundingBox 10 | 11 | from chunkflow.flow.load_precomputed import LoadPrecomputedOperator 12 | 13 | 14 | class TestLoadPrecomputed(unittest.TestCase): 15 | def setUp(self): 16 | print('test volume cutout...') 17 | # compute parameters 18 | self.mip = 0 19 | self.size = (36, 448, 448) 20 | 21 | # create image dataset using cloud-volume 22 | img = np.random.randint(0, 256, size=self.size) 23 | self.img = img.astype(np.uint8) 24 | # save the input to disk 25 | self.volume_path = 'file:///tmp/test/cutout/' + generate_random_string( 26 | ) 27 | CloudVolume.from_numpy(np.transpose(self.img), 28 | vol_path=self.volume_path) 29 | 30 | # prepare blackout section ids 31 | self.blackout_section_ids = [17, 20] 32 | ids = {'section_ids': self.blackout_section_ids} 33 | stor = CloudFiles(self.volume_path) 34 | stor.put_json('blackout_section_ids.json', ids) 35 | 36 | def test_cutout(self): 37 | print('test volume cutout...') 38 | operator = LoadPrecomputedOperator(self.volume_path, mip=self.mip) 39 | 40 | offset = (4, 64, 64) 41 | shape = (28, 320, 320) 42 | output_bbox = BoundingBox.from_delta(offset, shape) 43 | chunk = operator(output_bbox) 44 | 45 | self.assertEqual(offset, chunk.voxel_offset) 46 | np.testing.assert_array_equal(chunk, self.img[4:-4, 64:-64, 64:-64]) 47 | 48 | shutil.rmtree('/tmp/test') 49 | 50 | def test_blackout_sections(self): 51 | print('test blackout sections...') 52 | operator = LoadPrecomputedOperator(self.volume_path, 53 | mip=self.mip, 54 | blackout_sections=True) 55 | 56 | offset = (4, 64, 64) 57 | shape = (28, 320, 320) 58 | output_bbox = BoundingBox.from_delta(offset, shape) 59 | chunk = operator(output_bbox) 60 | 61 | img = np.copy(self.img) 62 | for z in self.blackout_section_ids: 63 | img[z, :, :] = 0 64 | 65 | img = img[4:-4, 64:-64, 64:-64] 66 | self.assertTrue(img == chunk) 67 | shutil.rmtree('/tmp/test') 68 | 69 | 70 | if __name__ == '__main__': 71 | unittest.main() 72 | -------------------------------------------------------------------------------- /tests/flow/test_load_save.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import numpy as np 3 | import os 4 | import shutil 5 | 6 | from chunkflow.chunk import Chunk 7 | from chunkflow.lib.cartesian_coordinate import Cartesian 8 | 9 | from chunkflow.flow.save_pngs import SavePNGsOperator 10 | 11 | 12 | def load_save_h5(chunk): 13 | assert isinstance(chunk, Chunk) 14 | file_name = 'test.h5' 15 | if os.path.exists(file_name): 16 | os.remove(file_name) 17 | 18 | chunk.to_h5(file_name, chunk_size=None) 19 | 20 | chunk2 = Chunk.from_h5(file_name) 21 | assert chunk == chunk2 22 | assert chunk.voxel_offset == chunk2.voxel_offset 23 | os.remove(file_name) 24 | 25 | 26 | def load_save_tif(chunk): 27 | """We'll lost global offset information using tif format!""" 28 | assert isinstance(chunk, Chunk) 29 | file_name = 'test.tif' 30 | if os.path.exists(file_name): 31 | os.remove(file_name) 32 | 33 | chunk.to_tif(file_name) 34 | chunk2 = Chunk.from_tif(file_name) 35 | 36 | # we can not preserve the global offset here 37 | # so chunk2's global offset will all be 0 38 | np.testing.assert_array_equal(chunk.array, chunk2.array) 39 | os.remove(file_name) 40 | 41 | 42 | def save_pngs(chunk): 43 | # test save images 44 | output_path = '/tmp/test/' 45 | save_pngs_operator = SavePNGsOperator(output_path) 46 | save_pngs_operator(chunk) 47 | print('remove the temporary directory.') 48 | shutil.rmtree(output_path) 49 | 50 | 51 | class TestReadSave(unittest.TestCase): 52 | def test_load_save_image(self): 53 | print('test image io...') 54 | arr = np.random.randint(0, 256, size=(8, 16, 16), dtype=np.uint8) 55 | chunk = Chunk(arr, voxel_offset=Cartesian(1, 2, 3)) 56 | load_save_h5(chunk) 57 | load_save_tif(chunk) 58 | 59 | def test_load_save_aff(self): 60 | print('test affinitymap io...') 61 | arr = np.random.rand(3, 8, 16, 16).astype(np.float32) 62 | chunk = Chunk(arr, voxel_offset=Cartesian(1, 2, 3)) 63 | load_save_h5(chunk) 64 | 65 | 66 | if __name__ == '__main__': 67 | unittest.main() 68 | -------------------------------------------------------------------------------- /tests/flow/test_mesh.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import numpy as np 4 | import tempfile 5 | 6 | from chunkflow.chunk import Chunk 7 | from chunkflow.flow.mesh import MeshOperator 8 | 9 | 10 | def test_mesh(): 11 | ck = Chunk.create(dtype=np.float32) 12 | # segment it with threshold 13 | ck = ck < 0.5 14 | ck = ck.astype(np.uint32) 15 | 16 | tempdir = tempfile.mkdtemp() 17 | #mesher = MeshOperator('file://' + tempdir, 'precomputed', manifest=True) 18 | mesher = MeshOperator('file://' + tempdir, 'obj') 19 | mesher(ck) 20 | 21 | mesher = MeshOperator('file://' + tempdir, 'ply') 22 | mesher(ck) 23 | -------------------------------------------------------------------------------- /tests/flow/test_normalize_section_contrast.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import numpy as np 4 | np.random.seed(0) 5 | 6 | from chunkflow.chunk.image import Image 7 | 8 | 9 | def test_normalize_section_contrast(): 10 | #histogram = np.random.randint(10000, size=256, dtype=np.uint32) 11 | image = np.arange(256, dtype=np.uint8).reshape(1, 16, 16) 12 | assert image.ndim == 3 13 | image = Image(image, voxel_offset=(26774, 234, 456)) 14 | 15 | normalized = image.normalize_contrast() 16 | -------------------------------------------------------------------------------- /tests/flow/test_save_precomputed.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import shutil 4 | import tempfile 5 | from time import sleep 6 | 7 | import numpy as np 8 | from cloudvolume import CloudVolume 9 | 10 | from chunkflow.chunk import Chunk 11 | from chunkflow.flow.save_precomputed import SavePrecomputedOperator 12 | 13 | 14 | 15 | def test_save_image(): 16 | mip = 0 17 | size = (8, 64, 64) 18 | voxel_offset = (2, 4, 3) 19 | 20 | chunk = Chunk.create(size=size, dtype=np.uint8, 21 | voxel_offset=voxel_offset) 22 | 23 | tempdir = tempfile.mkdtemp() 24 | volume_path = 'file://' + tempdir 25 | print('construct volume from numpy array in ', tempdir) 26 | vol = CloudVolume.from_numpy( 27 | chunk.transpose(), 28 | vol_path=volume_path, 29 | voxel_offset=voxel_offset[::-1], 30 | chunk_size=(32, 32, 4), 31 | max_mip=1, 32 | layer_type='image' 33 | ) 34 | 35 | print('construct save operator') 36 | op = SavePrecomputedOperator( 37 | volume_path, mip, 38 | upload_log = True, 39 | create_thumbnail = False, 40 | name = 'save' 41 | ) 42 | 43 | print('really save the chunk.') 44 | op(chunk, log={'timer': {'save': 43}}) 45 | 46 | sleep(2) 47 | shutil.rmtree(tempdir) 48 | -------------------------------------------------------------------------------- /tests/lib/aws/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/seung-lab/chunkflow/3c883ef95efb595d4172fe044ca3959b9337621d/tests/lib/aws/__init__.py -------------------------------------------------------------------------------- /tests/lib/aws/test_sqs_queue.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import unittest 3 | from chunkflow.lib.aws.sqs_queue import SQSQueue 4 | 5 | 6 | @pytest.mark.skip(reason="need credential") 7 | class TestSQSQueue(unittest.TestCase): 8 | def setUp(self): 9 | queue_name = 'chunkflow-test' 10 | print('test sqs queue using queue of {}'.format(queue_name)) 11 | # use small fetch_wait_time_seconds to make the test faster 12 | # we should use 20 seconds in production run to use long pull 13 | # see more details in the class documentation. 14 | self.queue = SQSQueue(queue_name, 15 | wait_if_empty=None, 16 | fetch_wait_time_seconds=1, 17 | retry_times=3) 18 | 19 | def test_send_and_receive_message_list(self): 20 | print('start sending messages...') 21 | message_list = [] 22 | for i in range(23): 23 | message_list.append(str(i)) 24 | self.queue.send_message_list(message_list) 25 | 26 | print('start receiving messages in iteration...') 27 | i = 0 28 | for receipt_handle, message in self.queue: 29 | print('received message: {}'.format(i)) 30 | self.assertTrue(int(message) in range(23)) 31 | i += 1 32 | self.queue.delete(receipt_handle) 33 | 34 | 35 | if __name__ == '__main__': 36 | unittest.main() 37 | -------------------------------------------------------------------------------- /tests/lib/sp3_bboxes.txt: -------------------------------------------------------------------------------- 1 | # the columes are x0,y0,z0,x1,y1,z1 2 | # we'll have to adjust the columns after loading as numpy array 3 | 2944, 5408, 2688, 3360, 5824, 3104 4 | 4928, 3744, 2368, 5344, 4160, 2784 5 | 4448, 2944, 3680, 4864, 3360, 4096 6 | 6528, 4832, 1920, 6944, 5248, 2336 7 | 9248, 6048, 3776, 9664, 6464, 4192 8 | 4320, 3520, 5632, 4736, 3936, 6048 9 | 8384, 3168, 5152, 8800, 3584, 5568 10 | 6304, 2976, 2304, 6720, 3392, 2720 11 | -------------------------------------------------------------------------------- /tests/lib/test_cartesian_coordinate.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | 5 | from cloudvolume.lib import Bbox, Vec 6 | 7 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian, to_cartesian, BoundingBoxes, PhysicalBoundingBox 8 | 9 | 10 | def test_cartesian(): 11 | assert to_cartesian(None) is None 12 | ct = (1,2,3) 13 | assert to_cartesian(ct) == Cartesian(1,2,3) 14 | 15 | ct = Cartesian(1,2,3) 16 | ct += 2 17 | assert ct == Cartesian(3,4,5) 18 | 19 | ct -= 2 20 | assert ct == Cartesian(1,2,3) 21 | 22 | np.testing.assert_equal(ct.vec, Vec(1,2,3)) 23 | 24 | ct = Cartesian(3,4,5) 25 | ct = ct // 2 26 | assert ct == Cartesian(1,2,2) 27 | 28 | # note that 2*ct will repeat the elements of ct! 29 | ct2 = ct*2 30 | assert ct2 > ct 31 | assert ct2 >= ct 32 | assert ct < ct2 33 | assert ct <= ct2 34 | 35 | ct3 = ct / 2 36 | assert ct3 == Cartesian(0.5, 1, 1) 37 | 38 | ct4 = Cartesian.from_collection((1,2,3)) 39 | assert ct4 == Cartesian(1, 2, 3) 40 | 41 | assert Cartesian(0, 0, 0)*Cartesian(1,2,3) == Cartesian(0, 0, 0) 42 | 43 | assert Cartesian(4,6,8) / Cartesian(2,3,2) == Cartesian(2,2,4) 44 | 45 | assert -Cartesian(1,-2,3) == Cartesian(-1, 2, -3) 46 | 47 | assert Cartesian(1,2,3).tuple == (1,2,3) 48 | assert Cartesian(1,2,3).vec is not None 49 | 50 | 51 | def test_bounding_box(): 52 | bbox = BoundingBox.from_string('3166-3766_7531-8131_2440-3040') 53 | assert bbox == BoundingBox(Cartesian(3166, 7531, 2440), Cartesian(3766, 8131, 3040)) 54 | 55 | bbox = BoundingBox.from_string('Sp1,3166-3766_7531-8131_2440-3040.h5') 56 | assert bbox == BoundingBox(Cartesian(3166, 7531, 2440), Cartesian(3766, 8131, 3040)) 57 | 58 | bbox = Bbox.from_delta((1,3,2), (64, 32, 8)) 59 | bbox = BoundingBox.from_bbox(bbox) 60 | assert bbox.start == Cartesian(1,3,2) 61 | assert bbox.stop == Cartesian(65, 35, 10) 62 | 63 | bbox = bbox.clone() 64 | assert isinstance(bbox, BoundingBox) 65 | 66 | minpt = Cartesian(1,2,3) 67 | maxpt = Cartesian(2,3,4) 68 | bbox = BoundingBox(minpt, maxpt) 69 | assert bbox.start == minpt 70 | assert bbox.stop == maxpt 71 | assert bbox.shape == Cartesian(1,1,1) 72 | 73 | bbox = BoundingBox.from_center(Cartesian(1,2,3), 3) 74 | assert bbox == BoundingBox.from_list([-2, -1, 0, 4, 5, 6]) 75 | 76 | bbox = BoundingBox.from_center(Cartesian(1,2,3), 3, even_size=False) 77 | assert bbox == BoundingBox.from_list([-2, -1, 0, 5, 6, 7]) 78 | 79 | bbox1 = BoundingBox.from_list([0,1,2, 2,3,4]) 80 | bbox2 = BoundingBox.from_list([1,2,3, 3,4,5]) 81 | assert bbox1.union(bbox2) == BoundingBox.from_list([0,1,2, 3,4,5]) 82 | assert bbox1.intersection(bbox2) == BoundingBox.from_list([1,2,3, 2,3,4]) 83 | 84 | minpt = Cartesian(1,2,3) 85 | maxpt = Cartesian(3,4,5) 86 | bbox = BoundingBox(minpt, maxpt) 87 | bbox_decomp = bbox.decompose(bbox.shape // 2) 88 | assert len(bbox_decomp) == 8 89 | assert bbox_decomp[0].start == minpt 90 | assert bbox_decomp[-1].stop == maxpt 91 | 92 | 93 | def test_bounding_boxes(): 94 | fname = os.path.join(os.path.dirname(__file__), 'sp3_bboxes.txt') 95 | bboxes = BoundingBoxes.from_file(fname) 96 | fname = os.path.join(os.path.dirname(__file__), 'sp3_bboxes.npy') 97 | bboxes.to_file(fname) 98 | os.remove(fname) 99 | 100 | 101 | def test_physical_bounding_box(): 102 | start = Cartesian(0, 1, 2) 103 | stop = Cartesian(2, 3, 4) 104 | voxel_size = Cartesian(2, 2, 2) 105 | pbbox = PhysicalBoundingBox(start, stop, voxel_size) 106 | 107 | pbbox2 = pbbox.to_other_voxel_size(Cartesian(1,1,1)) 108 | assert pbbox2.start == Cartesian(0,2,4) 109 | -------------------------------------------------------------------------------- /tests/lib/test_synapses.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tempfile 3 | import numpy as np 4 | from chunkflow.synapses import Synapses 5 | 6 | 7 | def test_synapses(): 8 | pre = np.arange(6).reshape((2,3)) 9 | post = np.array([[0, 1,2,3]]) 10 | syns = Synapses(pre, post=post) 11 | 12 | assert syns.pre_num == 2 13 | assert syns.post_num == 1 14 | 15 | temp_dir = tempfile.mkdtemp() 16 | temp_fname = os.path.join(temp_dir, 'test.h5') 17 | syns.to_h5(temp_fname) 18 | syns2 = Synapses.from_h5(temp_fname) 19 | assert syns == syns2 20 | 21 | syns.remove_synapses_without_post() 22 | assert syns.pre_num == 1 23 | os.remove(temp_fname) 24 | -------------------------------------------------------------------------------- /tests/requirements.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | coverage 3 | pytest-cov 4 | -------------------------------------------------------------------------------- /tests/test_command_line.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import os 5 | 6 | 7 | def test_composable_command_line_interface(): 8 | os.system('./command_line.sh') 9 | -------------------------------------------------------------------------------- /tests/test_volume.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | 3 | import numpy as np 4 | 5 | from cloudvolume.lib import generate_random_string 6 | from chunkflow.lib.cartesian_coordinate import BoundingBox, Cartesian 7 | from chunkflow.volume import PrecomputedVolume 8 | 9 | 10 | def test_volume(): 11 | print('test volume cutout...') 12 | # compute parameters 13 | size = (36, 448, 448) 14 | 15 | # create image dataset using cloud-volume 16 | img = np.random.randint(0, 256, size=size) 17 | img = img.astype(np.uint8) 18 | # save the input to disk 19 | volume_path = 'file:///tmp/test/volume/' + \ 20 | generate_random_string() 21 | 22 | vol = PrecomputedVolume.from_numpy( 23 | img, 24 | volume_path 25 | ) 26 | 27 | offset = Cartesian(4, 64, 64) 28 | shape = (28, 320, 320) 29 | bbox = BoundingBox.from_delta(offset, shape) 30 | chunk = vol.cutout(bbox) 31 | # chunk = chunk.squeeze_channel() 32 | 33 | assert offset == chunk.voxel_offset 34 | np.testing.assert_array_equal(chunk, img[4:-4, 64:-64, 64:-64]) 35 | 36 | shutil.rmtree('/tmp/test') --------------------------------------------------------------------------------