├── .github
    └── workflows
    │   ├── docker-publish.yml
    │   ├── python-package.yml
    │   ├── python-publish.yml
    │   └── stale.yml
├── .gitignore
├── .readthedocs.yaml
├── LICENSE
├── MANIFEST.in
├── README.md
├── assets
    ├── bbbc039-cpn-u22-demo-arrow.png
    ├── bbbc039-cpn-u22-demo.png
    ├── bbbc041-cpn-u22-demo.png
    └── neurips-cellseg-demo.png
├── celldetection
    ├── __init__.py
    ├── __meta__.py
    ├── callbacks
    │   ├── __init__.py
    │   ├── dropout.py
    │   └── keepalive.py
    ├── data
    │   ├── __init__.py
    │   ├── cpn.py
    │   ├── datasets
    │   │   ├── __init__.py
    │   │   ├── bbbc038.py
    │   │   ├── bbbc039.py
    │   │   ├── bbbc041.py
    │   │   ├── generic.py
    │   │   └── synth.py
    │   ├── instance_eval.py
    │   ├── misc.py
    │   ├── segmentation.py
    │   ├── toydata.py
    │   └── transforms.py
    ├── models
    │   ├── __init__.py
    │   ├── commons.py
    │   ├── convnext.py
    │   ├── convnextv2.py
    │   ├── cpn.py
    │   ├── densenet.py
    │   ├── features.py
    │   ├── filters.py
    │   ├── fpn.py
    │   ├── hosted.py
    │   ├── inference.py
    │   ├── lightning_base.py
    │   ├── lightning_cpn.py
    │   ├── loss.py
    │   ├── mamba.py
    │   ├── manet.py
    │   ├── mobilenetv3.py
    │   ├── normalization.py
    │   ├── ppm.py
    │   ├── resnet.py
    │   ├── smp.py
    │   ├── timmodels.py
    │   └── unet.py
    ├── mpi
    │   ├── __init__.py
    │   └── mpi.py
    ├── ops
    │   ├── __init__.py
    │   ├── boxes.py
    │   ├── commons.py
    │   ├── cpn.py
    │   ├── draw.py
    │   ├── features.py
    │   ├── loss.py
    │   └── normalization.py
    ├── optim
    │   ├── __init__.py
    │   └── lr_scheduler.py
    ├── util
    │   ├── __init__.py
    │   ├── logging.py
    │   ├── schedule.py
    │   ├── shm_cache.py
    │   ├── timer.py
    │   └── util.py
    └── visualization
    │   ├── __init__.py
    │   ├── cmaps.py
    │   └── images.py
├── celldetection_scripts
    ├── __init__.py
    └── cpn_inference.py
├── demos
    ├── Cell Detection with Contour Proposal Networks.ipynb
    ├── demo-binary.ipynb
    └── demo-multiclass.ipynb
├── docker
    ├── Dockerfile
    ├── DockerfileNvidia
    └── README.md
├── docs
    ├── Makefile
    ├── make.bat
    ├── mk.sh
    ├── requirements.txt
    └── source
    │   ├── _static
    │       ├── css
    │       │   └── theme.css
    │       └── img
    │       │   └── magnify.svg
    │   ├── _templates
    │       └── layout.html
    │   ├── celldetection.callbacks.rst
    │   ├── celldetection.data.rst
    │   ├── celldetection.models.rst
    │   ├── celldetection.mpi.rst
    │   ├── celldetection.ops.rst
    │   ├── celldetection.optim.rst
    │   ├── celldetection.rst
    │   ├── conf.py
    │   ├── index.rst
    │   ├── inference.rst
    │   └── install.rst
├── hubconf.py
├── pytest.ini
├── requirements.txt
├── setup.py
└── tests
    └── test_cpn_inference.py


/.github/workflows/docker-publish.yml:
--------------------------------------------------------------------------------
 1 | name: Publish Docker image
 2 | 
 3 | on:
 4 |   release:
 5 |     types: [published]
 6 |   workflow_dispatch:
 7 |     inputs:
 8 |       logLevel:
 9 |         description: 'Log level (e.g., debug, info, warning, error)'
10 |         required: false
11 |         default: 'info'
12 |         type: choice
13 |         options:
14 |           - debug
15 |           - info
16 |           - warning
17 |           - error
18 | 
19 | jobs:
20 |   push_to_registry:
21 |     name: Push Docker image to Docker Hub
22 |     runs-on: ubuntu-latest
23 |     steps:
24 |       - name: Delete tools cache
25 |         run: rm -rf /opt/hostedtoolcache
26 |       - name: Check out the repo
27 |         uses: actions/checkout@v3
28 |       - name: Set up QEMU
29 |         uses: docker/setup-qemu-action@v2
30 |       - name: Set up Docker Buildx
31 |         uses: docker/setup-buildx-action@v2
32 |       - name: Log in to Docker Hub
33 |         uses: docker/login-action@f4ef78c080cd8ba55a85445d5b36e214a81df20a
34 |         with:
35 |           username: ${{ secrets.DOCKER_USERNAME }}
36 |           password: ${{ secrets.DOCKER_PASSWORD }}
37 |       - name: Extract metadata (tags, labels) for Docker
38 |         id: meta
39 |         uses: docker/metadata-action@9ec57ed1fcdbf14dcef7dfbe97b2010124a938b7
40 |         with:
41 |           images: ericup/celldetection
42 |       - name: Build and push Docker image
43 |         uses: docker/build-push-action@v4
44 |         with:
45 |           context: .
46 |           file: ./docker/Dockerfile
47 |           platforms: linux/amd64,linux/arm64
48 |           push: true
49 |           tags: ${{ steps.meta.outputs.tags }}
50 |           labels: ${{ steps.meta.outputs.labels }}


--------------------------------------------------------------------------------
/.github/workflows/python-package.yml:
--------------------------------------------------------------------------------
 1 | name: Test
 2 | 
 3 | on:
 4 |   push:
 5 |     branches: [ main ]
 6 |   pull_request:
 7 |     branches: [ main ]
 8 | 
 9 | jobs:
10 |   test:
11 | 
12 |     runs-on: ${{ matrix.os }}
13 |     strategy:
14 |       fail-fast: false
15 |       matrix:
16 |         os: [ubuntu-latest, macos-latest, windows-latest]
17 |         python-version: [3.8, 3.9, '3.10', '3.11']
18 | 
19 |     steps:
20 |     - uses: actions/checkout@v2
21 |     - name: Set up Python ${{ matrix.python-version }}
22 |       uses: actions/setup-python@v2
23 |       with:
24 |         python-version: ${{ matrix.python-version }}
25 |     - name: Install dependencies
26 |       run: |
27 |         python -m pip install --upgrade pip
28 |         python -m pip install flake8 pytest
29 |         pip install -r requirements.txt
30 |     - name: Lint with flake8
31 |       run: |
32 |         # stop the build if there are Python syntax errors or undefined names
33 |         flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
34 |         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
35 |         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
36 |     - name: Install package
37 |       run: python -m pip install .
38 |     - name: Print version
39 |       run: python -c "import celldetection; print(celldetection.__version__)"
40 |     - name: Run tests
41 |       run: pytest
42 | 


--------------------------------------------------------------------------------
/.github/workflows/python-publish.yml:
--------------------------------------------------------------------------------
 1 | # This workflow will upload a Python Package using Twine when a release is created
 2 | # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
 3 | 
 4 | # This workflow uses actions that are not certified by GitHub.
 5 | # They are provided by a third-party and are governed by
 6 | # separate terms of service, privacy policy, and support
 7 | # documentation.
 8 | 
 9 | name: Upload Python Package
10 | 
11 | on:
12 |   release:
13 |     types: [published]
14 | 
15 | jobs:
16 |   deploy:
17 | 
18 |     runs-on: ubuntu-latest
19 | 
20 |     steps:
21 |     - uses: actions/checkout@v2
22 |     - name: Set up Python
23 |       uses: actions/setup-python@v2
24 |       with:
25 |         python-version: '3.x'
26 |     - name: Install dependencies
27 |       run: |
28 |         python -m pip install --upgrade pip
29 |         pip install build
30 |     - name: Build package
31 |       run: python -m build
32 |     - name: Publish package
33 |       uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
34 |       with:
35 |         user: __token__
36 |         password: ${{ secrets.PYPI_API_TOKEN }}
37 | 


--------------------------------------------------------------------------------
/.github/workflows/stale.yml:
--------------------------------------------------------------------------------
 1 | # This workflow warns and then closes issues and PRs that have had no activity for a specified amount of time.
 2 | #
 3 | # You can adjust the behavior by modifying this file.
 4 | # For more information, see:
 5 | # https://github.com/actions/stale
 6 | name: Mark stale issues and pull requests
 7 | 
 8 | on:
 9 |   schedule:
10 |   - cron: '30 19 * * *'  # Runs daily at 19:30 UTC
11 |   workflow_dispatch:      # Allows the workflow to be triggered manually from the GitHub UI
12 | 
13 | jobs:
14 |   stale:
15 |     runs-on: ubuntu-latest
16 |     permissions:
17 |       issues: write
18 |       pull-requests: write
19 | 
20 |     steps:
21 |     - uses: actions/stale@v5
22 |       with:
23 |         repo-token: ${{ secrets.GITHUB_TOKEN }}
24 |         stale-issue-message: 'This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs.'
25 |         stale-pr-message: 'This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs.'
26 |         stale-issue-label: 'no-issue-activity'
27 |         stale-pr-label: 'no-pr-activity'
28 |         days-before-stale: 60  # Number of days of inactivity before an issue or PR is marked stale
29 |         days-before-close: 7   # Number of days to close the issue/PR after it has been marked as stale
30 |         exempt-issue-labels: 'pinned, security'  # Labels which exempt an issue from being marked as stale
31 |         exempt-pr-labels: 'pinned, security'     # Labels which exempt a PR from being marked as stale
32 |         operations-per-run: 30                  # Limit the number of operations per run
33 |         remove-stale-when-updated: true          # Remove stale label when new activity occurs
34 |         debug-only: true                        # Set to false to perform actions, true to only log them
35 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # PyCharm
  2 | .idea/
  3 | 
  4 | # Byte-compiled / optimized / DLL files
  5 | __pycache__/
  6 | *.py[cod]
  7 | *$py.class
  8 | 
  9 | # C extensions
 10 | *.so
 11 | 
 12 | # Distribution / packaging
 13 | .Python
 14 | build/
 15 | develop-eggs/
 16 | dist/
 17 | downloads/
 18 | eggs/
 19 | .eggs/
 20 | lib/
 21 | lib64/
 22 | parts/
 23 | sdist/
 24 | var/
 25 | wheels/
 26 | pip-wheel-metadata/
 27 | share/python-wheels/
 28 | *.egg-info/
 29 | .installed.cfg
 30 | *.egg
 31 | MANIFEST
 32 | 
 33 | # PyInstaller
 34 | #  Usually these files are written by a python script from a template
 35 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 36 | *.manifest
 37 | *.spec
 38 | 
 39 | # Installer logs
 40 | pip-log.txt
 41 | pip-delete-this-directory.txt
 42 | 
 43 | # Unit test / coverage reports
 44 | htmlcov/
 45 | .tox/
 46 | .nox/
 47 | .coverage
 48 | .coverage.*
 49 | .cache
 50 | nosetests.xml
 51 | coverage.xml
 52 | *.cover
 53 | *.py,cover
 54 | .hypothesis/
 55 | .pytest_cache/
 56 | 
 57 | # Translations
 58 | *.mo
 59 | *.pot
 60 | 
 61 | # Django stuff:
 62 | *.log
 63 | local_settings.py
 64 | db.sqlite3
 65 | db.sqlite3-journal
 66 | 
 67 | # Flask stuff:
 68 | instance/
 69 | .webassets-cache
 70 | 
 71 | # Scrapy stuff:
 72 | .scrapy
 73 | 
 74 | # Sphinx documentation
 75 | docs/_build/
 76 | 
 77 | # PyBuilder
 78 | target/
 79 | 
 80 | # Jupyter Notebook
 81 | .ipynb_checkpoints
 82 | 
 83 | # IPython
 84 | profile_default/
 85 | ipython_config.py
 86 | 
 87 | # pyenv
 88 | .python-version
 89 | 
 90 | # pipenv
 91 | #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 92 | #   However, in case of collaboration, if having platform-specific dependencies or dependencies
 93 | #   having no cross-platform support, pipenv may install dependencies that don't work, or not
 94 | #   install all needed dependencies.
 95 | #Pipfile.lock
 96 | 
 97 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow
 98 | __pypackages__/
 99 | 
100 | # Celery stuff
101 | celerybeat-schedule
102 | celerybeat.pid
103 | 
104 | # SageMath parsed files
105 | *.sage.py
106 | 
107 | # Environments
108 | .env
109 | .venv
110 | env/
111 | venv/
112 | ENV/
113 | env.bak/
114 | venv.bak/
115 | 
116 | # Spyder project settings
117 | .spyderproject
118 | .spyproject
119 | 
120 | # Rope project settings
121 | .ropeproject
122 | 
123 | # mkdocs documentation
124 | /site
125 | 
126 | # mypy
127 | .mypy_cache/
128 | .dmypy.json
129 | dmypy.json
130 | 
131 | # Pyre type checker
132 | .pyre/
133 | 


--------------------------------------------------------------------------------
/.readthedocs.yaml:
--------------------------------------------------------------------------------
 1 | # .readthedocs.yaml
 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 | # Set the version of Python and other tools you might need
 9 | build:
10 |   os: ubuntu-22.04
11 |   tools:
12 |     python: "3.9"
13 |     # You can also specify other tool versions:
14 |     # nodejs: "19"
15 |     # rust: "1.64"
16 |     # golang: "1.19"
17 | 
18 | # Build documentation in the docs/ directory with Sphinx
19 | sphinx:
20 |    configuration: docs/source/conf.py
21 | 
22 | # If using Sphinx, optionally build your docs in additional formats such as PDF
23 | # formats:
24 | #    - pdf
25 | 
26 | # Optionally declare the Python requirements required to build your docs
27 | python:
28 |    install:
29 |    - requirements: docs/requirements.txt
30 |    - method: pip
31 |      path: .
32 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
  1 |                                  Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "[]"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright [yyyy] [name of copyright owner]
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/MANIFEST.in:
--------------------------------------------------------------------------------
1 | include *.txt


--------------------------------------------------------------------------------
/assets/bbbc039-cpn-u22-demo-arrow.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/FZJ-INM1-BDA/celldetection/92ad83f63185187068c57595362c8e4a550d37cc/assets/bbbc039-cpn-u22-demo-arrow.png


--------------------------------------------------------------------------------
/assets/bbbc039-cpn-u22-demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/FZJ-INM1-BDA/celldetection/92ad83f63185187068c57595362c8e4a550d37cc/assets/bbbc039-cpn-u22-demo.png


--------------------------------------------------------------------------------
/assets/bbbc041-cpn-u22-demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/FZJ-INM1-BDA/celldetection/92ad83f63185187068c57595362c8e4a550d37cc/assets/bbbc041-cpn-u22-demo.png


--------------------------------------------------------------------------------
/assets/neurips-cellseg-demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/FZJ-INM1-BDA/celldetection/92ad83f63185187068c57595362c8e4a550d37cc/assets/neurips-cellseg-demo.png


--------------------------------------------------------------------------------
/celldetection/__init__.py:
--------------------------------------------------------------------------------
 1 | """
 2 | Primary symbols include classes and functions for visualization, configuration, timing and generally useful utilities.
 3 | """
 4 | from . import models
 5 | from . import ops
 6 | from . import util
 7 | from . import visualization as vis
 8 | from . import data
 9 | from . import mpi
10 | from . import callbacks
11 | from . import optim
12 | from .data import toydata
13 | from .data.misc import universal_dict_collate_fn, to_tensor
14 | from .util import *
15 | from .visualization import *
16 | from .__meta__ import __version__
17 | 


--------------------------------------------------------------------------------
/celldetection/__meta__.py:
--------------------------------------------------------------------------------
1 | __title__ = 'CellDetection'
2 | __version__ = '0.4.9'
3 | __author__ = 'Eric Upschulte'
4 | __email__ = 'e.upschulte@fz-juelich.de'
5 | __license__ = 'Apache License, Version 2.0'
6 | __summary__ = 'Cell Detection with PyTorch.'
7 | __url__ = __uri__ = 'https://celldetection.org'
8 | 


--------------------------------------------------------------------------------
/celldetection/callbacks/__init__.py:
--------------------------------------------------------------------------------
1 | from .dropout import *
2 | from .keepalive import *


--------------------------------------------------------------------------------
/celldetection/callbacks/dropout.py:
--------------------------------------------------------------------------------
 1 | from pytorch_lightning import Trainer, LightningModule, Callback
 2 | from torch.nn.modules.dropout import _DropoutNd
 3 | from typing import Any
 4 | 
 5 | __all__ = ['StepDropout']
 6 | 
 7 | 
 8 | class StepDropout(Callback):
 9 |     def __init__(
10 |             self,
11 |             step_size,
12 |             base_drop_rate,
13 |             gamma=0.,
14 |             update_interval='epoch',
15 |             log=True,
16 |             log_name='drop_rate',
17 |             ascending=False,
18 |             **kwargs
19 |     ):
20 |         """Step Dropout.
21 | 
22 |         A simple Dropout Scheduler.
23 | 
24 |         References:
25 |             - https://arxiv.org/abs/2303.01500
26 | 
27 |         Examples:
28 |             >>> from pytorch_lightning import Trainer
29 |             >>> # Early Dropout (drop rate from .1 to 0 after 50 epochs)
30 |             >>> trainer = Trainer(callbacks=[StepDropout(50, base_drop_rate=.1, gamma=0.)])
31 | 
32 |             >>> # Late Dropout (drop rate from 0 to .1 after 50 epochs)
33 |             >>> trainer = Trainer(callbacks=[StepDropout(50, base_drop_rate=.1, gamma=0., ascending=True)])
34 | 
35 |         Args:
36 |             step_size: Period of drop rate decay.
37 |             base_drop_rate: Base drop rate.
38 |             gamma: Multiplicative factor of drop rate decay. Default: 0. to replicate "Early Dropout".
39 |             update_interval: One of ``('step', 'epoch')``.
40 |             log: Whether to log drop rates using ``module.log(log_name, drop_rate)``.
41 |             log_name: Name for logging.
42 |             logger: If ``True`` logs to the logger.
43 |             ascending: If ``True`` drop rate decays from right to left, i.e. it starts at ``0`` and
44 |                 ascends towards ``base_drop_rate``. Using ``ascending=True, gamma=0.`` replicates "Late Dropout".
45 |             **kwargs: Keyword arguments for ``module.log``.
46 |         """
47 |         super().__init__()
48 |         self.step_size = step_size
49 |         self.gamma = gamma
50 |         self.base_drop_rate = base_drop_rate
51 |         assert update_interval in ('epoch', 'step')
52 |         self.update_interval = update_interval
53 |         self.last_rate = -1
54 |         self.log = log
55 |         self.log_name = log_name
56 |         self.log_kwargs = {**dict(
57 |             on_step=self.update_interval == 'step',
58 |             on_epoch=self.update_interval == 'epoch',
59 |         ), **kwargs}
60 |         self.ascending = ascending
61 | 
62 |     def update_drop_rate(self, pl_module: "LightningModule", drop_rate: float):
63 |         self.last_rate = drop_rate
64 |         for mod in pl_module.modules():
65 |             if isinstance(mod, _DropoutNd):
66 |                 mod.p = drop_rate
67 |         if self.log:
68 |             pl_module.log(self.log_name, drop_rate, **self.log_kwargs)
69 | 
70 |     @staticmethod
71 |     def get_rate(base, gamma, step, step_size, ascending):
72 |         return base * (ascending + (-1 if ascending else 1) * gamma ** (step // step_size))
73 | 
74 |     def on_train_epoch_start(self, trainer: "Trainer", pl_module: "LightningModule") -> None:
75 |         if self.update_interval == 'epoch':
76 |             self.update_drop_rate(pl_module, self.get_rate(
77 |                 self.base_drop_rate, self.gamma, trainer.current_epoch, self.step_size, self.ascending))
78 | 
79 |     def on_train_batch_start(self, trainer: "Trainer", pl_module: "LightningModule", batch: Any,
80 |                              batch_idx: int) -> None:
81 |         if self.update_interval == 'step':
82 |             self.update_drop_rate(pl_module, self.get_rate(
83 |                 self.base_drop_rate, self.gamma, trainer.global_step, self.step_size, self.ascending))
84 | 


--------------------------------------------------------------------------------
/celldetection/callbacks/keepalive.py:
--------------------------------------------------------------------------------
  1 | import threading
  2 | import time
  3 | import os
  4 | import warnings
  5 | from typing import Any
  6 | import pytorch_lightning as pl
  7 | 
  8 | __all__ = ['KeepAlive']
  9 | 
 10 | 
 11 | def terminate_program(log_fn=None):
 12 |     from ..mpi.mpi import has_mpi, MPI
 13 |     import signal
 14 | 
 15 |     log_fn = log_fn or print
 16 | 
 17 |     if 'SLURM_JOB_ID' in os.environ:
 18 |         try:
 19 |             from subprocess import Popen, PIPE
 20 |             cmd = ['scancel', os.environ['SLURM_JOB_ID']]
 21 |             log_fn(' '.join(cmd), flush=True)
 22 |             process = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
 23 |             stdout, stderr = process.communicate()
 24 |             log_fn(f'scancel stdout, stderr:\n{stdout.decode()}\n\n{stderr.decode()}')
 25 |         except Exception as e:
 26 |             log_fn(f'Failed terminating the SLURM job: {e}')
 27 | 
 28 |     if has_mpi():
 29 |         try:
 30 |             log_fn('MPI.COMM_WORLD.Abort()', flush=True)
 31 |             MPI.COMM_WORLD.Abort()  # May not work if comm is down
 32 |         except Exception as e:
 33 |             log_fn(f'Failed terminating the MPI job: {e}')
 34 | 
 35 |     log_fn('Sending SIGTERM signal to the current process...')
 36 |     os.kill(os.getpid(), signal.SIGTERM)
 37 |     time.sleep(20)
 38 | 
 39 |     log_fn('Forcefully exiting...')
 40 |     os._exit(1)
 41 | 
 42 | 
 43 | class KeepAlive(pl.Callback):
 44 |     def __init__(self, timeout=60 * 15, interval=None, start=False, verbose=True):  # timeout in seconds
 45 |         super().__init__()
 46 |         self.timeout = timeout
 47 |         self.interval = timeout / 4 if interval is None else interval
 48 |         self.last_signal = time.time()
 49 |         self.monitor_thread = None
 50 |         self.running = False
 51 |         self.verbose = verbose
 52 |         if start:
 53 |             self.start()
 54 | 
 55 |     def log(self, *args, **kwargs):
 56 |         if self.verbose:
 57 |             print(*args, **kwargs, flush=True)
 58 | 
 59 |     def start(self):
 60 |         self.log("Starting KeepAlive messages")
 61 | 
 62 |         if not self.running:
 63 |             self.running = True
 64 |             self.monitor_thread = threading.Thread(target=self._monitor)
 65 |             self.monitor_thread.daemon = True  # exits when main thread exits
 66 |             self.monitor_thread.start()
 67 | 
 68 |     def stop(self):
 69 |         if self.running:
 70 |             self.running = False
 71 |             self.monitor_thread.join()
 72 | 
 73 |     def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"):
 74 |         if trainer.local_rank == 0:
 75 |             self.start()
 76 | 
 77 |     def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"):
 78 |         if trainer.local_rank == 0:
 79 |             self.stop()
 80 | 
 81 |     def on_train_batch_start(
 82 |             self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int,
 83 |             dataset_idx: int = 0
 84 |     ) -> None:
 85 |         self.keep_alive_signal()
 86 | 
 87 |     def on_terminate(self):
 88 |         terminate_program(log_fn=self.log)
 89 | 
 90 |     def _monitor(self):
 91 |         while self.running:
 92 |             time.sleep(self.interval)  # Check every second
 93 |             if time.time() - self.last_signal > self.timeout:
 94 |                 msg = "Keepalive signal timeout. Terminating the program."
 95 |                 warnings.warn(msg)
 96 |                 self.log(msg)
 97 |                 self.on_terminate()
 98 | 
 99 |     def keep_alive_signal(self):
100 |         self.last_signal = time.time()
101 | 
102 |     on_validate_start = on_test_start = on_predict_start = on_train_start
103 |     on_validate_end = on_test_end = on_predict_end = on_train_end
104 |     on_predict_batch_start = on_validation_batch_start = on_test_batch_start = on_train_batch_start
105 | 


--------------------------------------------------------------------------------
/celldetection/data/__init__.py:
--------------------------------------------------------------------------------
 1 | """
 2 | This submodule contains data related code and numpy operations.
 3 | """
 4 | from .cpn import *
 5 | from .cpn import labels2contour_list as labels2contours
 6 | from .misc import *
 7 | from .instance_eval import LabelMatcherList, LabelMatcher
 8 | from .segmentation import *
 9 | from .datasets import *
10 | 


--------------------------------------------------------------------------------
/celldetection/data/datasets/__init__.py:
--------------------------------------------------------------------------------
1 | from .bbbc039 import *
2 | from .bbbc041 import *
3 | from .generic import *
4 | from .synth import *
5 | 


--------------------------------------------------------------------------------
/celldetection/data/datasets/bbbc038.py:
--------------------------------------------------------------------------------
 1 | from os.path import join, dirname, basename, isdir
 2 | from os import makedirs
 3 | from imageio import imread
 4 | from glob import glob
 5 | import numpy as np
 6 | from ..cpn import masks2labels
 7 | import torchvision
 8 | 
 9 | __all__ = ['download_bbbc038', 'BBBC038Train']
10 | 
11 | 
12 | def download_bbbc038(directory):
13 |     """Download BBBC038.
14 | 
15 |     Download and extract the BBBC038 dataset to given directory.
16 | 
17 |     References:
18 |         https://bbbc.broadinstitute.org/BBBC038
19 | 
20 |     Args:
21 |         directory: Root directory.
22 | 
23 |     """
24 |     for url in [
25 |         'https://data.broadinstitute.org/bbbc/BBBC038/stage1_train.zip',
26 |         'https://data.broadinstitute.org/bbbc/BBBC038/stage1_test.zip',
27 |         'https://data.broadinstitute.org/bbbc/BBBC038/stage2_test_final.zip'
28 |     ]:
29 |         directory_ = join(directory, basename(url).split('.')[0])
30 |         makedirs(directory_, exist_ok=True)
31 |         torchvision.datasets.utils.download_and_extract_archive(url, directory, extract_root=directory_)
32 | 
33 | 
34 | class BBBC038Train:
35 |     def __init__(self, directory, download=False):
36 |         if download and not isdir(join(directory, 'stage1_train')):
37 |             download_bbbc038(directory)
38 |         self.image_f = sorted(glob(join(directory, 'stage1_train', '*', 'images', '*.*')))
39 |         self.label_f = [sorted(glob(join(dirname(dirname(f)), 'masks', '*.*'))) for f in self.image_f]
40 | 
41 |     def __getitem__(self, item):
42 |         img_f = self.image_f[item]
43 |         lbl_f = self.label_f[item]
44 | 
45 |         img = imread(img_f)
46 |         masks = np.stack([imread(f) for f in lbl_f])
47 |         lbl = masks2labels(masks)
48 |         return img, lbl, (img_f, lbl_f)
49 | 
50 |     def __len__(self):
51 |         return len(self.image_f)
52 | 


--------------------------------------------------------------------------------
/celldetection/data/datasets/bbbc039.py:
--------------------------------------------------------------------------------
  1 | from os.path import join
  2 | from os import makedirs
  3 | from imageio import imread
  4 | from skimage import measure
  5 | import numpy as np
  6 | from ..cpn import labels2contour_list as labels2contours
  7 | from ...visualization.images import show_detection
  8 | import torchvision
  9 | 
 10 | __all__ = ['download_bbbc039', 'BBBC039Train', 'BBBC039Test', 'BBBC039Val']
 11 | 
 12 | 
 13 | def read_all(directory, filename):
 14 |     if filename is None or directory is None:
 15 |         return None
 16 |     with open(join(directory, filename), 'r') as f:
 17 |         return [i.strip() for i in f.readlines()]
 18 | 
 19 | 
 20 | def load(images_directory, masks_directory, names, **label_kwargs):
 21 |     if None in (images_directory, masks_directory):
 22 |         return None, None, None
 23 |     images = [imread(join(images_directory, f.replace('.png', '.tif'))) for f in names]
 24 |     masks = [imread(join(masks_directory, f)) for f in names]
 25 |     return images, masks, [measure.label(m[:, :, 0], **label_kwargs) for m in masks]
 26 | 
 27 | 
 28 | def download_bbbc039(directory):
 29 |     """Download BBBC039.
 30 | 
 31 |     Download and extract the BBBC039 dataset to given directory.
 32 | 
 33 |     References:
 34 |         https://bbbc.broadinstitute.org/BBBC039
 35 | 
 36 |     Args:
 37 |         directory: Root directory.
 38 | 
 39 |     """
 40 |     makedirs(directory, exist_ok=True)
 41 |     for url in [
 42 |         'https://data.broadinstitute.org/bbbc/BBBC039/images.zip',
 43 |         'https://data.broadinstitute.org/bbbc/BBBC039/metadata.zip',
 44 |         'https://data.broadinstitute.org/bbbc/BBBC039/masks.zip'
 45 |     ]:
 46 |         torchvision.datasets.utils.download_and_extract_archive(url, directory)
 47 | 
 48 | 
 49 | class _BBBC039:
 50 |     def __init__(self, directory, download, mode: str):
 51 |         assert mode in ('train', 'test', 'val')
 52 | 
 53 |         meta_directory = join(directory, 'metadata')
 54 |         masks_directory = join(directory, 'masks')
 55 |         images_directory = join(directory, 'images')
 56 | 
 57 |         if download:
 58 |             download_bbbc039(directory)
 59 | 
 60 |         self.names = read_all(meta_directory, {
 61 |             'train': 'training.txt',
 62 |             'val': 'validation.txt',
 63 |             'test': 'test.txt'
 64 |         }[mode])
 65 | 
 66 |         self.images, self.masks, self.labels = load(images_directory, masks_directory, self.names)
 67 | 
 68 |     def plot(self, num=1, figsize=(20, 15)):
 69 |         for i in np.random.randint(0, len(self), num):
 70 |             show_detection(image=self.images[i], contours=labels2contours(self.labels[i]),
 71 |                            figsize=figsize, contour_linestyle='-')
 72 | 
 73 |     def __getitem__(self, item):
 74 |         return self.names[item], self.images[item], self.masks[item], self.labels[item]
 75 | 
 76 |     def __len__(self):
 77 |         return len(self.images)
 78 | 
 79 | 
 80 | class BBBC039Train(_BBBC039):
 81 |     def __init__(self, directory, download=False):
 82 |         """BBBC039 Train.
 83 | 
 84 |         Training split of the BBBC039 dataset.
 85 | 
 86 |         References:
 87 |             https://bbbc.broadinstitute.org/BBBC039
 88 | 
 89 |         Args:
 90 |             directory: Root directory.
 91 |             download: Whether to download the dataset.
 92 |         """
 93 |         super().__init__(directory, download=download, mode='train')
 94 | 
 95 | 
 96 | class BBBC039Val(_BBBC039):
 97 |     def __init__(self, directory, download=False):
 98 |         """BBBC039 Validation.
 99 | 
100 |         Validation split of the BBBC039 dataset.
101 | 
102 |         References:
103 |             https://bbbc.broadinstitute.org/BBBC039
104 | 
105 |         Args:
106 |             directory: Root directory.
107 |             download: Whether to download the dataset.
108 |         """
109 |         super().__init__(directory, download=download, mode='val')
110 | 
111 | 
112 | class BBBC039Test(_BBBC039):
113 |     def __init__(self, directory, download=False):
114 |         """BBBC039 Test.
115 | 
116 |         Test split of the BBBC039 dataset.
117 | 
118 |         References:
119 |             https://bbbc.broadinstitute.org/BBBC039
120 | 
121 |         Args:
122 |             directory: Root directory.
123 |             download: Whether to download the dataset.
124 |         """
125 |         super().__init__(directory, download=download, mode='test')
126 | 


--------------------------------------------------------------------------------
/celldetection/data/datasets/bbbc041.py:
--------------------------------------------------------------------------------
  1 | from torchvision.datasets import utils
  2 | from os.path import join, basename, isfile
  3 | from os import makedirs, listdir
  4 | from imageio import imread
  5 | import numpy as np
  6 | import json
  7 | from matplotlib import pyplot as plt
  8 | import shutil
  9 | from ...visualization.images import plot_box, plot_text
 10 | 
 11 | __all__ = ['download_bbbc041', 'BBBC041Train', 'BBBC041Test']
 12 | 
 13 | 
 14 | def download_bbbc041(directory, url='https://data.broadinstitute.org/bbbc/BBBC041/malaria.zip'):
 15 |     """Download BBBC041.
 16 | 
 17 |     Download and extract the BBBC041 dataset to given directory.
 18 | 
 19 |     References:
 20 |         https://bbbc.broadinstitute.org/BBBC041
 21 | 
 22 |     Args:
 23 |         directory: Root directory.
 24 |         url: Download URL (this dataset is distributed in a single zip file).
 25 | 
 26 |     """
 27 |     makedirs(directory, exist_ok=True)
 28 |     utils.download_and_extract_archive(url, directory)
 29 |     malaria_dir = join(directory, 'malaria')
 30 |     for src in listdir(malaria_dir):
 31 |         src = join(malaria_dir, src)
 32 |         dst = join(directory, basename(src))
 33 |         print(src, '->', dst)
 34 |         shutil.move(src, dst)
 35 | 
 36 | 
 37 | class _BBBC041:
 38 |     def __init__(self, directory, download, mode: str):
 39 |         assert mode in ('train', 'test')  # , 'val'
 40 | 
 41 |         json_file = join(directory, {
 42 |             'train': 'training.json',
 43 |             # 'val': 'validation.json',
 44 |             'test': 'test.json'
 45 |         }[mode])
 46 | 
 47 |         if download and not isfile(json_file):
 48 |             download_bbbc041(directory)
 49 | 
 50 |         with open(json_file, 'r') as f:
 51 |             meta = json.load(f)
 52 | 
 53 |         self.filenames = []
 54 |         self.images = []
 55 |         self.boxes = []
 56 |         self.categories = []
 57 |         for item in meta:
 58 |             image_item = item['image']
 59 |             object_items = item['objects']
 60 |             pn = image_item['pathname']
 61 |             if pn[0] == '/':
 62 |                 pn = pn[1:]
 63 |             self.filenames.append(join(directory, pn))
 64 |             self.images.append(imread(self.filenames[-1]))
 65 |             boxes = []
 66 |             categories = []
 67 |             for box in object_items:
 68 |                 bb = box['bounding_box']
 69 |                 # Boxes as x_min, y_min, x_max, y_max
 70 |                 boxes.append([bb['minimum']['c'], bb['minimum']['r'], bb['maximum']['c'], bb['maximum']['r']])
 71 |                 categories.append(box['category'])
 72 |             self.boxes.append(boxes)
 73 |             self.categories.append(categories)
 74 | 
 75 |     def plot(self, item=None, num=1, figsize=(16, 9)):
 76 |         if item is None:
 77 |             item = np.random.randint(0, len(self), num)
 78 |         else:
 79 |             item = item,
 80 |         for i in item:
 81 |             name, img, boxes, c = self[i]
 82 |             plt.figure(None, figsize)
 83 |             plt.imshow(img)
 84 |             for j, box in enumerate(boxes):
 85 |                 plot_box(*box)
 86 |                 plot_text(c[j], box[0] + 5, box[1] - 5)
 87 |             plt.show()
 88 | 
 89 |     def __len__(self):
 90 |         return len(self.images)
 91 | 
 92 |     def __getitem__(self, item):
 93 |         return self.filenames[item], self.images[item], self.boxes[item], self.categories[item]
 94 | 
 95 | 
 96 | class BBBC041Train(_BBBC041):
 97 |     def __init__(self, directory, download=False):
 98 |         """BBBC041 Train data.
 99 | 
100 |         References:
101 |             https://bbbc.broadinstitute.org/BBBC041
102 | 
103 |         Args:
104 |             directory: Data directory.
105 |             download: Whether to download data.
106 |         """
107 |         super().__init__(directory, download=download, mode='train')
108 | 
109 | 
110 | class BBBC041Test(_BBBC041):
111 |     def __init__(self, directory, download=False):
112 |         """BBBC041 Test data.
113 | 
114 |         References:
115 |             https://bbbc.broadinstitute.org/BBBC041
116 | 
117 |         Args:
118 |             directory: Data directory.
119 |             download: Whether to download data.
120 |         """
121 |         super().__init__(directory, download=download, mode='test')
122 | 


--------------------------------------------------------------------------------
/celldetection/data/datasets/generic.py:
--------------------------------------------------------------------------------
 1 | from h5py import File
 2 | from os.path import isfile
 3 | 
 4 | __all__ = ['GenericH5']
 5 | 
 6 | 
 7 | class GenericH5:
 8 |     def __init__(self, filenames, keys, cache=False):
 9 |         """Generic hdf5 dataset.
10 | 
11 |         Basic interface for a set of hdf5 files.
12 | 
13 |         Examples:
14 |             >>> import celldetection as cd
15 |             ... from glob import glob
16 |             ... data = GenericH5(glob('./files/*.h5'), ('image', 'labels'))
17 |             ... name, (img, lbl) = data[0]
18 | 
19 |         Args:
20 |             filenames: Hdf5 file names.
21 |             keys: Keys. Used to retrieve data from each file.
22 |             cache: Whether to hold data in memory.
23 |         """
24 |         self._filenames = list(filenames)
25 |         self._filenames.sort()
26 |         self.content = list(self._filenames)
27 |         for f in self.content:
28 |             if not isfile(f):
29 |                 raise FileNotFoundError(f'File not found: {f}')
30 |         self._single = isinstance(keys, str)
31 |         self.keys = (keys,) if self._single else keys
32 |         self.cache = cache
33 | 
34 |     def __getitem__(self, item):
35 |         it = self.content[item]
36 |         if isinstance(it, str):
37 |             with File(it, 'r') as h:
38 |                 it = [h[k][:] for k in self.keys]
39 |             if self._single:
40 |                 it, = it
41 |             if self.cache:
42 |                 self.content[item] = it
43 |         return self._filenames[item], it
44 | 
45 |     def __len__(self):
46 |         return len(self.content)
47 | 
48 |     def __str__(self):
49 |         return f'{len(self)} hdf5 files'
50 | 
51 |     __repr__ = __str__
52 | 


--------------------------------------------------------------------------------
/celldetection/data/datasets/synth.py:
--------------------------------------------------------------------------------
 1 | from .generic import GenericH5
 2 | from glob import glob
 3 | from os.path import join, isdir
 4 | from os import makedirs
 5 | from torchvision import datasets
 6 | 
 7 | __all__ = ['download_synth', 'SynthTrain', 'SynthVal', 'SynthTest']
 8 | 
 9 | 
10 | def download_synth(directory, url='https://celldetection.org/data/synth.zip'):
11 |     """Download Synth.
12 | 
13 |     Download and extract the Synth dataset to given directory.
14 | 
15 |     Args:
16 |         directory: Root directory.
17 |         url: Download URL (this dataset is distributed in a single zip file).
18 | 
19 |     """
20 |     makedirs(directory, exist_ok=True)
21 |     datasets.utils.download_and_extract_archive(url, directory)
22 | 
23 | 
24 | class _Synth(GenericH5):
25 |     def __init__(self, directory, download, mode: str, cache=True):
26 |         assert mode in ('train', 'val', 'test')
27 | 
28 |         if download and not isdir(join(directory, mode)):
29 |             download_synth(directory)
30 | 
31 |         super().__init__(
32 |             filenames=glob(join(directory, mode, '*.h5')),
33 |             keys=('image', 'labels'), cache=cache
34 |         )
35 | 
36 | 
37 | class SynthTrain(_Synth):
38 |     def __init__(self, directory, download=False, cache=True):
39 |         """Synth Train data.
40 | 
41 |         Args:
42 |             directory: Data directory. E.g. ``directory='data/synth'``.
43 |             download: Whether to download data.
44 |             cache: Whether to hold data in memory.
45 |         """
46 |         super().__init__(directory, download, 'train', cache=cache)
47 | 
48 | 
49 | class SynthVal(_Synth):
50 |     def __init__(self, directory, download=False, cache=True):
51 |         super().__init__(directory, download, 'val', cache=cache)
52 | 
53 |     __init__.__doc__ = SynthTrain.__init__.__doc__.replace('Synth Train data.', 'Synth Val data.')
54 | 
55 | 
56 | class SynthTest(_Synth):
57 |     def __init__(self, directory, download=False, cache=True):
58 |         super().__init__(directory, download, 'test', cache=cache)
59 | 
60 |     __init__.__doc__ = SynthTrain.__init__.__doc__.replace('Synth Train data.', 'Synth Test data.')
61 | 


--------------------------------------------------------------------------------
/celldetection/data/segmentation.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | from skimage import morphology
  3 | import cv2
  4 | from .misc import rgb_to_scalar
  5 | 
  6 | __all__ = ['remove_partials_', 'fill_label_gaps_', 'filter_instances_', 'relabel_', 'stack_labels',
  7 |            'unary_masks2labels', 'boxes2masks', 'fill_padding_', 'remove_padding']
  8 | 
  9 | 
 10 | def remove_partials_(label_stack, border=1, constant=-1):
 11 |     if border < 1:
 12 |         return label_stack, None
 13 |     bad_labels = set(np.unique(label_stack[:, :border]))
 14 |     bad_labels.update(np.unique(label_stack[:, -border:]))
 15 |     bad_labels.update(np.unique(label_stack[:border, :]))
 16 |     bad_labels.update(np.unique(label_stack[-border:, :]))
 17 |     mask = np.isin(label_stack, list(bad_labels - {0}))
 18 |     label_stack[mask] = constant
 19 |     return label_stack, mask
 20 | 
 21 | 
 22 | def fill_label_gaps_(labels):
 23 |     """Fill label gaps.
 24 | 
 25 |     Ensure that labels greater zero are within interval [1, num_unique_labels_greater_zero].
 26 |     Works fast if gaps are unlikely, slow otherwise. Alternatively consider using np.vectorize.
 27 |     Labels <= 0 are preserved as is.
 28 | 
 29 |     Args:
 30 |         labels:
 31 | 
 32 |     Returns:
 33 | 
 34 |     """
 35 |     uni = np.unique(labels)
 36 |     uniques = list(set(uni) - set(uni[uni <= 0]))  # ignore zeros and negative values
 37 |     uniques.sort()
 38 |     gaps = list(set(range(1, len(uniques) + 1)) - set(uniques))
 39 |     while len(gaps) > 0:
 40 |         labels[labels == uniques.pop()] = gaps.pop()
 41 | 
 42 | 
 43 | def fill_padding_(inputs, padding: int, constant=-1, preserve_existing=True, axes=(0, 1)):
 44 |     if padding <= 0:
 45 |         return None
 46 |     if isinstance(inputs, (list, tuple)):
 47 |         [fill_padding_(i, padding, constant) for i in inputs]
 48 |         return
 49 |     for ax in axes:
 50 |         if ax < 0:
 51 |             ax = ax + inputs.ndim
 52 |         for sl in (slice(0, padding), slice(-padding, None)):  # , (slice(None), slice(0, padding)), (slice(None), slice(-padding, None))
 53 |             sl = (slice(None),) * ax + (sl,)
 54 |             if preserve_existing:
 55 |                 mask = ~inputs[sl].any(-1)
 56 |             else:
 57 |                 mask = slice(None)
 58 |             inputs[sl][mask] = constant
 59 | 
 60 | 
 61 | def remove_padding(inputs, padding: int):
 62 |     if isinstance(inputs, (list, tuple)):
 63 |         return [remove_padding(i, padding) for i in inputs]
 64 |     return inputs[padding:-padding, padding:-padding]
 65 | 
 66 | 
 67 | def filter_instances_(labels, partials=True, partials_border=1, min_area=4, max_area=None, constant=-1,
 68 |                       continuous=True):
 69 |     """Filter instances from label image.
 70 | 
 71 |     Note:
 72 |         Filtered instance labels are set to `constant`.
 73 |         Labels might not be continuous afterwards.
 74 | 
 75 |     Args:
 76 |         labels:
 77 |         partials:
 78 |         partials_border:
 79 |         min_area:
 80 |         max_area:
 81 |         constant:
 82 |         continuous:
 83 | 
 84 |     Returns:
 85 | 
 86 |     """
 87 |     if partials:
 88 |         remove_partials_(labels, border=partials_border, constant=constant)
 89 | 
 90 |     if max_area is not None or min_area is not None:
 91 |         # Unique labels with counts: 13.5 ms ± 242 µs ((576, 576, 3), 763 instances)
 92 |         uni_labels, uni_counts = np.unique(labels, return_counts=True)
 93 |         uni_labels, uni_counts = uni_labels[1:], uni_counts[1:]
 94 |         bad_labels = []
 95 |         if max_area:
 96 |             bad_labels += list(uni_labels[uni_counts > max_area].ravel())
 97 |         if min_area:
 98 |             bad_labels += list(uni_labels[uni_counts < min_area].ravel())
 99 |         for label in bad_labels:
100 |             labels[labels == label] = constant
101 | 
102 |     if continuous:
103 |         fill_label_gaps_(labels)
104 | 
105 | 
106 | def relabel_(label_stack, axis=2):
107 |     """Relabel.
108 | 
109 |     Inplace relabeling of a label stack.
110 |     After applying this op the labels in label_stack are continuous, starting at 1.
111 |     Negative labels remain untouched.
112 | 
113 |     Notes:
114 |         - Uses label function from sklearn.morphology
115 | 
116 |     Args:
117 |         label_stack: Array[height, width, channels].
118 |         axis: Channel axis.
119 |     """
120 |     assert label_stack.ndim == 3
121 |     neg_m = label_stack < 0
122 |     cur_max = 0
123 |     for channel in range(label_stack.shape[axis]):
124 |         stack_ = morphology.label(label_stack[(slice(None),) * axis + (channel,)])
125 |         for u in set(np.unique(stack_)) - {0}:
126 |             mask = stack_ == u
127 |             if np.any(np.logical_and(mask, neg_m[(slice(None),) * axis + (channel,)])):
128 |                 continue
129 |             cur_max += 1
130 |             label_stack[mask, channel] = cur_max
131 | 
132 | 
133 | def stack_labels(*maps, axis=2, dtype='int32', relabel=True):
134 |     """Stack labels.
135 | 
136 |     Args:
137 |         *maps: List[Union[Array[h, w], Array[h, w, 3]]. Grayscale or rgb label maps.
138 |             Rgb labels are assumed to encode labels with color and are converted to grayscale labels before stacking.
139 |         axis: Stacking axis.
140 |         dtype: Output data type.
141 |         relabel: Whether to assign new labels or not.
142 | 
143 |     Returns:
144 |         Array[h, w, c]. Label image.
145 |     """
146 |     maps = [(rgb_to_scalar(m, dtype=dtype) if (m.ndim == 3 and m.shape[2] == 3) else m.astype(dtype)) for m in maps]
147 |     stack = np.stack(maps, axis=axis)
148 |     if relabel:
149 |         relabel_(stack, axis)
150 |     return stack
151 | 
152 | 
153 | def unary_masks2labels(unary_masks, transpose=True):
154 |     """Unary masks to labels.
155 | 
156 |     Args:
157 |         unary_masks: List[Array[height, width]] or Array[num_objects, height, width]
158 |             List of masks. Each mask is assumed to contain exactly one object.
159 |         transpose: If True label images are in channels last format, otherwise channels first.
160 | 
161 |     Returns:
162 |         Label image. Array[height, width, num_objects] if `transpose` else Array[num_objects, height, width].
163 |     """
164 |     lbl = (unary_masks > 0) * np.arange(1, len(unary_masks) + 1)[:, None, None]
165 |     if transpose:
166 |         lbl = lbl.transpose((1, 2, 0))
167 |     return lbl
168 | 
169 | 
170 | def boxes2masks(boxes, size):
171 |     masks = [None] * len(boxes)
172 |     for idx, b in enumerate(boxes):
173 |         mask = masks[idx] = np.zeros(size, dtype='uint8')
174 |         xmin, ymin, xmax, ymax = b
175 |         cv2.rectangle(mask, (xmin, ymin), (xmax, ymax), idx + 1, thickness=-1)
176 |     return masks
177 | 


--------------------------------------------------------------------------------
/celldetection/data/toydata.py:
--------------------------------------------------------------------------------
 1 | import cv2
 2 | import numpy as np
 3 | 
 4 | CLASS_NAMES_GEOMETRIC = {
 5 |     1: 'rectangle',
 6 |     2: 'triangle',
 7 |     3: 'ellipse'
 8 | }
 9 | 
10 | 
11 | def random_triangle(image, mask, x, y, color, radius_range=(3, 28)):
12 |     a, b, c, d, e = np.random.randint(*radius_range, size=5)
13 |     triangle_cnt = np.array([[x, y - a], [x + b, y + c], [x - d, y + e]])
14 |     cv2.drawContours(image, [triangle_cnt], 0, color, -1)
15 |     cv2.drawContours(mask, [triangle_cnt], 0, 1, -1)
16 |     return image, mask
17 | 
18 | 
19 | def random_ellipse(image, mask, x, y, color, radius_range=(3, 28)):
20 |     rh, rw = np.random.randint(*radius_range, size=2)
21 |     angle = np.random.randint(0, 360)
22 |     cv2.ellipse(image, (x, y), axes=(rh, rw), angle=angle, startAngle=0, endAngle=360, color=color, thickness=-1)
23 |     cv2.ellipse(mask, (x, y), axes=(rh, rw), angle=angle, startAngle=0, endAngle=360, color=1, thickness=-1)
24 |     return image, mask
25 | 
26 | 
27 | def random_rectangle(image, mask, x, y, color, radius_range=(3, 28)):
28 |     rh, rw = np.random.randint(*radius_range, size=2)
29 |     angle = np.random.randint(0, 360)
30 |     tmp = np.zeros(image.shape[:2], dtype=np.uint8)
31 |     cv2.ellipse(tmp, (x, y), axes=(rh, rw), angle=angle, startAngle=0, endAngle=360, color=1, thickness=-1)
32 |     c = cv2.findContours(tmp, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
33 |     c = c[0] if len(c) == 2 else c[1]
34 |     box = np.int0(cv2.boxPoints(cv2.minAreaRect(c[0])))
35 |     cv2.drawContours(image, [box], 0, color, -1)
36 |     cv2.drawContours(mask, [box], 0, 1, -1)
37 |     return image, mask
38 | 
39 | 
40 | def random_circle(image, mask, x, y, color, radius_range=(3, 28)):
41 |     rad = np.random.randint(*radius_range)
42 |     cv2.circle(image, center=(x, y), radius=rad, color=color, thickness=-1)
43 |     cv2.circle(mask, center=(x, y), radius=rad, color=1, thickness=-1)
44 |     return image, mask
45 | 
46 | 
47 | def random_geometric_objects(height=256, width=256, radius_range=(3, 28), intensity_range=(0, 180), margin=13):
48 |     img = np.zeros((height, width, 3), dtype='uint8') + 255
49 |     mrad = np.max(radius_range)
50 |     xa, xb = margin + mrad, img.shape[1] - mrad - margin
51 |     ya, yb = margin + mrad, img.shape[0] - mrad - margin
52 |     step = int(mrad * 1.5)
53 |     xy = np.mgrid[xa:xb:step, ya:yb:step].reshape((2, -1))
54 |     xy_rad = np.random.randint(*radius_range, xy.shape[1:])
55 |     masks = []
56 |     labels = []
57 |     classes = []
58 |     for idx, (x, y) in enumerate(xy.T):
59 |         rad = xy_rad[idx]
60 |         x = np.clip(int(x), 0, img.shape[1]) + np.random.randint(0, int(rad * .5))
61 |         y = np.clip(int(y), 0, img.shape[0]) + np.random.randint(0, int(rad * .5))
62 |         color = np.random.randint(*intensity_range, 3)
63 |         variant = np.random.choice([1, 2, 3])
64 |         mask = np.zeros((height, width), dtype='uint8')
65 |         classes.append(variant)
66 |         if variant == 1:
67 |             img, mask = random_rectangle(img, mask, x, y, color.tolist(), radius_range=radius_range)
68 |         elif variant == 2:
69 |             img, mask = random_triangle(img, mask, x, y, color.tolist(), radius_range=radius_range)
70 |         else:
71 |             img, mask = random_ellipse(img, mask, x, y, color.tolist(), radius_range=radius_range)
72 |         masks.append(mask)
73 |         label = np.copy(mask).astype('int')
74 |         label[label > .5] += idx
75 |         labels.append(label)
76 |     return img, np.array(masks), np.stack(labels, -1), np.array(classes)
77 | 


--------------------------------------------------------------------------------
/celldetection/data/transforms.py:
--------------------------------------------------------------------------------
 1 | from abc import abstractmethod
 2 | from pytorch_lightning.core.mixins import HyperparametersMixin
 3 | import numpy as np
 4 | import cv2
 5 | 
 6 | from .misc import random_crop, normalize_percentile
 7 | 
 8 | __all__ = ['Transforms', 'BasicTransforms']
 9 | 
10 | 
11 | class Transforms(HyperparametersMixin):
12 |     def __init__(self, **kwargs):
13 |         """Transforms.
14 | 
15 |         Defines all transforms for all stages.
16 |         """
17 |         super().__init__()
18 |         self.save_hyperparameters()
19 | 
20 |     @abstractmethod
21 |     def static(self, **kwargs):
22 |         # All kwargs from __init__ file available via `self.hparams.setting_name`.
23 |         return kwargs
24 | 
25 |     def fit(self, **kwargs):
26 |         kwargs = self.static(**kwargs)
27 |         return kwargs
28 | 
29 |     def validate(self, **kwargs):
30 |         kwargs = self.static(**kwargs)
31 |         return kwargs
32 | 
33 |     def test(self, **kwargs):
34 |         kwargs = self.static(**kwargs)
35 |         return kwargs
36 | 
37 |     def predict(self, **kwargs):
38 |         kwargs = self.static(**kwargs)
39 |         return kwargs
40 | 
41 |     def __call__(self, stage='fit', **kwargs):
42 |         assert stage in ('fit', 'validate', 'predict', 'test')
43 |         return getattr(self, stage)(**kwargs)
44 | 
45 | 
46 | class BasicTransforms(Transforms):
47 |     def static(self, crop=False, **kwargs):
48 |         image = kwargs['image']
49 |         labels = kwargs['labels']
50 |         if crop and self.hparams.crop_size:
51 |             image, labels = random_crop((image, labels), size=(self.hparams.crop_size,) * 2)
52 | 
53 |         if image.dtype != np.uint8:
54 |             image = normalize_percentile(image)
55 |         if image.ndim == 2 or (image.ndim == 3 and image.shape[-1] == 1):
56 |             image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
57 | 
58 |         kwargs['image'] = image / 255
59 |         kwargs['labels'] = labels
60 |         return kwargs
61 | 
62 |     def fit(self, **kwargs):
63 |         kwargs = self.static(crop=True, **kwargs)
64 |         return kwargs
65 | 
66 |     def validate(self, **kwargs):
67 |         kwargs = self.static(**kwargs)
68 |         return kwargs
69 | 
70 |     def test(self, **kwargs):
71 |         kwargs = self.static(**kwargs)
72 |         return kwargs
73 | 
74 |     def predict(self, **kwargs):
75 |         kwargs = self.static(**kwargs)
76 |         return kwargs
77 | 


--------------------------------------------------------------------------------
/celldetection/models/__init__.py:
--------------------------------------------------------------------------------
 1 | """
 2 | This submodule contains model definitions and common modules.
 3 | """
 4 | from .commons import *
 5 | from .unet import *
 6 | from .resnet import *
 7 | from .cpn import *
 8 | from .fpn import *
 9 | from .inference import *
10 | from .mobilenetv3 import *
11 | from .normalization import *
12 | from .filters import *
13 | from .features import *
14 | from .ppm import *
15 | from .loss import *
16 | from .convnext import *
17 | from .convnextv2 import *
18 | from .densenet import *
19 | from .mamba import *
20 | from .manet import *
21 | from .smp import *
22 | from .timmodels import *
23 | from .lightning_cpn import *
24 | 
25 | from ..util.util import NormProxy
26 | 


--------------------------------------------------------------------------------
/celldetection/models/features.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import torch.nn as nn
 3 | import numpy as np
 4 | from ..models.filters import EdgeFilter2d, GaussianFilter2d
 5 | from ..ops.features import texture_filter
 6 | 
 7 | __all__ = ['MultiscaleBasicFeatures']
 8 | 
 9 | 
10 | class MultiscaleBasicFeatures(nn.Module):
11 |     def __init__(
12 |             self,
13 |             in_channels,
14 |             intensity=True,
15 |             edges=True,
16 |             texture=True,
17 |             sigma_min=.5,
18 |             sigma_max=16,
19 |             num_sigma=None,
20 |             padding_mode='zeros',
21 |             trainable=False
22 |     ):
23 |         """Multiscale Basic Features.
24 | 
25 |         Notes:
26 |             - Multiscale aspect is implemented using Gaussian blur over different scales.
27 | 
28 |         References:
29 |             - https://github.com/scikit-image/scikit-image/blob/91f26e9c0a00522137ddb638fa8362ec01d2d7b2/skimage/feature/_basic_features.py#L101
30 | 
31 |         Args:
32 |             in_channels: Number of input channels.
33 |             intensity: Whether to include intensities to feature set.
34 |             edges: Whether to include results of edge filters to feature set.
35 |             texture: Whether to include eigenvalues of the Hessian matrix after Gaussian blurring to the feature set.
36 |             sigma_min: Smallest sigma for the Gaussian blur kernel.
37 |             sigma_max: Largest sigma for the Gaussian blur kernel.
38 |             num_sigma: Number of different scales.
39 |             padding_mode: One of ``'zeros'``, ``'reflect'``, ``'replicate'`` or ``'circular'``.
40 |                 Default: ``'zeros'``.
41 |             trainable: Whether the kernels should be trainable.
42 |         """
43 |         super().__init__()
44 |         self.intensity = intensity
45 |         self.edges = edges
46 |         self.texture = texture
47 |         self.num_sigma = num_sigma or int(np.log2(sigma_max) - np.log2(sigma_min) + 1)
48 |         self.sigmas = np.logspace(np.log2(sigma_min), np.log2(sigma_max), num=self.num_sigma, base=2, endpoint=True)
49 |         self.gaussians = nn.ModuleList()
50 |         for sigma in self.sigmas:
51 |             kernel_size = int(2 * np.ceil(3 * sigma) + 1)
52 |             g = GaussianFilter2d(in_channels, kernel_size, sigma=sigma, padding=kernel_size // 2,
53 |                                  padding_mode=padding_mode, trainable=trainable)
54 |             self.gaussians.add_module(f"sigma{round(sigma, 4)}".replace('.', '_'), g)
55 |         if self.edges:
56 |             self.edge_filter = EdgeFilter2d(in_channels, padding=1, padding_mode=padding_mode, magnitude=True,
57 |                                             trainable=trainable)
58 | 
59 |     def forward(self, x):
60 |         results = []
61 |         for gaussian in self.gaussians:
62 |             filtered = gaussian(x)
63 |             if self.intensity:
64 |                 results.append(filtered)
65 |             if self.edges:
66 |                 results.append(self.edge_filter(filtered))
67 |             if self.texture:
68 |                 results.append(texture_filter(filtered))
69 |         return torch.concat(results, 1)
70 | 


--------------------------------------------------------------------------------
/celldetection/models/hosted.py:
--------------------------------------------------------------------------------
1 | HOST_TEMPLATE = 'https://celldetection.org/torch/models/{name}'
2 | HOSTED_MODELS = dict(
3 |     ginoro='ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c',
4 | )
5 | 


--------------------------------------------------------------------------------
/celldetection/models/inference.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from ..util.util import asnumpy
 3 | 
 4 | __all__ = ['Inference']
 5 | 
 6 | 
 7 | class Inference:
 8 |     def __init__(self, model, device=None, amp=False, transforms=None):
 9 |         self.transforms = transforms
10 |         self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
11 |         self.model = model
12 |         self.model.eval()
13 |         self.model.requires_grad_(False)
14 |         self.use_amp = amp
15 | 
16 |     def __call__(self, inputs):
17 |         if self.transforms is not None:
18 |             inputs = self.transforms(inputs)
19 |         inputs = torch.as_tensor(inputs, device=self.device, dtype=torch.float32)
20 |         if inputs.ndim == 2:
21 |             inputs = inputs[None]
22 |         if inputs.ndim == 3:
23 |             inputs = inputs[None]
24 |         with torch.cuda.amp.autocast(self.use_amp):
25 |             out = self.model(inputs)
26 |         return asnumpy(out)
27 | 


--------------------------------------------------------------------------------
/celldetection/models/lightning_cpn.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from typing import Any, Dict, List, Union, Sequence
  3 | from ..util.util import asnumpy, get_tiling_slices
  4 | from torch import Tensor, nn
  5 | from collections import OrderedDict
  6 | from lightning_fabric.utilities.rank_zero import rank_zero_only
  7 | from ..data.instance_eval import LabelMatcher
  8 | from ..data.cpn import contours2labels
  9 | from ..data.misc import channels_first2channels_last
 10 | from ..ops.cpn import remove_border_contours
 11 | import numpy as np
 12 | from torchvision.ops.boxes import remove_small_boxes, nms
 13 | import pytorch_lightning as pl
 14 | from .lightning_base import LitBase
 15 | 
 16 | __all__ = ['LitCpn']
 17 | 
 18 | 
 19 | STEP_OUTPUT = Union[Tensor, Dict[str, Any]]
 20 | EPOCH_OUTPUT = List[STEP_OUTPUT]
 21 | 
 22 | 
 23 | class LitCpn(LitBase):
 24 |     def __init__(
 25 |             self,
 26 |             model: Union[str, dict, nn.Module],
 27 |             losses_prog_bar=True,
 28 |             optimizer=None,
 29 |             scheduler=None,
 30 |             scheduler_conf=None,
 31 |             warmup_steps=512,
 32 |             lr_scale='sqrt',
 33 |             weight_decay_scale=None,
 34 |             **kwargs
 35 |     ):
 36 |         kwargs['val_hparams'] = {
 37 |             'score_thresh': [.5, .86, .88, .9, .92],
 38 |             **(kwargs.get('val_hparams') or {})
 39 |         }
 40 | 
 41 |         super().__init__(
 42 |             model=model,
 43 |             losses_prog_bar=losses_prog_bar,
 44 |             optimizer=optimizer,
 45 |             scheduler=scheduler,
 46 |             scheduler_conf=scheduler_conf,
 47 |             warmup_steps=warmup_steps,
 48 |             lr_scale=lr_scale,
 49 |             weight_decay_scale=weight_decay_scale,
 50 |             **kwargs
 51 |         )
 52 | 
 53 |         self.scores_lower_bound_key = 'scores_lower_bound'
 54 |         self.scores_upper_bound_key = 'scores_upper_bound'
 55 | 
 56 |     def _training_step(self, batch: dict, batch_idx: int) -> STEP_OUTPUT:
 57 |         inputs = batch.pop(self.inputs_key)
 58 |         outputs: dict = self.model(inputs, targets=batch, rank=self.global_rank)
 59 |         return outputs
 60 | 
 61 |     def _evaluation_log(self, prefix, batch_idx, inputs, contours, global_step):
 62 |         self.log_contour_figures(tag=f'{prefix}/batch{batch_idx}', inputs=inputs, contours=contours,
 63 |                                  global_step=global_step)
 64 | 
 65 |     def _evaluation_step(self, batch: dict, batch_idx: int, prefix: str, hparams_key, inputs, indices, matches,
 66 |                          log_step: bool):
 67 |         with torch.no_grad():
 68 |             outputs: dict = self(inputs, targets=batch, feature_cache=True)  # todo: val loss
 69 |         contours = asnumpy(outputs['contours'])
 70 |         targets = asnumpy(batch[self.targets_key])
 71 |         if log_step and self._log_figures:
 72 |             self._evaluation_log(prefix, batch_idx, inputs=inputs, contours=contours, global_step=self.global_step)
 73 |         matches[hparams_key] = matches_ = matches.get(hparams_key, {})
 74 | 
 75 |         for i, (cons, target, index) in enumerate(zip(contours, targets, indices)):
 76 |             prediction = contours2labels(cons, size=inputs[i].shape[-2:], initial_depth=3)
 77 |             target = channels_first2channels_last(target)
 78 |             matches_[index] = LabelMatcher(prediction, target, zero_division=self._eval_zero_division)
 79 | 
 80 |     def _predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
 81 |         # Note: Remove everything that is not required for postprocessing from batch (e.g. via batch.pop).
 82 |         inputs = batch.pop(self.inputs_key)  # remove inputs
 83 |         assert inputs is not None
 84 |         scores_lower_bound = batch.pop(self.scores_lower_bound_key, None)
 85 |         scores_upper_bound = batch.pop(self.scores_upper_bound_key, None)
 86 |         return self(inputs, scores_upper_bound=scores_upper_bound, scores_lower_bound=scores_lower_bound, **batch)
 87 | 
 88 |     def forward_tiled(
 89 |             self,
 90 |             inputs: Tensor,
 91 |             crop_size: Union[int, Sequence[int]] = 1024,
 92 |             stride: Union[int, Sequence[int]] = 512,
 93 |             **kwargs
 94 |     ) -> Dict[str, List[Tensor]]:
 95 |         assert np.array(crop_size) <= np.array(stride) * 2
 96 |         slices, slices_by_dim = get_tiling_slices(inputs.shape[2:], crop_size, stride)  # ordered
 97 |         prod = np.prod(slices_by_dim)
 98 |         results: List[List[Dict[str, Tensor]]] = [[None] * prod for _ in torch.arange(0, inputs.shape[0])]
 99 |         h_tiles, w_tiles = slices_by_dim
100 |         device = self.device
101 |         kwargs.pop('max_imsize', None)
102 |         extra_keys = kwargs.get('extra_keys', ())  # extra output keys
103 |         extra_nms = kwargs.get('extra_nms', {})  # specify which extra outputs need to be filtered by nms
104 |         border_removal = kwargs.get('border_removal', 6)
105 |         box_min_size = kwargs.get('min_box_size', 1.)
106 |         nms_thresh = kwargs.get('nms_thresh', self.model.__dict__.get('nms_thresh', None))
107 |         inputs_mask = kwargs.get('inputs_mask')
108 |         assert nms_thresh is not None, 'Could not retrieve nms_thresh from model. Please specify it in forward method.'
109 |         targets = kwargs.pop('targets', None)
110 |         for i, slices_ in enumerate(slices):
111 |             crop = inputs[(...,) + tuple(slices_)].to(device)
112 |             if inputs_mask is not None:
113 |                 crop_m = inputs_mask[(...,) + tuple(slices_)].to(device)
114 |                 if not torch.any(crop_m):
115 |                     continue  # skip masked out tile
116 |             outputs = self.forward(crop, targets=targets, max_imsize=False, **kwargs)
117 |             h_i, w_i = np.unravel_index(i, slices_by_dim)
118 |             h_start, w_start = [s.start for s in slices_]
119 | 
120 |             top, bottom = h_i > 0, h_i < (h_tiles - 1)
121 |             right, left = w_i < (w_tiles - 1), w_i > 0
122 |             for j in torch.arange(0, inputs.shape[0]):
123 |                 contours = outputs['contours'][j]
124 |                 boxes = outputs['boxes'][j]
125 |                 scores = outputs['scores'][j]
126 |                 extra = [outputs[k][j] for k in extra_keys]
127 | 
128 |                 # Remove small boxes (default min_size: 1.)
129 |                 keep = remove_small_boxes(boxes, box_min_size)
130 |                 contours, scores, boxes = (c[keep] for c in (contours, scores, boxes))
131 |                 extra = [e[keep] for e in extra]
132 | 
133 |                 # Remove partial detections to avoid tiling artifacts
134 |                 keep = remove_border_contours(contours, crop.shape[2:], border_removal,
135 |                                               top=top, right=right, bottom=bottom, left=left)
136 |                 contours, scores, boxes = (c[keep] for c in (contours, scores, boxes))
137 |                 extra = [e[keep] for e in extra]
138 | 
139 |                 # Add offset  # TODO: Replace with cpn internal offsets
140 |                 contours[..., 1] += h_start
141 |                 contours[..., 0] += w_start
142 |                 boxes[..., [0, 2]] += w_start
143 |                 boxes[..., [1, 3]] += h_start
144 | 
145 |                 results[j][i] = dict(
146 |                     contours=contours,
147 |                     boxes=boxes,
148 |                     scores=scores,
149 |                     extra=extra,  # list
150 |                     keep=torch.ones(contours.shape[0], dtype=torch.bool)
151 |                 )
152 | 
153 |         final = OrderedDict(
154 |             contours=[torch.cat([res_['contours'] for res_ in res if res_ is not None]) for res in results],
155 |             scores=[torch.cat([res_['scores'] for res_ in res if res_ is not None]) for res in results],
156 |             boxes=[torch.cat([res_['boxes'] for res_ in res if res_ is not None]) for res in results],
157 |             **{k: [torch.cat([res_['extra'][i] for res_ in res if res_ is not None]) for res in results] for i, k in
158 |                enumerate(extra_keys)}
159 |         )
160 | 
161 |         if not self.training:
162 |             nms_method = self.model.__dict__.get('nms_method', 'nms')
163 |             for n in torch.arange(inputs.shape[0]):
164 |                 boxes = final['boxes'][n]
165 |                 reference = boxes.shape[0]
166 |                 if nms_method == 'nms':
167 |                     keep = nms(boxes, final['scores'][n], iou_threshold=nms_thresh)
168 |                 else:
169 |                     raise ValueError(f'Unknown nms method: {self.nms_method}')
170 |                 for k in final:
171 |                     v = final[k][n]
172 |                     if extra_nms.get(k, True):
173 |                         assert v.shape[0] == reference, f'Output `{k}` is not compatible with nms. ' \
174 |                                                         f'Specify extra_nms=dict({k}=False) when calling ' \
175 |                                                         f'forward method.'
176 |                         final[k][n] = v[keep]
177 |         return final
178 | 
179 |     @rank_zero_only
180 |     def log_batch(self: 'pl.LightningModule', batch: dict, stage: str, keys=('inputs', 'labels'), global_step=None):
181 |         if global_step is None:
182 |             global_step = self.global_step
183 |         super().log_batch(batch=batch, stage=stage, keys=keys)
184 |         for logger in self._iter_loggers():
185 |             if hasattr(logger, 'add_image'):
186 |                 if 'sampled_contours' in batch:
187 |                     # Prepare contours (contours may be zero padded)
188 |                     cons = [c[:cl.max()] for c, cl in zip(batch['sampled_contours'], batch['labels'])]
189 |                     self.log_contour_figures(f'{stage}/contours', batch['inputs'], cons, logger=logger,
190 |                                              global_step=global_step)
191 |                 if 'resampled_contours' in batch:
192 |                     # Prepare contours (contours may be zero padded)
193 |                     cons = [c[:cl.max()] for c, cl in zip(batch['resampled_contours'], batch['labels'])]
194 |                     self.log_contour_figures(f'{stage}/contours_resampled', batch['inputs'], cons, logger=logger,
195 |                                              global_step=global_step)
196 | 


--------------------------------------------------------------------------------
/celldetection/models/loss.py:
--------------------------------------------------------------------------------
 1 | from torch.nn.modules.loss import _Loss
 2 | from torch import Tensor
 3 | from torchvision.ops.focal_loss import sigmoid_focal_loss
 4 | from ..ops.loss import iou_loss, box_npll_loss
 5 | 
 6 | __all__ = ['SigmoidFocalLoss', 'IoULoss', 'BoxNpllLoss']
 7 | 
 8 | 
 9 | class _FocalLoss(_Loss):
10 |     __constants__ = ['reduction']
11 | 
12 |     def __init__(self, size_average=None, reduce=None, reduction: str = 'mean', alpha=.5, gamma=2) -> None:
13 |         super().__init__(size_average, reduce, reduction)
14 |         self.alpha = alpha
15 |         self.gamma = gamma
16 | 
17 | 
18 | class SigmoidFocalLoss(_FocalLoss):
19 |     def forward(self, input: Tensor, target: Tensor) -> Tensor:
20 |         return sigmoid_focal_loss(input, target, alpha=self.alpha, gamma=self.gamma, reduction=self.reduction)
21 | 
22 | 
23 | class IoULoss(_Loss):
24 |     __constants__ = ['reduction']
25 | 
26 |     def __init__(self, generalized=True, method='linear', min_size=None, size_average=None, reduce=None,
27 |                  reduction: str = 'mean') -> None:
28 |         super().__init__(size_average, reduce, reduction)
29 |         self.generalized = generalized
30 |         self.method = method
31 |         self.min_size = min_size
32 | 
33 |     def forward(self, input: Tensor, target: Tensor) -> Tensor:
34 |         return iou_loss(input, target, self.reduction, generalized=self.generalized, method=self.method,
35 |                         min_size=self.min_size)
36 | 
37 |     def extra_repr(self) -> str:
38 |         return f"generalized={self.generalized}, method='{self.method}'"
39 | 
40 | 
41 | class BoxNpllLoss(_Loss):
42 |     __constants__ = ['reduction']
43 | 
44 |     def __init__(self, factor=10., sigmoid=False, min_size=None, epsilon=1e-8, size_average=None, reduce=None,
45 |                  reduction: str = 'mean') -> None:
46 |         super().__init__(size_average, reduce, reduction)
47 |         self.factor = factor
48 |         self.sigmoid = sigmoid
49 |         self.min_size = min_size
50 |         self.epsilon = epsilon
51 | 
52 |     def forward(self, uncertainty: Tensor, input: Tensor, target: Tensor) -> Tensor:
53 |         return box_npll_loss(uncertainty, input, target, factor=self.factor, sigmoid=self.sigmoid, epsilon=self.epsilon,
54 |                              reduction=self.reduction, min_size=self.min_size)
55 | 


--------------------------------------------------------------------------------
/celldetection/models/mamba.py:
--------------------------------------------------------------------------------
 1 | try:
 2 |     from mamba_ssm import Mamba
 3 | except ModuleNotFoundError:
 4 |     Mamba = False
 5 | 
 6 | 
 7 | from torch.cuda.amp import autocast
 8 | import torch
 9 | import torch.nn as nn
10 | 
11 | __all__ = ['MambaLayer']
12 | 
13 | 
14 | class MambaLayer(nn.Module):
15 |     def __init__(self, channels, d_state=16, kernel_size=4, expand=2, nd=None):
16 |         """Mamba Layer.
17 | 
18 |         References:
19 |             - https://arxiv.org/abs/2401.04722
20 |             - https://github.com/bowang-lab/U-Mamba
21 |             - https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py
22 | 
23 |         Examples:
24 |             ```python
25 |             model = cd.models.ResNet50(3, secondary_block=MambaLayer).cuda()
26 |             ```
27 | 
28 |         Note:
29 |             Requires model to run on Cuda!
30 | 
31 |         Args:
32 |             channels: Input channels. Equal to output channels.
33 |             d_state: SSM state expansion factor.
34 |             kernel_size: Kernel size of 1D convolution.
35 |             expand: Block expansion factor. (inner_channels=expand*channels)
36 |             nd: Spatial dimension. Just here for compatibility.
37 |         """
38 |         super().__init__()
39 |         self.norm = nn.LayerNorm(channels)
40 |         assert Mamba, ('The Python package `mamba_ssm` must be installed to use this package. '
41 |                        'Since the installation may not work on every system it is not a mandatory '
42 |                        'requirement for `celldetection`. '
43 |                        'Please follow the installation instructions here: https://github.com/state-spaces/mamba.')
44 |         self.mamba = Mamba(d_model=channels, d_state=d_state, d_conv=kernel_size, expand=expand)
45 | 
46 |     @autocast(enabled=False)
47 |     def forward(self, x):
48 |         if x.dtype == torch.float16:
49 |             x = x.type(torch.float32)
50 |         n, c, *spatial = x.shape
51 |         x_ = x.flatten(2).transpose(1, 2)  # Tensor[n, h*w(*d), c]
52 |         x_ = self.norm(x_)
53 |         x_ = self.mamba(x_)
54 |         x = x_.transpose(1, 2).reshape(n, c, *spatial)  # Tensor[n, c, h, w(, d)]
55 |         return x
56 | 


--------------------------------------------------------------------------------
/celldetection/models/manet.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | from torch.nn.functional import softmax
  4 | from typing import Optional, Union, List, Type, Dict
  5 | from torch import Tensor
  6 | from torch.nn import functional as F
  7 | from functools import partial
  8 | from ..util.util import lookup_nn, dict2model
  9 | from .unet import UNet, IntermediateUNetBlock
 10 | from .commons import ConvNormRelu, SqueezeExcitation as SE, _ni_3d
 11 | from .smp import SmpEncoder
 12 | from .timmodels import TimmEncoder
 13 | 
 14 | __all__ = [
 15 |     'MaNet', 'MultiscaleFusionAttention', 'PositionWiseAttention',
 16 |     'SmpMaNet', 'TimmMaNet'
 17 | ]
 18 | 
 19 | 
 20 | class SqueezeExcitation(SE):
 21 |     def forward(self, inputs):
 22 |         scale = super(SE, self).forward(inputs)
 23 |         return scale
 24 | 
 25 | 
 26 | class PositionWiseAttention(nn.Module):
 27 |     def __init__(self, in_channels, out_channels, mid_channels=64, kernel_size=3, padding=1, beta=False, nd=2):
 28 |         super().__init__()
 29 |         Conv = lookup_nn('Conv2d', nd=nd, call=False)
 30 |         self.beta = nn.Parameter(torch.zeros(1)) if beta else 1.
 31 |         if in_channels != out_channels:
 32 |             self.in_conv = Conv(in_channels, out_channels, kernel_size=3, padding=1)
 33 |         else:
 34 |             self.in_conv = None
 35 |         self.proj_b, self.proj_a = [Conv(out_channels, mid_channels, kernel_size=1) for _ in range(2)]
 36 |         self.proj = Conv(out_channels, out_channels, kernel_size=kernel_size, padding=padding)
 37 |         self.out_conv = Conv(out_channels, out_channels, kernel_size=3, padding=1)
 38 | 
 39 |     def forward(self, inputs):
 40 |         x = inputs if self.in_conv is None else self.in_conv(inputs)
 41 |         a = self.proj_a(x).flatten(2)
 42 |         b = self.proj_b(x).flatten(2)  # Tensor[n, c', hw]
 43 |         p = torch.matmul(a.permute(0, 2, 1), b)  # Tensor[n, hw, hw]
 44 |         p = softmax(p.flatten(1), dim=1).view(p.shape)  # Tensor[n, hw, hw]
 45 |         c = self.proj(x).flatten(2)  # Tensor[n, c, hw]
 46 |         out = torch.matmul(p, c.permute(0, 2, 1)).view(*c.shape[:2], *inputs.shape[2:])  # T[n, c, hw] -> T[n, c, h, w]
 47 |         out = self.out_conv(self.beta * out + x)
 48 |         return out
 49 | 
 50 | 
 51 | class PAB(IntermediateUNetBlock):
 52 |     def __init__(self, in_channels, out_channels, mid_channels=64, kernel_size=3, padding=1, nd=2, replace=False):
 53 |         kwargs = dict(out_channels=(out_channels,), out_strides=(1,))
 54 |         if replace:
 55 |             kwargs = dict(out_channels=(), out_strides=())
 56 |         super().__init__(**kwargs)
 57 |         self.module = PositionWiseAttention(in_channels=in_channels, out_channels=out_channels,
 58 |                                             mid_channels=mid_channels, kernel_size=kernel_size, padding=padding, nd=nd)
 59 |         self.replace = replace
 60 | 
 61 |     def forward(
 62 |             self,
 63 |             x: Dict[str, Tensor],
 64 |     ) -> Dict[str, Tensor]:
 65 |         in_key = list(x.keys())[-1]
 66 |         out_key = in_key if self.replace else str(len(x))
 67 |         x[out_key] = self.module(x[in_key])
 68 |         return x
 69 | 
 70 | 
 71 | class MultiscaleFusionAttention(nn.Module):
 72 |     def __init__(
 73 |             self,
 74 |             in_channels,  # main
 75 |             in_channels2,  # lateral
 76 |             out_channels,
 77 |             norm_layer='BatchNorm2d',
 78 |             activation='relu',
 79 |             compression=16,
 80 |             interpolation=None,
 81 |             nd=2,
 82 |     ):
 83 |         super().__init__()
 84 |         kw = dict(activation=activation, norm_layer=norm_layer, nd=nd, bias=False)
 85 |         self.in_block = nn.Sequential(
 86 |             ConvNormRelu(in_channels, in_channels, **kw),
 87 |             ConvNormRelu(in_channels, in_channels2, kernel_size=1, padding=0, **kw),
 88 |         )
 89 |         self.se_high = SqueezeExcitation(in_channels2, compression=compression, activation=activation, nd=nd)
 90 |         self.se_low = SqueezeExcitation(in_channels2, compression=compression, activation=activation, nd=nd)
 91 |         self.out_block = nn.Sequential(
 92 |             ConvNormRelu(in_channels2 + in_channels2, out_channels, **kw),
 93 |             ConvNormRelu(out_channels, out_channels, **kw),
 94 |         )
 95 |         if interpolation is True:
 96 |             interpolation = 'nearest'
 97 |         elif interpolation is False:
 98 |             interpolation = None
 99 |         self.interpolation = interpolation
100 | 
101 |     def forward(self, x, x2=None):
102 |         if isinstance(x, (tuple, list)):
103 |             assert x2 is None
104 |             x, x2 = x
105 |         x = self.in_block(x)
106 |         if self.interpolation is not None:
107 |             x = F.interpolate(x, x2.shape[2:], mode=self.interpolation)
108 |         if x2 is not None:
109 |             a = self.se_high(x)  # main
110 |             b = self.se_low(x2)  # lateral
111 |             x = x * (a + b)  # scaled x
112 |             x = torch.cat((x, x2), 1)
113 |         return self.out_block(x)
114 | 
115 | 
116 | class MaNet(UNet):
117 |     def __init__(self, backbone, out_channels: int = 0, pab_channels=64,
118 |                  block: Type[nn.Module] = None, block_kwargs: dict = None, final_activation=None,
119 |                  interpolate='nearest', nd=2, **kwargs):
120 |         """Multi-Scale Attention Network.
121 | 
122 |         A U-Net variant using a generic encoder and a special decoder that includes a
123 |         Position-wise Attention Block (PAB) and several Multi-scale Fusion Attention Blocks (MFAB).
124 | 
125 |         References:
126 |             - https://ieeexplore.ieee.org/document/9201310
127 | 
128 |         Args:
129 |             backbone: Backbone instance.
130 |             out_channels: Output channels. If set to ``0``, the output layer is omitted.
131 |             pab_channels: Channels of the Position-wise Attention Block (PAB).
132 |             block: Main block. Default: Multi-scale Fusion Attention Block (MFAB).
133 |             block_kwargs: Block keyword arguments.
134 |             final_activation: Final activation function.
135 |             interpolate: Interpolation.
136 |             nd: Spatial dimensions.
137 |             **kwargs: Additional keyword arguments.
138 |         """
139 |         if isinstance(backbone, dict):
140 |             backbone = dict2model(backbone)
141 | 
142 |         oc = backbone.out_channels
143 |         intermediate_blocks = None
144 |         if pab_channels:
145 |             intermediate_blocks = PAB(oc[-1], oc[-1], mid_channels=pab_channels, nd=nd,
146 |                                       replace=True, **kwargs.get('pwa_kwargs', {}))
147 |         kwargs['block_interpolate'] = block_interpolate = kwargs.get('block_interpolate', True)
148 |         if block is None:
149 |             block = partial(MultiscaleFusionAttention, interpolation=block_interpolate if block_interpolate else None)
150 |             kwargs['block_cat'] = kwargs.get('block_cat', True)
151 |         super().__init__(backbone=backbone, out_channels=out_channels, block=block,
152 |                          block_kwargs=block_kwargs, final_activation=final_activation, interpolate=interpolate,
153 |                          nd=nd, intermediate_blocks=intermediate_blocks, **kwargs)
154 | 
155 | 
156 | def _default_timm_kwargs(backbone_kwargs, pretrained=False):
157 |     kw = dict(pretrained=pretrained)
158 |     kw.update({} if backbone_kwargs is None else backbone_kwargs)
159 |     return kw
160 | 
161 | 
162 | class TimmMaNet(MaNet):
163 |     def __init__(self, in_channels, out_channels, model_name, final_activation=None, backbone_kwargs=None,
164 |                  pretrained=True, block_cls=None, nd=2, **kwargs):
165 |         _ni_3d(nd)
166 |         super().__init__(TimmEncoder(model_name=model_name, in_channels=in_channels,
167 |                                      **_default_timm_kwargs(backbone_kwargs, pretrained)),
168 |                          out_channels, final_activation=final_activation, block=block_cls, **kwargs)
169 |         self.save_hyperparameters()
170 | 
171 | 
172 | def _default_smp_kwargs(backbone_kwargs, pretrained=False):
173 |     if pretrained is True:
174 |         pretrained = 'imagenet'
175 |     elif pretrained is False:
176 |         pretrained = None
177 |     kw = dict(weights=pretrained)
178 |     kw.update({} if backbone_kwargs is None else backbone_kwargs)
179 |     return kw
180 | 
181 | 
182 | class SmpMaNet(MaNet):
183 |     def __init__(self, in_channels, out_channels, model_name, final_activation=None, backbone_kwargs=None,
184 |                  pretrained=True, block_cls=None, nd=2, **kwargs):
185 |         _ni_3d(nd)
186 |         super().__init__(SmpEncoder(model_name=model_name, in_channels=in_channels,
187 |                                     **_default_smp_kwargs(backbone_kwargs, pretrained)),
188 |                          out_channels, final_activation=final_activation, block=block_cls, **kwargs)
189 |         self.save_hyperparameters()
190 | 


--------------------------------------------------------------------------------
/celldetection/models/mobilenetv3.py:
--------------------------------------------------------------------------------
 1 | import torch.nn as nn
 2 | from torchvision.models.mobilenetv3 import InvertedResidualConfig, InvertedResidual, _mobilenet_v3_conf
 3 | from typing import Any, Callable, List, Optional, Sequence
 4 | from functools import partial
 5 | 
 6 | try:
 7 |     from torchvision.ops.misc import Conv2dNormActivation as ConvBNActivation
 8 | except ImportError:
 9 |     from torchvision.models.mobilenetv2 import ConvBNActivation  # for older versions
10 | 
11 | __all__ = ['MobileNetV3Large', 'MobileNetV3Small']
12 | 
13 | 
14 | def init_modules_(mod: nn.Module):
15 |     for m in mod.modules():
16 |         if isinstance(m, nn.Conv2d):
17 |             nn.init.kaiming_normal_(m.weight, mode='fan_out')
18 |             if m.bias is not None:
19 |                 nn.init.zeros_(m.bias)
20 |         elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
21 |             nn.init.ones_(m.weight)
22 |             nn.init.zeros_(m.bias)
23 |         elif isinstance(m, nn.Linear):
24 |             nn.init.normal_(m.weight, 0, 0.01)
25 |             nn.init.zeros_(m.bias)
26 | 
27 | 
28 | class MobileNetV3Base(nn.Sequential):
29 |     """Adaptation of torchvision.models.mobilenetv3.MobileNetV3"""
30 | 
31 |     def __init__(
32 |             self,
33 |             in_channels,
34 |             inverted_residual_setting: List[InvertedResidualConfig],
35 |             block: Optional[Callable[..., nn.Module]] = None,
36 |             norm_layer: Optional[Callable[..., nn.Module]] = None,
37 |             **kwargs: Any
38 |     ) -> None:
39 |         super().__init__()
40 | 
41 |         if not inverted_residual_setting:
42 |             raise ValueError("The inverted_residual_setting should not be empty")
43 |         elif not (isinstance(inverted_residual_setting, Sequence) and
44 |                   all([isinstance(s, InvertedResidualConfig) for s in inverted_residual_setting])):
45 |             raise TypeError("The inverted_residual_setting should be List[InvertedResidualConfig]")
46 | 
47 |         if block is None:
48 |             block = InvertedResidual
49 |         if norm_layer is None:
50 |             norm_layer = partial(nn.BatchNorm2d, eps=0.001, momentum=0.01)
51 | 
52 |         layers: List[nn.Sequential] = [nn.Sequential()]
53 |         cbna_kw = {}  # needed because torchvision changed behaviour
54 |         if norm_layer is not None:
55 |             cbna_kw['norm_layer'] = norm_layer
56 | 
57 |         # building first layer
58 |         firstconv_output_channels = inverted_residual_setting[0].input_channels
59 |         self.out_channels = [firstconv_output_channels]
60 |         layers[-1].add_module(str(len(layers[-1])),
61 |                               ConvBNActivation(in_channels, firstconv_output_channels, kernel_size=3, stride=2,
62 |                                                activation_layer=nn.Hardswish, **cbna_kw))
63 | 
64 |         # building inverted residual blocks
65 |         for cnf in inverted_residual_setting:
66 |             if cnf.stride > 1:
67 |                 layers.append(nn.Sequential())
68 |                 self.out_channels.append(cnf.out_channels)
69 |             else:
70 |                 self.out_channels[-1] = cnf.out_channels
71 |             layers[-1].add_module(str(len(layers[-1])), block(cnf, norm_layer))
72 | 
73 |         # building last several layers
74 |         lastconv_input_channels = inverted_residual_setting[-1].out_channels
75 |         lastconv_output_channels = 6 * lastconv_input_channels
76 |         self.out_channels[-1] = lastconv_output_channels
77 |         assert len(self.out_channels) == len(layers)
78 |         layers[-1].add_module(str(len(layers[-1])),
79 |                               ConvBNActivation(lastconv_input_channels, lastconv_output_channels, kernel_size=1,
80 |                                                activation_layer=nn.Hardswish, **cbna_kw))
81 | 
82 |         super().__init__(*layers)
83 |         init_modules_(self)
84 | 
85 | 
86 | class MobileNetV3Large(MobileNetV3Base):
87 |     def __init__(self, in_channels, width_mult: float = 1.0, reduced_tail: bool = False, dilated: bool = False):
88 |         super().__init__(in_channels=in_channels, inverted_residual_setting=_mobilenet_v3_conf(
89 |             'mobilenet_v3_large', width_mult=width_mult, reduced_tail=reduced_tail, dilated=dilated)[0])
90 | 
91 | 
92 | class MobileNetV3Small(MobileNetV3Base):
93 |     def __init__(self, in_channels, width_mult: float = 1.0, reduced_tail: bool = False, dilated: bool = False):
94 |         super().__init__(in_channels=in_channels, inverted_residual_setting=_mobilenet_v3_conf(
95 |             'mobilenet_v3_small', width_mult=width_mult, reduced_tail=reduced_tail, dilated=dilated)[0])
96 | 


--------------------------------------------------------------------------------
/celldetection/models/normalization.py:
--------------------------------------------------------------------------------
 1 | import torch.nn as nn
 2 | from ..ops.normalization import pixel_norm
 3 | 
 4 | __all__ = ['PixelNorm']
 5 | 
 6 | 
 7 | class PixelNorm(nn.Module):
 8 |     def __init__(self, dim=1, eps=1e-8):
 9 |         """Pixel normalization.
10 | 
11 |         References:
12 |             - https://papers.nips.cc/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf ("Local Repose Norm")
13 |             - https://arxiv.org/pdf/1710.10196.pdf
14 | 
15 |         Args:
16 |             dim: Dimension to normalize.
17 |             eps: Epsilon.
18 |         """
19 |         super().__init__()
20 |         self.dim = dim
21 |         self.eps = eps
22 | 
23 |     def forward(self, x):
24 |         return pixel_norm(x, dim=self.dim, eps=self.eps)
25 | 


--------------------------------------------------------------------------------
/celldetection/models/ppm.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import torch.nn as nn
 3 | import torch.nn.functional as F
 4 | from typing import Union
 5 | from torchvision.models.segmentation.deeplabv3 import ASPP
 6 | from ..util.util import lookup_nn, get_nd_conv
 7 | 
 8 | __all__ = ['Ppm']
 9 | 
10 | 
11 | class Ppm(nn.Module):
12 |     def __init__(
13 |             self,
14 |             in_channels, out_channels,
15 |             scales: Union[list, tuple] = (1, 2, 3, 6),
16 |             kernel_size=1,
17 |             norm='BatchNorm2d',
18 |             activation='relu',
19 |             concatenate=True,
20 |             nd=2,
21 |             **kwargs
22 |     ):
23 |         """Pyramid Pooling Module.
24 | 
25 |         References:
26 |             - https://ieeexplore.ieee.org/document/8100143
27 | 
28 |         Args:
29 |             in_channels: Number of input channels.
30 |             out_channels: Number of output channels per pyramid scale.
31 |             scales: Pyramid scales. Default: (1, 2, 3, 6).
32 |             kernel_size: Kernel size.
33 |             norm: Normalization.
34 |             activation: Activation.
35 |             concatenate: Whether to concatenate module inputs to pyramid pooling output before returning results.
36 |             **kwargs: Keyword arguments for ``nn.Conv2d``.
37 |         """
38 |         super().__init__()
39 |         self.blocks = nn.ModuleList()
40 |         self.concatenate = concatenate
41 |         self.out_channels = out_channels * len(scales) + in_channels * int(concatenate)
42 |         Conv = get_nd_conv(nd)
43 |         AdaptiveAvgPool = lookup_nn(nn.AdaptiveAvgPool2d, call=False, nd=nd)
44 |         norm = lookup_nn(norm, call=False, nd=nd)
45 |         activation = lookup_nn(activation, call=False, nd=nd)
46 |         for scale in scales:
47 |             self.blocks.append(nn.Sequential(
48 |                 AdaptiveAvgPool(output_size=scale),
49 |                 Conv(in_channels, out_channels, kernel_size, **kwargs),
50 |                 norm(out_channels),
51 |                 activation(),
52 |             ))
53 | 
54 |     def forward(self, x):
55 |         prefix = [x] if self.concatenate else []
56 |         return torch.cat(prefix + [
57 |             F.interpolate(m(x), x.shape[2:], mode='bilinear', align_corners=False) for m in self.blocks
58 |         ], 1)
59 | 
60 | 
61 | def append_pyramid_pooling_(module: nn.Sequential, out_channels, scales=(1, 2, 3, 6), method='ppm', in_channels=None,
62 |                             **kwargs):
63 |     if in_channels is None:
64 |         in_channels = module.out_channels[-1]
65 |     method = method.lower()
66 |     if method == 'ppm':
67 |         assert (out_channels % len(scales)) == 0
68 |         p = Ppm(in_channels, out_channels, scales=scales, **kwargs)
69 |         out_channels = p.out_channels
70 |     elif method == 'aspp':
71 |         scales = sorted(tuple(set(scales) - {1}))
72 |         nd = kwargs.pop('nd', 2)
73 |         assert nd == 2, NotImplementedError('Only nd=2 supported.')
74 |         p = ASPP(in_channels, scales, out_channels, **kwargs)
75 |     else:
76 |         raise ValueError
77 |     module.append(p)
78 |     if hasattr(module, 'out_channels'):
79 |         module.out_channels += (out_channels,)
80 |     if hasattr(module, 'out_strides'):
81 |         module.out_strides += module.out_strides[-1:]
82 | 


--------------------------------------------------------------------------------
/celldetection/models/smp.py:
--------------------------------------------------------------------------------
 1 | import segmentation_models_pytorch as smp
 2 | from torch import nn
 3 | from collections import OrderedDict
 4 | from typing import List, Callable
 5 | from pytorch_lightning.core.mixins import HyperparametersMixin
 6 | 
 7 | __all__ = ['SmpEncoder']
 8 | 
 9 | 
10 | class ExternBase(nn.Module):
11 |     def __init__(self, model_name: str):
12 |         super().__init__()
13 |         self.check_model_name(model_name)
14 |         assert not model_name.startswith('_')
15 |         self.model_name = model_name
16 | 
17 |     @staticmethod
18 |     def check_model_name(model_name: str, model_list_callback: Callable[[], List[str]] = None):
19 |         if model_name is None:
20 |             msg = 'Please specify `model_name`.'
21 |             if model_list_callback is not None:
22 |                 msg += '\n  - '.join([' These are all available models, some of them may not be supported:'
23 |                                       ] + model_list_callback())
24 |             raise ValueError(msg)
25 | 
26 |     def _get_name(self):
27 |         return self.model_name.title()
28 | 
29 | 
30 | class SmpEncoder(ExternBase, HyperparametersMixin):
31 |     def __init__(self, model_name: str, in_channels: int = 3, depth: int = 5, pretrained=False,
32 |                  output_stride: int = 32, **kwargs):
33 |         """Smp Encoder.
34 | 
35 |         A wrapper that provides compatibility with "segmentation_models_pytorch (smp)".
36 | 
37 |         References:
38 |             - smp GitHub: https://github.com/qubvel/segmentation_models.pytorch
39 |             - Encoders: https://smp.readthedocs.io/en/latest/encoders.html
40 |             - Documentation: https://smp.readthedocs.io/en/latest/
41 | 
42 |         Args:
43 |             model_name: Encoder name. Find available encoders here: https://smp.readthedocs.io/en/latest/encoders.html
44 |             in_channels: Input channels.
45 |             depth: Encoder dpeth.
46 |             pretrained: Encoder weights. Find available weights here: https://smp.readthedocs.io/en/latest/encoders.html
47 |             output_stride: Output stride.
48 |             **kwargs:
49 |         """
50 |         self.check_model_name(model_name, smp.encoders.get_encoder_names)
51 |         super().__init__(model_name)
52 |         self.save_hyperparameters()
53 | 
54 |         # Map pretrained for consistency
55 |         if pretrained is True:
56 |             pretrained = 'imagenet'  # best guess
57 |         elif pretrained is False:
58 |             pretrained = None
59 |         if 'weights' in kwargs:
60 |             pretrained = kwargs.pop('weights')
61 |         self.module = smp.encoders.get_encoder(self.model_name, in_channels=in_channels, depth=depth,
62 |                                                weights=pretrained, output_stride=output_stride, **kwargs)
63 |         if pretrained:
64 |             self.pretrained_cfg = smp.encoders.get_preprocessing_params(self.model_name, pretrained)
65 |         self._skips = 0
66 |         if self.module.out_channels[0] == in_channels:
67 |             self.out_channels = self.module.out_channels[1:]  # skip input image
68 |             self._skips += 1
69 |         while self.out_channels[0] <= 0:
70 |             self.out_channels = self.out_channels[1:]  # skip dummies
71 |             self._skips += 1
72 |         self.out_strides = tuple(
73 |             [2 ** i for i in range(self._skips, len(self.module.out_channels))])  # derive strides from rank
74 | 
75 |     def forward(self, x, *args, **kwargs):
76 |         out = self.module(x, *args, **kwargs)
77 |         assert isinstance(out, list), type(out)
78 |         out = out[self._skips:]
79 |         out = OrderedDict([(str(k), v) for k, v in zip(range(len(out)), out)])
80 |         return out
81 | 


--------------------------------------------------------------------------------
/celldetection/models/timmodels.py:
--------------------------------------------------------------------------------
  1 | import timm
  2 | import torch
  3 | from torch import nn
  4 | from collections import OrderedDict
  5 | from .smp import ExternBase
  6 | from ..util.util import get_device
  7 | import re
  8 | from typing import List, Tuple
  9 | from warnings import warn
 10 | 
 11 | __all__ = ['TimmEncoder']
 12 | 
 13 | 
 14 | def get_channels_strides(model, example_input=None, verbose=True) -> Tuple[List[int], List[int]]:
 15 |     """Get channels and strides.
 16 | 
 17 |     A helper function to probe the channels and strides of the output of ``model``.
 18 | 
 19 |     Examples:
 20 |         >>> m = cd.models.TimmEncoder(
 21 |         ...     'cspdarknet53',
 22 |         ...     return_layers=['stem', 'stages.0', 'stages.1', 'stages.2', 'stages.3', 'stages.4'],
 23 |         ... )
 24 |         >>> out_channels, out_strides = get_channels_strides(m, torch.rand(1, 3, 512, 512), verbose=True)
 25 |         layer 0 output shape: torch.Size([1, 32, 512, 512]) , stride: 1
 26 |         layer 1 output shape: torch.Size([1, 64, 256, 256]) , stride: 2
 27 |         layer 2 output shape: torch.Size([1, 128, 128, 128]) , stride: 4
 28 |         layer 3 output shape: torch.Size([1, 256, 64, 64]) , stride: 8
 29 |         layer 4 output shape: torch.Size([1, 512, 32, 32]) , stride: 16
 30 |         layer 5 output shape: torch.Size([1, 1024, 16, 16]) , stride: 32
 31 |         >>> out_strides
 32 |         [1, 2, 4, 8, 16, 32]
 33 |         >>> out_channels
 34 |         [32, 64, 128, 256, 512, 1024]
 35 | 
 36 |     Args:
 37 |         model: Model.
 38 |         example_input: Example input.
 39 |         verbose: Whether to print output names, shapes and strides.
 40 | 
 41 |     Returns:
 42 |         Strides as List[int].
 43 |     """
 44 |     if example_input is None:
 45 |         reference = 512
 46 |         example_input = torch.rand(1, 3, *(reference,) * 2).to(get_device(model))
 47 |     else:
 48 |         reference = example_input.shape[-1]
 49 |     training = model.training
 50 |     model.eval()
 51 |     with torch.no_grad():
 52 |         out = model(example_input)
 53 |     model.train(training)
 54 | 
 55 |     strides = []
 56 |     channels = []
 57 |     if isinstance(out, dict):
 58 |         for k, v in out.items():
 59 |             strides.append(reference // v.shape[-1])
 60 |             channels.append(v.shape[1])
 61 |             if verbose:
 62 |                 print('layer', k, 'output shape:', v.shape, ', stride:', strides[-1])
 63 |     elif isinstance(out, torch.Tensor):
 64 |         assert out.ndim == example_input.ndim, f'Module output has different number of dimensions: ' \
 65 |                                                f'{out.shape, example_input.shape}. ' \
 66 |                                                f'Consider updating ``return_layers`` in __init__ call.'
 67 |         strides.append(reference // out.shape[-1])
 68 |         channels.append(out.shape[1])
 69 |     else:
 70 |         raise ValueError('Cannot handle output.')
 71 |     return channels, strides
 72 | 
 73 | 
 74 | def get_names(model, *pattern) -> List[str]:
 75 |     """Get names.
 76 | 
 77 |     A helper function to discover module names for intermediate feature extraction.
 78 |     A list of module names can be passed to the constructor of this class (``return_layers``)
 79 |     to convert the timm model to a custom feature extractor.
 80 | 
 81 |     Examples:
 82 |         >>> m = cd.models.TimmEncoder('convnext_nano')
 83 |         >>> get_names(m, 'stem$', 'stages.[0-9]$')
 84 |         ['stem', 'stages.0', 'stages.1', 'stages.2', 'stages.3']
 85 | 
 86 |     Args:
 87 |         model: Model.
 88 |         *pattern: Name patterns.
 89 |     """
 90 |     return [n for n, _ in model.module.named_modules() if
 91 |             len(pattern) == 0 or any(list(bool(len(re.findall(p, n))) for p in pattern))]
 92 | 
 93 | 
 94 | class TimmEncoder(ExternBase):
 95 |     def __init__(
 96 |             self,
 97 |             model_name: str,
 98 |             in_channels: int = 3,
 99 |             return_layers: List[str] = None,
100 |             out_channels: List[str] = None,
101 |             out_strides: List[str] = None,
102 |             pretrained: bool = False,
103 |             pretrained_cfg=None,
104 |             keep_names: bool = False,
105 |             output_stride: int = None,
106 |             depth: int = None,
107 |             convert_lists: bool = True,
108 |             **kwargs
109 |     ):
110 |         """Timm Encoder.
111 | 
112 |         A wrapper that provides compatibility with "PyTorch Image Models (timm)".
113 | 
114 |         Notes:
115 |             It is possible to customize TimmEncoder`s return_layers to return specific layer outputs.
116 |             In that case ``out_channels`` and ``out_strides`` must be provided.
117 |             To find appropriate settings the helper functions ``get_names`` and ``get_channels_strides`` can be used.
118 | 
119 |         References:
120 |             - timm GitHub: https://github.com/huggingface/pytorch-image-models
121 |             - timm Documentation: https://huggingface.co/docs/timm/index
122 | 
123 |         Examples:
124 |             >>> import timm, celldetection as cd
125 |             >>> timm.list_models('*darknet*')  # discover models
126 |             ['cs3darknet_focus_l',
127 |              'cs3darknet_focus_m',
128 |              'cs3darknet_focus_s',
129 |              'cs3darknet_focus_x',
130 |              'cs3darknet_l',
131 |              'cs3darknet_m',
132 |              'cs3darknet_s',
133 |              'cs3darknet_x',
134 |              'cs3sedarknet_l',
135 |              'cs3sedarknet_x',
136 |              'cs3sedarknet_xdw',
137 |              'cspdarknet53',
138 |              'darknet17',
139 |              'darknet21',
140 |              'darknet53',
141 |              'darknetaa53',
142 |              'sedarknet21']
143 |              >>> encoder = cd.models.TimmEncoder('darknet21')
144 |              >>> encoder.out_channels
145 |              ... [32, 64, 128, 256, 512, 1024]
146 |              >>> encoder.out_strides
147 |              ... [1, 2, 4, 8, 16, 32]
148 |              >>> output: Dict[str, Tensor] = encoder(torch.rand(1, 3, 512, 512))
149 |              >>> for key, layer_output in output.items():
150 |              ...    print(key, layer_output.shape)
151 |              0 torch.Size([1, 32, 512, 512])
152 |              1 torch.Size([1, 64, 256, 256])
153 |              2 torch.Size([1, 128, 128, 128])
154 |              3 torch.Size([1, 256, 64, 64])
155 |              4 torch.Size([1, 512, 32, 32])
156 |              5 torch.Size([1, 1024, 16, 16])
157 | 
158 |         Args:
159 |             model_name: Name of model to instantiate.
160 |             return_layers: List of layer names used for intermediate feature retrieval.
161 |             out_channels: List of output channels per return layer.
162 |             out_strides: List of output strides per return layer.
163 |             pretrained: Whether to load pretrained ImageNet-1k weights.
164 |             pretrained_cfg: External pretrained_cfg for model.
165 |             keep_names: Whether to keep layer names for model output. If ``False``, names are replaced
166 |                 with enumeration for consistency.
167 |             output_stride: Some models support different output stride (e.g. 16 instead of 32).
168 |                 This is achieved by using ``stride=1`` with dilation instead of downsampling with strides.
169 |             depth: Custom encoder depth. If return_layers provided, this acts as number of return layers.
170 |             convert_lists: Whether to return output lists to dictionaries for consistency.
171 |             **kwargs: Keyword arguments for ``timm.create_model`` call.
172 |         """
173 |         super().__init__(model_name=model_name)
174 |         assert depth is None or depth > 0
175 |         depth_ = slice(None) if depth is None else slice(None, depth)
176 |         if output_stride is not None:
177 |             kwargs['output_stride'] = output_stride
178 | 
179 |         # Without custom return layers, use default layers from timm
180 |         if return_layers is None:
181 |             if depth is not None:
182 |                 kwargs['out_indices'] = tuple(range(depth))
183 |             try:
184 |                 self.module: nn.Module = timm.create_model(
185 |                     model_name=self.model_name,
186 |                     in_chans=in_channels,
187 |                     pretrained=pretrained,
188 |                     pretrained_cfg=pretrained_cfg,
189 |                     features_only=kwargs.get('features_only', True),
190 |                     **kwargs
191 |                 )
192 |             except RuntimeError as e:
193 |                 raise ValueError('This model does not support automatic feature extraction. '
194 |                                  'To retrieve features nonetheless, use `return_layers` to specify a list of names of '
195 |                                  'layers whose output is to be returned.\n\n' + str(e))
196 |             self.pretrained_cfg = self.module.__dict__.get('pretrained_cfg', {})
197 |             self.return_layers = [i['module'] for i in self.module.feature_info[depth_]]
198 |             if out_channels is None:
199 |                 out_channels = [i['num_chs'] for i in self.module.feature_info[depth_]]
200 |             if out_strides is None:
201 |                 out_strides = [i['reduction'] for i in self.module.feature_info[depth_]]
202 | 
203 |         # Custom return layers
204 |         else:
205 |             from torchvision.models.feature_extraction import create_feature_extractor
206 |             self.module: nn.Module = timm.create_model(
207 |                 model_name=self.model_name, pretrained=pretrained,
208 |                 in_chans=in_channels, pretrained_cfg=pretrained_cfg, **kwargs
209 |             )
210 |             self.pretrained_cfg = self.module.__dict__.get('pretrained_cfg', {})
211 |             if out_channels is None:
212 |                 warn('For custom return_layers setting out_channels must be specified.')
213 |             if out_strides is None:
214 |                 warn('For custom return_layers setting out_strides must be specified.')
215 |             self.return_layers = return_layers[depth_]
216 |             self.module = create_feature_extractor(self.module, self.return_layers)
217 |         self.out_channels = out_channels[depth_] if out_channels is not None else None
218 |         self.out_strides = out_strides[depth_] if out_strides is not None else None
219 |         self.keep_names = keep_names
220 |         self.convert_lists = convert_lists
221 | 
222 |     def forward(self, x):
223 |         out = self.module(x)  # assuming ordered
224 |         if self.convert_lists and isinstance(out, list):
225 |             out = OrderedDict([(str(k), v) for k, v in enumerate(out)])
226 |         if isinstance(out, dict):
227 |             if not self.keep_names:
228 |                 out = OrderedDict([(str(k), v) for k, (_, v) in enumerate(out.items())])
229 |         return out
230 | 


--------------------------------------------------------------------------------
/celldetection/mpi/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | This submodule contains MPI operations. For more information about MPI visit: https://mpi4py.readthedocs.io/en/stable/
3 | """
4 | from .mpi import *
5 | 


--------------------------------------------------------------------------------
/celldetection/ops/__init__.py:
--------------------------------------------------------------------------------
 1 | """
 2 | This submodule contains PyTorch operations.
 3 | """
 4 | from .commons import *
 5 | from .loss import *
 6 | from .draw import *
 7 | from .boxes import *
 8 | from .normalization import *
 9 | from .features import *
10 | from .cpn import *
11 | 


--------------------------------------------------------------------------------
/celldetection/ops/boxes.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from torch import Tensor
  3 | import torchvision.ops.boxes as bx
  4 | from typing import Tuple, List
  5 | import numpy as np
  6 | from os import name
  7 | 
  8 | WINDOWS = name == 'nt'
  9 | 
 10 | __all__ = ['nms', 'contours2boxes', 'pairwise_box_iou', 'pairwise_generalized_box_iou', 'filter_by_box_voting']
 11 | 
 12 | 
 13 | def torch_compile(*args, **kwargs):
 14 |     def decorator(func):
 15 |         if WINDOWS:
 16 |             return func  # compile not supported on Windows, yet
 17 |         else:
 18 |             return torch.compile(func, *args, **kwargs)
 19 | 
 20 |     return decorator
 21 | 
 22 | 
 23 | @torch_compile(dynamic=True)
 24 | def nms(boxes, scores, thresh=.5) -> torch.Tensor:
 25 |     """Non-maximum suppression.
 26 | 
 27 |     Perform non-maximum suppression (NMS) on the boxes according to their intersection-over-union (IoU).
 28 | 
 29 |     Notes:
 30 |         - Use ``torchvision.ops.boxes.nms`` if possible; This is just a "pure-python" alternative
 31 |         - ``cd.ops.boxes.nms`` for 8270 boxes: 13.9 ms ± 131 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
 32 |         - ``tv.ops.boxes.nms`` for 8270 boxes: 1.84 ms ± 4.91 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
 33 |         - ``cd.ops.boxes.nms`` for 179 boxes: 265 µs ± 1.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
 34 |         - ``tv.ops.boxes.nms`` for 179 boxes: 103 µs ± 2.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
 35 | 
 36 |     Args:
 37 |         boxes: Boxes. Tensor[num_boxes, 4] in (x0, y0, x1, y1) format.
 38 |         scores: Scores. Tensor[num_boxes].
 39 |         thresh: Threshold. Discards all overlapping boxes with ``IoU > thresh``.
 40 | 
 41 |     Returns:
 42 |         Keep indices. Tensor[num_keep].
 43 |     """
 44 |     indices = torch.argsort(scores, descending=True)
 45 |     boxes = boxes[indices]
 46 |     mask = bx.box_iou(boxes, boxes) > thresh
 47 |     vals, idx = torch.max(mask, 1)
 48 |     torch.logical_not(vals, out=vals)
 49 |     torch.logical_or(vals, idx >= torch.arange(mask.shape[0], device=mask.device), out=vals)
 50 |     return indices[vals]
 51 | 
 52 | 
 53 | @torch_compile(dynamic=True)
 54 | def get_iou_voting(boxes: Tensor, thresh: float):
 55 |     iou = bx.box_iou(boxes, boxes)
 56 |     iou *= iou > thresh  # consistent with nms thresh
 57 |     votes = iou.sum(-1)
 58 |     return votes
 59 | 
 60 | 
 61 | def filter_by_box_voting(boxes, thresh, min_vote, return_votes: bool = False):
 62 |     """Filter by box voting.
 63 | 
 64 |     Filter boxes by popular vote. A box receives a vote if it has an IoU larger than `thresh` with another box.
 65 |     Each box also votes for itself, hence, the smallest possible vote is 1.
 66 | 
 67 |     Args:
 68 |         boxes: Boxes.
 69 |         thresh: IoU threshold for two boxes to be considered redundant, counting as a vote for both boxes.
 70 |         min_vote: Minimum voting for a box to be accepted. A vote is the sum of IoUs of a box compared to all `boxes`,
 71 |             including itself. Hence, the smallest possible vote is 1.
 72 |         return_votes: Whether to return voting results.
 73 | 
 74 |     Returns:
 75 |         Keep indices and optionally voting results.
 76 |     """
 77 |     keep_indices = torch.arange(len(boxes), device=boxes.device, dtype=torch.int)
 78 |     votes = get_iou_voting(boxes, thresh)
 79 |     votes_mask = votes >= min_vote
 80 |     keep_indices = keep_indices[votes_mask]
 81 |     if return_votes:
 82 |         return keep_indices, votes[votes_mask]
 83 |     return keep_indices
 84 | 
 85 | 
 86 | def contours2boxes(contours, axis=-2):
 87 |     """Contours to boxes.
 88 | 
 89 |     Converts contours to bounding boxes in (x0, y0, x1, y1) format.
 90 | 
 91 |     Args:
 92 |         contours: Contours as Tensor[(..., )num_points, 2]
 93 |         axis: The ``num_points`` axis.
 94 | 
 95 |     Returns:
 96 | 
 97 |     """
 98 |     return torch.cat((contours.min(axis).values, contours.max(axis).values), axis + (axis < 0))
 99 | 
100 | 
101 | # implementation adapted from torchvision.ops.boxes._box_inter_union
102 | def _pairwise_box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]:
103 |     area1 = bx.box_area(boxes1)
104 |     area2 = bx.box_area(boxes2)
105 |     lt = torch.maximum(boxes1[:, :2], boxes2[:, :2])
106 |     rb = torch.minimum(boxes1[:, 2:], boxes2[:, 2:])
107 |     wh = bx._upcast(rb - lt).clamp(min=0)
108 |     intersection = torch.prod(wh, dim=1)
109 |     union = area1 + area2 - intersection
110 |     return intersection, union
111 | 
112 | 
113 | def pairwise_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
114 |     inter, union = _pairwise_box_inter_union(boxes1, boxes2)
115 |     return torch.abs(inter / union)
116 | 
117 | 
118 | # implementation adapted from torchvision.ops.boxes.generalized_box_iou
119 | def pairwise_generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
120 |     inter, union = _pairwise_box_inter_union(boxes1, boxes2)
121 |     iou = inter / union
122 |     lti = torch.minimum(boxes1[:, :2], boxes2[:, :2])
123 |     rbi = torch.maximum(boxes1[:, 2:], boxes2[:, 2:])
124 |     whi = bx._upcast(rbi - lti).clamp(min=0)
125 |     areai = torch.prod(whi, dim=1)
126 |     return iou - (areai - union) / areai
127 | 


--------------------------------------------------------------------------------
/celldetection/ops/commons.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from torch import Tensor
  3 | import torch.nn.functional as F
  4 | from typing import List
  5 | 
  6 | __all__ = ['downsample_labels', 'padded_stack2d', 'split_spatially', 'minibatch_std_layer', 'strided_upsampling2d',
  7 |            'interpolate_vector', 'pad_to_div', 'pad_to_size', 'spatial_mean', 'process_scores', 'equal_size']
  8 | 
  9 | 
 10 | def equal_size(x, reference, mode='bilinear', align_corners=False):
 11 |     if reference.shape[2:] != x.shape[2:]:  # 337 ns
 12 |         # bilinear: 3.79 ms for (128, 128) to (512, 512)
 13 |         # bicubic: 11.5 ms for (128, 128) to (512, 512)
 14 |         x = F.interpolate(x, reference.shape[2:],
 15 |                           mode=mode, align_corners=align_corners)
 16 |     return x
 17 | 
 18 | 
 19 | def _apply_score_bounds(scores, scores_lower_bound, scores_upper_bound):
 20 |     if scores_upper_bound is not None:
 21 |         assert scores_upper_bound.ndim >= 4, (f'Please make sure scores_upper_bound comes in NCHW format: '
 22 |                                               f'{scores_upper_bound.shape}')
 23 |         assert scores_upper_bound.dtype.is_floating_point, (f'Please make sure to pass scores_upper_bound as float '
 24 |                                                             f'instead of {scores_upper_bound.dtype}')
 25 |         scores = torch.minimum(scores, equal_size(scores_upper_bound, scores))
 26 |     if scores_lower_bound is not None:
 27 |         assert scores_lower_bound.ndim >= 4, (f'Please make sure scores_upper_bound comes in NCHW format: '
 28 |                                               f'{scores_lower_bound.shape}')
 29 |         assert scores_lower_bound.dtype.is_floating_point, (f'Please make sure to pass scores_upper_bound as float '
 30 |                                                             f'instead of {scores_lower_bound.dtype}')
 31 |         scores = torch.maximum(scores, equal_size(scores_lower_bound, scores))
 32 |     return scores
 33 | 
 34 | 
 35 | def process_scores(scores, score_channels, score_thresh, scores_lower_bound, scores_upper_bound):
 36 |     score_bounds = scores_lower_bound, scores_upper_bound
 37 |     if score_channels == 1:
 38 |         scores = _apply_score_bounds(torch.sigmoid(scores), *score_bounds)
 39 |         classes = torch.squeeze((scores > score_thresh).long(), 1)
 40 |     elif score_channels == 2:
 41 |         scores = _apply_score_bounds(F.softmax(scores, dim=1)[:, 1:2], *score_bounds)
 42 |         classes = torch.squeeze((scores > score_thresh).long(), 1)
 43 |     elif score_channels > 2:
 44 |         scores = _apply_score_bounds(F.softmax(scores, dim=1), *score_bounds)
 45 |         classes = torch.argmax(scores, dim=1).long()
 46 |     else:
 47 |         raise ValueError
 48 |     return scores, classes
 49 | 
 50 | 
 51 | def downsample_labels(inputs, size: List[int]):
 52 |     """
 53 | 
 54 |     Down-sample via max-pooling and interpolation
 55 | 
 56 |     Notes:
 57 |         - Downsampling can lead to loss of labeled instances, both during max pooling and interpolation.
 58 |         - Typical timing: 0.08106 ms for 256x256
 59 | 
 60 |     Args:
 61 |         inputs: Label Tensor to resize. Shape (n, c, h, w)
 62 |         size: Tuple containing target height and width.
 63 | 
 64 |     Returns:
 65 | 
 66 |     """
 67 |     sizeh, sizew = size  # for torchscript
 68 |     if inputs.shape[-2:] == (sizeh, sizew):
 69 |         return inputs
 70 |     if inputs.dtype != torch.float32:
 71 |         inputs = inputs.float()
 72 |     h, w = inputs.shape[-2:]
 73 |     th, tw = size
 74 |     k = h // th, w // tw
 75 |     r = F.max_pool2d(inputs, k, k)
 76 |     if r.shape[-2:] != (sizeh, sizew):
 77 |         r = F.interpolate(r, size, mode='nearest')
 78 |     return r
 79 | 
 80 | 
 81 | def padded_stack2d(*images, dim=0) -> Tensor:
 82 |     """Padding stack.
 83 | 
 84 |     Stacks 2d images along given axis.
 85 |     Spatial dimensions are padded according to largest height/width.
 86 | 
 87 |     Args:
 88 |         *images: Tensor[..., h, w]
 89 |         dim: Stack dimension.
 90 | 
 91 |     Returns:
 92 |         Tensor
 93 |     """
 94 |     ts = tuple(max((i.shape[j] for i in images)) for j in range(-2, 0))
 95 |     images = [F.pad(i, [0, ts[1] - i.shape[-1], 0, ts[0] - i.shape[-2]]) for i in images]
 96 |     return torch.stack(images, dim=dim)
 97 | 
 98 | 
 99 | def split_spatially(x, size):
100 |     """Split spatially.
101 | 
102 |     Splits spatial dimensions of Tensor ``x`` into patches of given ``size`` and adds the patches
103 |     to the batch dimension.
104 | 
105 |     Args:
106 |         x: Input Tensor[n, c, h, w, ...].
107 |         size: Patch size of the splits.
108 | 
109 |     Returns:
110 |         Tensor[n * h//height * w//width, c, height, width].
111 |     """
112 |     n, c = x.shape[:2]
113 |     spatial = x.shape[2:]
114 |     nd = len(spatial)
115 |     assert len(spatial) == len(size)
116 |     v = n, c
117 |     for cur, new in zip(spatial, size):
118 |         v += (cur // new, new)
119 |     perm = (0,) + tuple(range(2, nd * 2 + 1, 2)) + tuple(range(1, nd * 3, 2))
120 |     return x.view(v).permute(perm).reshape((-1, c) + tuple(size))
121 | 
122 | 
123 | def minibatch_std_layer(x, channels=1, group_channels=None, epsilon=1e-8):
124 |     """Minibatch standard deviation layer.
125 | 
126 |     The minibatch standard deviation layer first splits the batch dimension into slices of size ``group_channels``.
127 |     The channel dimension is split into ``channels`` slices. For the groups the standard deviation is calculated and
128 |     averaged over spatial dimensions and channel slice depth. The result is broadcasted to the spatial dimensions,
129 |     repeated for the batch dimension and then concatenated to the channel dimension of ``x``.
130 | 
131 |     References:
132 |         - https://arxiv.org/pdf/1710.10196.pdf
133 | 
134 |     Args:
135 |         x: Input Tensor[n, c, h, w].
136 |         channels: Number of averaged standard deviation channels.
137 |         group_channels: Number of channels per group. Default: batch size.
138 |         epsilon: Epsilon.
139 | 
140 |     Returns:
141 |         Tensor[n, c + channels, h, w].
142 |     """
143 |     n, c, h, w = x.shape
144 |     gc = min(group_channels or n, n)
145 |     cc, g = c // channels, n // gc
146 |     y = x.view(gc, g, channels, cc, h, w)
147 |     y = y.var(0, False).add(epsilon).sqrt().mean([2, 3, 4], True).squeeze(-1).repeat(gc, 1, h, w)
148 |     return torch.cat([x, y], 1)
149 | 
150 | 
151 | def strided_upsampling2d(x, factor=2, const=0):
152 |     """Strided upsampling.
153 | 
154 |     Upsample by inserting rows and columns filled with ``constant``.
155 | 
156 |     Args:
157 |         x: Tensor[n, c, h, w].
158 |         factor: Upsampling factor.
159 |         const: Constant used to fill inserted rows and columns.
160 | 
161 |     Returns:
162 |         Tensor[n, c, h*factor, w*factor].
163 |     """
164 |     n, c, h, w = x.shape
165 |     x_ = torch.zeros((n, c, h * factor, w * factor), dtype=x.dtype, device=x.device)
166 |     if const != 0:
167 |         x_.fill_(const)
168 |     x_[..., ::factor, ::factor] = x
169 |     return x_
170 | 
171 | 
172 | def interpolate_vector(v, size, **kwargs):
173 |     """Interpolate vector.
174 | 
175 |     Args:
176 |         v: Vector as ``Tensor[d]``.
177 |         size: Target size.
178 |         **kwargs: Keyword arguments for ``F.interpolate``
179 | 
180 |     Returns:
181 | 
182 |     """
183 |     return torch.squeeze(torch.squeeze(
184 |         F.interpolate(v[None, None], size, **kwargs), 0
185 |     ), 0)
186 | 
187 | 
188 | def pad_to_size(v, size, return_pad=False, **kwargs):
189 |     """Pad tp size.
190 | 
191 |     Applies padding to end of each dimension.
192 | 
193 |     Args:
194 |         v: Input Tensor.
195 |         size: Size tuple. Last element corresponds to last dimension of input `v`.
196 |         return_pad: Whether to return padding values.
197 |         **kwargs: Additional keyword arguments for `F.pad`.
198 | 
199 |     Returns:
200 |         Padded Tensor.
201 |     """
202 |     pad = []
203 |     for a, b in zip(size, v.shape[-len(size):]):
204 |         pad += [max(0, a - b), 0]
205 |     if any(pad):
206 |         v = F.pad(v, pad[::-1], **kwargs)
207 |     if return_pad:
208 |         return v, pad
209 |     return v
210 | 
211 | 
212 | def pad_to_div(v, div=32, nd=2, return_pad=False, **kwargs):
213 |     """Pad to div.
214 | 
215 |     Applies padding to input Tensor to make it divisible by `div`.
216 | 
217 |     Args:
218 |         v: Input Tensor.
219 |         div: Div tuple. If single integer, `nd` is used to define number of dimensions to pad.
220 |         nd: Number of dimensions to pad. Only used if `div` is not a tuple or list.
221 |         return_pad: Whether to return padding values.
222 |         **kwargs: Additional keyword arguments for `F.pad`.
223 | 
224 |     Returns:
225 |         Padded Tensor.
226 |     """
227 |     if not isinstance(div, (tuple, list)):
228 |         div = (div,) * nd
229 |     size = [(i // d + bool(i % d)) * d for i, d in zip(v.shape[-len(div):], div)]
230 |     return pad_to_size(v, size, return_pad=return_pad, **kwargs)
231 | 
232 | 
233 | def spatial_mean(x, keepdim=False):
234 |     spatial = tuple(range(2, x.ndim))
235 |     return torch.mean(x, spatial, keepdim=keepdim)
236 | 


--------------------------------------------------------------------------------
/celldetection/ops/draw.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | 
 3 | 
 4 | def draw_contours_(canvas, contours, close=True):
 5 |     """Draw contours.
 6 | 
 7 |     Draw ``contours`` on ``canvas``.
 8 | 
 9 |     Note:
10 |         This is an inplace operation.
11 | 
12 |     Args:
13 |         canvas: Tensor[h, w].
14 |         contours: Contours in (x, y) format. Tensor[num_contours, num_points, 2].
15 |         close: Whether to close contours. This is necessary if the last point of a contour is not equal to the first.
16 | 
17 |     """
18 |     if close:
19 |         contours = torch.cat((contours, contours[..., :1, :]), -2)
20 |     diff = torch.diff(contours, axis=1)
21 |     sign, diff = torch.sign(diff), torch.abs(diff)
22 |     err = diff[..., 0] - diff[..., 1]
23 |     x, y = contours[:, :-1, 0] + 0, contours[:, :-1, 1] + 0  # start point
24 |     x_, y_ = contours[:, 1:, 0], contours[:, 1:, 1]  # end point
25 |     labels = torch.broadcast_to(torch.arange(1, 1 + len(contours), device=canvas.device).to(canvas.dtype)[:, None],
26 |                                 x.shape)
27 |     m = torch.ones(x.shape, dtype=torch.bool, device=canvas.device)
28 |     while True:
29 |         canvas[y[m], x[m]] = labels[m]
30 |         # m: Select lines that are not finished; m_: Remove contours that are finished
31 |         m = m & (((x != x_) | (y != y_)) & (x >= 0) & (y >= 0) & (x < canvas.shape[1]) & (y < canvas.shape[0]))
32 |         m_ = torch.any(m, axis=-1)
33 |         m = m[m_]
34 |         if len(m) <= 0:
35 |             break
36 |         x, y, x_, y_, err, diff, sign, labels = (i[m_] for i in (x, y, x_, y_, err, diff, sign, labels))
37 |         err_ = 2 * err
38 |         sel = err_ > -diff[..., 1]
39 |         err[sel] -= diff[sel][..., 1]  # Note: torch cannot handle arr[mask, index] properly; not equivalent to numpy
40 |         x[sel] += sign[sel][..., 0]
41 |         sel = err_ < diff[..., 0]
42 |         err[sel] += diff[sel][..., 0]
43 |         y[sel] += sign[sel][..., 1]
44 | 


--------------------------------------------------------------------------------
/celldetection/ops/features.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from itertools import combinations_with_replacement
 3 | 
 4 | __all__ = ['texture_filter']
 5 | 
 6 | 
 7 | def symmetric_image(images: list):
 8 |     nd = images[0].ndim - 2
 9 |     # Array[..., nd, nd] required e.g. for eigvalsh
10 |     sym = torch.zeros(images[0].shape + (nd,) * 2, dtype=images[0].dtype, device=images[0].device)
11 |     for idx, (row, col) in enumerate(combinations_with_replacement(range(nd), 2)):
12 |         sym[..., row, col] = sym[..., col, row] = images[idx]
13 |     return sym
14 | 
15 | 
16 | def symmetric_compute_eigenvalues(images: list):
17 |     # List[Tensor[..., n, n]]
18 |     matrices = symmetric_image(images)
19 |     eigs = torch.flip(torch.linalg.eigvalsh(matrices), (-1,))  # eigenvals in decending order
20 |     eigs = torch.permute(eigs, (0, eigs.ndim - 1,) + tuple(range(1, eigs.ndim - 1)))
21 |     return eigs
22 | 
23 | 
24 | def texture_filter(gaussian_filtered, reshape=True):
25 |     """Texture filter.
26 | 
27 |     References:
28 |         - https://github.com/scikit-image/scikit-image/blob/0e28f6397a475db3f6755f03674c75ce02142bc3/skimage/feature/_basic_features.py#L10
29 | 
30 |     Args:
31 |         gaussian_filtered: Tensor[n, c, h, w]. Typically a batch of gaussian filtered images.
32 |         reshape: Whether to reshape result from Tensor[n, d, c, h, w] to Tensor[n, d * c, h, w] before returning it.
33 | 
34 |     Returns:
35 |         Tensor[n, d, c, h, w] or Tensor[n, d * c, h, w].
36 |     """
37 |     f_gf = torch.gradient(gaussian_filtered, dim=list(range(2, gaussian_filtered.ndim)))
38 |     axes = combinations_with_replacement(range(2, gaussian_filtered.ndim), 2)
39 |     elems = [torch.gradient(f_gf[ax0 - 2], dim=ax1)[0] for ax0, ax1 in axes]
40 |     eigvals = symmetric_compute_eigenvalues(elems)
41 |     if reshape:
42 |         n, d, c, h, w = eigvals.shape
43 |         eigvals = eigvals.reshape(n, d * c, h, w)
44 |     return eigvals
45 | 


--------------------------------------------------------------------------------
/celldetection/ops/loss.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn.functional as F
  3 | from torch import Tensor
  4 | import numpy as np
  5 | import torchvision.ops.boxes as bx
  6 | from .boxes import pairwise_box_iou, pairwise_generalized_box_iou
  7 | 
  8 | __all__ = ['reduce_loss', 'log_margin_loss', 'margin_loss', 'r1_regularization', 'iou_loss']
  9 | 
 10 | 
 11 | def reduce_loss(x: Tensor, reduction: str, **kwargs):
 12 |     """Reduce loss.
 13 | 
 14 |     Reduces Tensor according to ``reduction``.
 15 | 
 16 |     Args:
 17 |         x: Input.
 18 |         reduction: Reduction method. Must be a symbol of ``torch``.
 19 |         **kwargs: Additional keyword arguments.
 20 | 
 21 |     Returns:
 22 |         Reduced Tensor.
 23 |     """
 24 |     if reduction == 'none':
 25 |         return x
 26 |     fn = getattr(torch, reduction, None)
 27 |     if fn is None:
 28 |         raise ValueError(f'Unknown reduction: {reduction}')
 29 |     return fn(x, **kwargs)
 30 | 
 31 | 
 32 | def log_margin_loss(inputs: Tensor, targets: Tensor, m_pos=.9, m_neg=None, exponent=1, reduction='mean', eps=1e-6):
 33 |     if m_neg is None:
 34 |         m_neg = 1 - m_pos
 35 | 
 36 |     pos = torch.pow(F.relu_(torch.log(m_pos / (inputs + eps))), exponent)
 37 |     neg = torch.pow(F.relu_(torch.log((1 - m_neg) / (1 - inputs + eps))), exponent)
 38 |     loss = targets * pos + (1 - targets) * neg
 39 |     return reduce_loss(loss, reduction)
 40 | 
 41 | 
 42 | def margin_loss(inputs: Tensor, targets: Tensor, m_pos=.9, m_neg=None, exponent=2, reduction='mean'):
 43 |     if m_neg is None:
 44 |         m_neg = 1 - m_pos
 45 | 
 46 |     pos = torch.pow(F.relu_(m_pos - inputs), exponent)
 47 |     neg = torch.pow(F.relu_(inputs - m_neg), exponent)
 48 |     loss = targets * pos + (1 - targets) * neg
 49 |     return reduce_loss(loss, reduction)
 50 | 
 51 | 
 52 | def r1_regularization(logits, inputs, gamma=1., reduction='sum'):
 53 |     r"""R1 regularization.
 54 | 
 55 |     A gradient penalty regularization.
 56 |     This regularization may for example be applied to a discriminator with real data:
 57 | 
 58 |     .. math::
 59 | 
 60 |         R_1(\psi) &:= \frac{\gamma}{2} \mathbb E_{ p_{\mathcal D}(x)} \left[\|\nabla D_\psi(x)\|^2\right]
 61 | 
 62 |     References:
 63 |         - https://arxiv.org/pdf/1801.04406.pdf (Eq. 9)
 64 |         - https://arxiv.org/pdf/1705.09367.pdf
 65 |         - https://arxiv.org/pdf/1711.09404.pdf
 66 | 
 67 |     Examples:
 68 |         >>> real.requires_grad_(True)
 69 |         ... real_logits = discriminator(real)
 70 |         ... loss_d_real = F.softplus(-real_logits)
 71 |         ... loss_d_r1 = r1_regularization(real_logits, real)
 72 |         ... loss_d_real = (loss_d_r1 + loss_d_real).mean()
 73 |         ... loss_d_real.backward()
 74 |         ... real.requires_grad_(False)
 75 | 
 76 |     Args:
 77 |         logits: Logits.
 78 |         inputs: Inputs.
 79 |         gamma: Gamma.
 80 |         reduction: How to reduce all non-batch dimensions. E.g. ``'sum'`` or ``'mean'``.
 81 | 
 82 |     Returns:
 83 |         Penalty Tensor[n].
 84 |     """
 85 |     grads = torch.autograd.grad(logits.sum(), inputs=inputs, create_graph=True, retain_graph=True, only_inputs=True)[0]
 86 |     penalty = reduce_loss(grads.square(), reduction, dim=list(range(1, grads.ndim)))
 87 |     return penalty * (gamma * .5)
 88 | 
 89 | 
 90 | def iou_loss(boxes, boxes_targets, reduction='mean', generalized=True, method='linear', min_size=None):
 91 |     if min_size is not None:  # eliminates invalid boxes
 92 |         keep = bx.remove_small_boxes(boxes, min_size)
 93 |         boxes, boxes_targets = (c[keep] for c in (boxes, boxes_targets))
 94 | 
 95 |     if generalized:
 96 |         iou = pairwise_generalized_box_iou(boxes, boxes_targets)  # Tensor[n]
 97 |     else:
 98 |         iou = pairwise_box_iou(boxes, boxes_targets)  # Tensor[n]
 99 | 
100 |     if method == 'log':
101 |         if generalized:
102 |             iou = iou * .5 + .5
103 |         loss = -torch.log(iou + 1e-8)
104 |     elif method == 'linear':
105 |         loss = 1 - iou
106 |     else:
107 |         raise ValueError
108 | 
109 |     loss = reduce_loss(loss, reduction=reduction)
110 |     return loss
111 | 
112 | 
113 | def box_npll_loss(uncertainty, boxes, boxes_targets, factor=10., sigmoid=False, epsilon=1e-8, reduction='mean',
114 |                   min_size=None):
115 |     """NPLL.
116 | 
117 |     References:
118 |         https://arxiv.org/abs/2006.15607
119 | 
120 |     Args:
121 |         uncertainty: Tensor[n, 4].
122 |         boxes: Tensor[n, 4].
123 |         boxes_targets: Tensor[n, 4].
124 |         sigmoid: Whether to apply the ``sigmoid`` function to ``uncertainty``.
125 |         factor: Uncertainty factor.
126 |         epsilon: Epsilon.
127 |         reduction: Loss reduction.
128 |         min_size: Minimum box size. May be used to remove degenerate boxes.
129 | 
130 |     Returns:
131 |         Loss.
132 |     """
133 |     if min_size is not None:  # eliminates invalid boxes
134 |         keep = bx.remove_small_boxes(boxes, min_size)
135 |         boxes, boxes_targets, uncertainty = (c[keep] for c in (boxes, boxes_targets, uncertainty))
136 |     delta_sq = torch.square((torch.sigmoid(uncertainty) if sigmoid else uncertainty) * factor)
137 |     a = torch.square(boxes - boxes_targets) / (2 * delta_sq + epsilon)
138 |     b = 0.5 * torch.log(delta_sq + epsilon)
139 |     iou = pairwise_box_iou(boxes, boxes_targets)  # Tensor[n]
140 |     loss = iou * ((a + b).sum(dim=1) + 2 * np.log(2 * np.pi))
141 |     loss = reduce_loss(loss, reduction=reduction)
142 |     return loss
143 | 


--------------------------------------------------------------------------------
/celldetection/ops/normalization.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | 
 3 | __all__ = ['pixel_norm']
 4 | 
 5 | 
 6 | def pixel_norm(x, dim=1, eps=1e-8):
 7 |     """Pixel normalization.
 8 | 
 9 |     References:
10 |         - https://papers.nips.cc/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf ("Local Repose Norm")
11 |         - https://arxiv.org/pdf/1710.10196.pdf
12 | 
13 |     Args:
14 |         x: Input Tensor.
15 |         dim: Dimension to normalize.
16 |         eps: Epsilon.
17 | 
18 |     Returns:
19 |         Normalized Tensor.
20 |     """
21 |     return x * torch.rsqrt(torch.mean(torch.square(x), dim=dim, keepdim=True) + eps)
22 | 


--------------------------------------------------------------------------------
/celldetection/optim/__init__.py:
--------------------------------------------------------------------------------
1 | from .lr_scheduler import *
2 | 


--------------------------------------------------------------------------------
/celldetection/optim/lr_scheduler.py:
--------------------------------------------------------------------------------
  1 | from torch.optim.lr_scheduler import MultiplicativeLR, SequentialLR as _SequentialLR, \
  2 |     ReduceLROnPlateau as _ReduceLROnPlateau
  3 | from torch.optim import Optimizer
  4 | from typing import Union, Callable, List
  5 | import warnings
  6 | from bisect import bisect_right
  7 | 
  8 | from ..util.util import has_argument
  9 | 
 10 | __all__ = ['WarmUp', 'SequentialLR', 'ReduceLROnPlateau']
 11 | 
 12 | 
 13 | def linear_schedule(step, steps):
 14 |     return 1. if step > steps else min(step / steps, 1.)
 15 | 
 16 | 
 17 | class WarmUp(MultiplicativeLR):
 18 |     def __init__(
 19 |             self,
 20 |             optimizer: Optimizer,
 21 |             steps: int,
 22 |             lr_lambda: Union[Callable[[int, int], float], List[Callable[[int, int], float]]] = None,
 23 |             last_epoch: int = -1,
 24 |             verbose: bool = False
 25 |     ):
 26 |         """WarmUp.
 27 | 
 28 |         Applies a scaling factor to learning rate for steps ``1`` to ``steps``.
 29 |         Applies no changes otherwise.
 30 |         ``WarmUp`` can be chained or used sequentally with other schedulers.
 31 | 
 32 |         Notes:
 33 |             - WarmUp sets learning rates based on captured initial learning rates (``base_lrs``).
 34 |             - Changes from other schedulers applied before WarmUp will be ovewritten.
 35 |             - Chaining ``WarmUp`` is equivalent to SequentialLR if other schedule dynamically manipulates
 36 |               optimizer.param_groups without relying on ``base_lrs``.
 37 | 
 38 |         Examples:
 39 |             >>> # First WarmUp, then other schedule
 40 |             >>> from torch.optim.lr_scheduler import SequentialLR, CosineAnnealingLR
 41 |             >>> from celldetection.optim import WarmUp
 42 |             >>> warmup_steps = 512
 43 |             >>> scheduler = SequentialLR(optimizer, [
 44 |             ...     WarmUp(optimizer, warmup_steps),  # warmup for 512 steps
 45 |             ...     CosineAnnealingLR(optimizer, T_max=512, eta_min=0.00001),  # after 512 steps switch to cosine ann.
 46 |             ... ], milestones=[warmup_steps])
 47 | 
 48 |             >>> # Chaining WarmUp and other schedule
 49 |             >>> from torch.optim.lr_scheduler import ChainedScheduler, StepLR
 50 |             >>> scheduler = ChainedScheduler([
 51 |             ...     StepLR(optimizer, 5, gamma=.99),
 52 |             ...     WarmUp(optimizer, 512),  # overwrites changes of previous scheduler during warmup
 53 |             ... ])
 54 | 
 55 |         Args:
 56 |             optimizer: Optimizer.
 57 |             steps: Number of warmup steps.
 58 |             lr_lambda: A function which computes a multiplicative factor given an integer parameter epoch, or a list
 59 |                 of such functions, one for each group in optimizer.param_groups.
 60 |             last_epoch: Last epoch.
 61 |             verbose: If ``True``, prints a message to stdout for each update. Default: ``False``.
 62 |         """
 63 |         if lr_lambda is None:
 64 |             lr_lambda = linear_schedule
 65 |         self.steps = steps
 66 |         super().__init__(optimizer, lr_lambda=lr_lambda, last_epoch=last_epoch, verbose=verbose)
 67 | 
 68 |     def get_lr(self):
 69 |         if not self._get_lr_called_within_step:
 70 |             warnings.warn("To get the last learning rate computed by the scheduler, "
 71 |                           "please use `get_last_lr()`.", UserWarning)
 72 |         if self.last_epoch <= self.steps:
 73 |             return [lr * lmbda(self.last_epoch, self.steps) for lmbda, lr in zip(self.lr_lambdas, self.base_lrs)]
 74 |         return [group['lr'] for group in self.optimizer.param_groups]
 75 | 
 76 | 
 77 | class SequentialLR(_SequentialLR):
 78 | 
 79 |     def step(self, metrics=None):  # fixes TypeError caused by use of metric
 80 |         self.last_epoch += 1
 81 |         idx = bisect_right(self._milestones, self.last_epoch)
 82 |         scheduler = self._schedulers[idx]
 83 |         if idx > 0 and self._milestones[idx - 1] == self.last_epoch:
 84 |             scheduler.step(0)
 85 |         else:
 86 | 
 87 |             if metrics is None:
 88 |                 scheduler.step()
 89 |             else:
 90 |                 # Some schedulers require metrics, others do not. If not handled here, it'll raise a TypeError
 91 |                 if has_argument(scheduler.step, 'metric', 'metrics', mode='any'):
 92 |                     scheduler.step(metrics)
 93 |                 else:
 94 |                     scheduler.step()
 95 | 
 96 |         self._last_lr = scheduler.get_last_lr()
 97 | 
 98 | 
 99 | class ReduceLROnPlateau(_ReduceLROnPlateau):
100 |     def __init__(self, optimizer, mode='min', factor=0.1, patience=10,
101 |                  threshold=1e-4, threshold_mode='rel', cooldown=0,
102 |                  min_lr=0, eps=1e-8, warmup=1, verbose="deprecated"):
103 |         """
104 |         Initializes the ReduceLROnPlateau object. This scheduler decreases the learning rate
105 |         when a metric has stopped improving, which is commonly used to fine-tune a model in
106 |         machine learning.
107 | 
108 |         Notes:
109 |             - Adds the warmup option to PyTorch's ``ReduceLROnPlateau``.
110 | 
111 |         Args:
112 |             optimizer (Optimizer): Wrapped optimizer.
113 |             mode (str): One of `min` or `max`. In `min` mode, the learning rate will be reduced
114 |                         when the quantity monitored has stopped decreasing; in `max` mode, it will
115 |                         be reduced when the quantity monitored has stopped increasing. Default: 'min'.
116 |             factor (float): Factor by which the learning rate will be reduced. `new_lr = lr * factor`.
117 |                             Default: 0.1.
118 |             patience (int): Number of epochs with no improvement after which learning rate will be
119 |                             reduced. Default: 10.
120 |             threshold (float): Threshold for measuring the new optimum, to only focus on significant
121 |                                changes. Default: 1e-4.
122 |             threshold_mode (str): One of `rel`, `abs`. In `rel` mode, dynamic_threshold = best * (1 + threshold)
123 |                                   in 'max' mode or best * (1 - threshold) in `min` mode. In `abs` mode,
124 |                                   dynamic_threshold = best + threshold in `max` mode or best - threshold in
125 |                                   `min` mode. Default: 'rel'.
126 |             cooldown (int): Number of epochs to wait before resuming normal operation after lr has been reduced.
127 |                             Default: 0.
128 |             min_lr (float or list): A scalar or a list of scalars. A lower bound on the learning rate of
129 |                                     all param groups or each group respectively. Default: 0.
130 |             eps (float): Minimal decay applied to lr. If the difference between new and old lr is smaller
131 |                          than eps, the update is ignored. Default: 1e-8.
132 |             warmup (int): Number of epochs to wait before initially starting normal operation. Default: 1.
133 |             verbose (str): Deprecated argument. Not used. Default: "deprecated".
134 |         """
135 |         super().__init__(optimizer=optimizer, mode=mode, factor=factor, patience=patience,
136 |                          threshold=threshold, threshold_mode=threshold_mode, cooldown=cooldown,
137 |                          min_lr=min_lr, eps=eps, verbose=verbose)
138 |         self.warmup_counter = int(warmup)  # ignores bad epochs for the first number of `warmup` steps
139 | 
140 |     def get_last_lr(self):  # required by PyTorch functions to be implemented right here
141 |         return self._last_lr
142 | 
143 |     def step(self, metrics, epoch=None):
144 |         best_ = None
145 |         if self.warmup_counter:
146 |             self.warmup_counter -= 1
147 |             best_ = self.best
148 | 
149 |         res = super().step(metrics, epoch)
150 |         if best_ is not None:
151 |             self.best = best_
152 |             self.num_bad_epochs = 0  # ignore any bad epochs in warmup
153 |         return res
154 | 


--------------------------------------------------------------------------------
/celldetection/util/__init__.py:
--------------------------------------------------------------------------------
1 | from .util import *
2 | from .timer import *
3 | from .schedule import *
4 | from .shm_cache import *
5 | from .logging import *
6 | 


--------------------------------------------------------------------------------
/celldetection/util/logging.py:
--------------------------------------------------------------------------------
 1 | import numpy as np
 2 | from typing import Union, Optional, Dict, List
 3 | from ..visualization.images import figure2img
 4 | 
 5 | __all__ = []
 6 | 
 7 | 
 8 | def register(obj):
 9 |     __all__.append(obj.__name__)
10 |     return obj
11 | 
12 | 
13 | @register
14 | def log_figure(
15 |         logger,
16 |         tag: str,
17 |         figure,
18 |         global_step: Optional[int] = None,
19 |         close: Optional[bool] = True,
20 |         walltime: Optional[float] = None
21 | ):
22 |     if hasattr(logger, 'add_image'):
23 |         if isinstance(figure, list):
24 |             # Note that tensorboard cannot handle RGBA in NCHW, hence A channel must be removed.
25 |             logger.add_image(tag, figure2img(figure, close=close, transpose=True, remove_alpha=True),
26 |                              global_step, walltime, dataformats='NCHW')
27 |         else:
28 |             logger.add_image(tag, figure2img(figure, close=close, transpose=True, remove_alpha=False), global_step,
29 |                              walltime, dataformats='CHW')
30 | 


--------------------------------------------------------------------------------
/celldetection/util/shm_cache.py:
--------------------------------------------------------------------------------
  1 | import hashlib
  2 | from os.path import isfile, dirname, basename, isdir, join, splitext, islink, abspath
  3 | from os import remove, makedirs, rmdir, unlink, symlink, cpu_count
  4 | from shutil import copy2
  5 | from concurrent.futures import ThreadPoolExecutor
  6 | from warnings import warn
  7 | import pytorch_lightning as pl
  8 | from .util import compare_file_hashes
  9 | 
 10 | __all__ = ['ShmCache']
 11 | 
 12 | 
 13 | def get_hash(filename):
 14 |     hash_obj = hashlib.sha256()
 15 |     hash_obj.update(abspath(filename).encode())
 16 |     return hash_obj.hexdigest()
 17 | 
 18 | 
 19 | def get_dst(filename, root):
 20 |     bn = basename(filename)
 21 |     dn = dirname(filename)
 22 |     cache_dir = join(root, get_hash(dn))
 23 |     dst = join(cache_dir, get_hash(bn)) + splitext(bn)[1]
 24 |     return cache_dir, dst
 25 | 
 26 | 
 27 | def _resolve_filename(filename, root):
 28 |     postfix = ''
 29 |     if isinstance(filename, dict):  # allow to pass dict, specifying postfix and filename
 30 |         postfix = filename.get('postfix', postfix)
 31 |         filename = filename['filename']
 32 |     cache_dir, dst = get_dst(filename[:-len(postfix)] if len(postfix) else filename, root)
 33 |     dst = dst + postfix
 34 |     return cache_dir, dst, filename
 35 | 
 36 | 
 37 | def copy_file(filename, root, check_hash=True, verbose=True):
 38 |     cache_dir, dst, filename = _resolve_filename(filename, root)
 39 |     makedirs(cache_dir, exist_ok=True)
 40 |     if verbose:
 41 |         print('Copy', filename, '-->', dst)
 42 |     copy2(filename, dst)
 43 |     if check_hash:
 44 |         assert compare_file_hashes(filename, dst), f'File hashes did not match after copying: {filename}, {dst}'
 45 |     return dst
 46 | 
 47 | 
 48 | def link_file(filename, root, verbose=True):
 49 |     cache_dir, dst, filename = _resolve_filename(filename, root)
 50 |     makedirs(cache_dir, exist_ok=True)
 51 |     if verbose:
 52 |         print('Link', abspath(filename), 'as', abspath(dst))
 53 |     if not islink(dst):
 54 |         try:
 55 |             symlink(abspath(filename), abspath(dst))
 56 |         except Exception as e:
 57 |             print(f"Error creating symlink for {filename} at {dst}: {e}")
 58 |             return None
 59 |     elif verbose:
 60 |         print('Already linked', abspath(filename), 'as', abspath(dst))
 61 |     return dst
 62 | 
 63 | 
 64 | def delete_file_or_directory(filename, root, verbose=True):
 65 |     if filename is None:
 66 |         return True
 67 |     import os
 68 |     assert filename.startswith(root) and '..' not in filename
 69 |     try:
 70 |         if islink(filename):
 71 |             if verbose:
 72 |                 print('Unlink', filename)
 73 |             unlink(filename)
 74 |         elif isdir(filename):
 75 |             assert os.path.realpath(filename).startswith(root)
 76 |             if verbose:
 77 |                 print('Delete', filename)
 78 |             rmdir(filename)
 79 |         elif isfile(filename):
 80 |             assert os.path.realpath(filename).startswith(root)
 81 |             if verbose:
 82 |                 print('Delete', filename)
 83 |             remove(filename)
 84 |         elif verbose:
 85 |             print('Already gone:', filename)
 86 |         return True
 87 |     except Exception as e:
 88 |         print(f"Error deleting {filename}: {e}")
 89 |         return False
 90 | 
 91 | 
 92 | class ShmCache:
 93 |     def __init__(self, filenames, link_filenames=None, directory='/dev/shm/celldetection', verbose=True):
 94 |         """Shared Memory Cache.
 95 | 
 96 |         Creates copies or symbolic links for all files in `filenames` and `link_filenames`, respectively.
 97 |         This serves as an easy-to-use interface for caching training data.
 98 | 
 99 |         Note:
100 |             When using `ShmCache` in a distributed environment, failing processes should call `teardown`
101 |             to remove all cached files before the script terminates. Note that other running processes
102 |             will not find any cached files after `teardown` has been called. Hence, it is expected
103 |             that exceptions from failing ranks are accompanied by `OSError: Unable to open file (open failed)`
104 |             or `FileNotFoundError` from running ranks.
105 | 
106 |         Args:
107 |             filenames: Filenames (to be copied).
108 |             link_filenames: Filenames (to be symbolically linked).
109 |             directory: Target directory.
110 |             verbose: Verbosity.
111 |         """
112 |         self.filenames = filenames
113 |         self.link_filenames = link_filenames
114 |         self.directory = directory
115 |         self.cached_files = None
116 |         self.linked_files = None
117 |         self.subs = None
118 |         self.verbose = verbose
119 | 
120 |     @property
121 |     def setup_complete(self):
122 |         return self.subs is not None
123 | 
124 |     def setup(self, max_workers=None):
125 |         """Setup.
126 | 
127 |         Copies all files specified in the constructor via `filenames`.
128 |         Also creates symlinks for all files specified in the constructor via `link_filenames`.
129 | 
130 |         Args:
131 |             max_workers: Max number of workers.
132 | 
133 |         Returns:
134 |             List of cached file names.
135 |         """
136 |         if max_workers is None:
137 |             max_workers = cpu_count() or 1
138 | 
139 |         if not isdir(self.directory):
140 |             if self.verbose:
141 |                 print('Creating', self.directory)
142 |             makedirs(self.directory, exist_ok=True)
143 | 
144 |         with ThreadPoolExecutor(max_workers=max_workers) as executor:
145 |             self.subs = set()
146 |             if self.filenames is not None:
147 |                 self.cached_files = list(
148 |                     executor.map(lambda f: copy_file(f, self.directory, verbose=self.verbose), self.filenames))
149 |                 self.subs.update({dirname(f) for f in self.cached_files})
150 |             if self.link_filenames is not None:
151 |                 self.linked_files = list(
152 |                     executor.map(lambda f: link_file(f, self.directory, verbose=self.verbose), self.link_filenames))
153 |                 self.subs.update({dirname(f) for f in self.linked_files})
154 | 
155 |         return self.cached_files
156 | 
157 |     def teardown(self, max_workers=None):
158 |         """Teardown.
159 | 
160 |         Removes all copies, unlinks all symlinks and removes all subdirectories that contained copies or symlinks.
161 | 
162 |         Args:
163 |             max_workers: Max number of workers.
164 |         """
165 |         if max_workers is None:
166 |             max_workers = cpu_count() or 1
167 | 
168 |         if self.cached_files is None:
169 |             return
170 | 
171 |         with ThreadPoolExecutor(max_workers=max_workers) as executor:
172 |             # Remove cached files
173 |             results = list(executor.map(lambda f: delete_file_or_directory(f, self.directory, verbose=self.verbose),
174 |                                         self.cached_files))
175 | 
176 |             # Remove linked files
177 |             if self.linked_files is not None:
178 |                 results = list(executor.map(lambda f: delete_file_or_directory(f, self.directory, verbose=self.verbose),
179 |                                             self.linked_files))
180 | 
181 |             # Remove cached dirs (must be empty)
182 |             if self.subs is not None:
183 |                 results += list(
184 |                     executor.map(lambda f: delete_file_or_directory(f, self.directory, verbose=self.verbose),
185 |                                  self.subs))
186 | 
187 |         self.cached_files = self.linked_files = self.subs = None
188 | 
189 |         if not all(results):
190 |             warn(f'Could not remove all cached files in {self.directory}. Removed {sum(results)}/{len(results)}')
191 | 


--------------------------------------------------------------------------------
/celldetection/util/timer.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import gc
 3 | from time import time
 4 | import numpy as np
 5 | 
 6 | __all__ = ['start_timer', 'stop_timer', 'print_timing']
 7 | 
 8 | TIMINGS = {}
 9 | 
10 | 
11 | def convert_seconds(seconds):
12 |     if seconds // 60:
13 |         minutes = seconds // 60
14 |         seconds %= 60
15 |     else:
16 |         minutes = 0
17 |     if minutes // 60:
18 |         hours = minutes // 60
19 |         minutes %= 60
20 |     else:
21 |         hours = 0
22 |     if hours // 24:
23 |         days = hours // 24
24 |         hours %= 24
25 |     else:
26 |         days = 0
27 |     return int(days), int(hours), int(minutes), int(np.round(seconds))
28 | 
29 | 
30 | def seconds_to_str(seconds):
31 |     s = [f'{i} {n[:-1] if i == 1 else n}' for i, n in
32 |          zip(convert_seconds(seconds), ('days', 'hours', 'minutes', 'seconds'))]
33 |     s = ', '.join(s)
34 |     return s
35 | 
36 | 
37 | def print_timing(name, seconds):
38 |     if seconds >= 1:
39 |         val = np.round(seconds, 3)
40 |         fill = 75 - len(name) - len(str(val))
41 |         print(name + ':', ' ' * fill, val, 's')
42 |         return
43 |     seconds *= 1000
44 |     if seconds >= 1:
45 |         val = np.round(seconds, 3)
46 |         fill = 75 - len(name) - len(str(val))
47 |         print(name + ':', ' ' * fill, val, 'ms')
48 |         return
49 |     seconds *= 1000
50 |     if seconds >= 1:
51 |         val = np.round(seconds, 3)
52 |         fill = 75 - len(name) - len(str(val))
53 |         print(name + ':', ' ' * fill, val, 'µs')
54 |         return
55 |     seconds *= 1000
56 |     val = np.round(seconds, 3)
57 |     fill = 75 - len(name) - len(str(val))
58 |     print(name + ':', ' ' * fill, val, 'ns')
59 | 
60 | 
61 | def start_timer(name, cuda=True, collect=True):
62 |     """Keyword PyTorch timer.
63 | 
64 |     Can be used to measure PyTorch GPU times.
65 | 
66 |     Args:
67 |         name: Keyword
68 | 
69 |     Returns:
70 | 
71 |     """
72 |     if collect:
73 |         gc.collect()
74 |     if cuda:
75 |         torch.cuda.empty_cache()
76 |         torch.cuda.reset_peak_memory_stats()
77 |         torch.cuda.synchronize()
78 |     TIMINGS[name] = time()
79 | 
80 | 
81 | def stop_timer(name, cuda=True, verbose=True):
82 |     if cuda:
83 |         torch.cuda.synchronize()
84 |     delta = time() - TIMINGS[name]
85 |     if verbose:
86 |         print_timing(name, delta)
87 |     return delta
88 | 


--------------------------------------------------------------------------------
/celldetection/visualization/__init__.py:
--------------------------------------------------------------------------------
1 | from .images import *
2 | from .cmaps import *
3 | 


--------------------------------------------------------------------------------
/celldetection/visualization/cmaps.py:
--------------------------------------------------------------------------------
 1 | from matplotlib import pyplot as plt
 2 | import numpy as np
 3 | from numpy import ndarray
 4 | from typing import Union
 5 | import cv2
 6 | 
 7 | __all__ = ['label_cmap', 'random_colors_hsv']
 8 | 
 9 | 
10 | def random_colors_hsv(num, hue_range=(0, 180), saturation_range=(60, 133), value_range=(180, 256), ubyte=False):
11 |     colors, = cv2.cvtColor(np.stack((
12 |         np.random.randint(*hue_range, num),
13 |         np.random.randint(*saturation_range, num),
14 |         np.random.randint(*value_range, num),
15 |     ), 1).astype('uint8')[None], cv2.COLOR_HSV2RGB)
16 |     if not ubyte:
17 |         colors = colors / 255
18 |     return colors
19 | 
20 | 
21 | def label_cmap(labels: ndarray, colors: Union[str, ndarray] = 'rand', zero_val: Union[float, tuple, list] = 0.,
22 |                rgba: bool = True, alpha: float = None, reduce_axis=2, neg_alpha_factor=.5, ubyte=False):
23 |     """Label colormap.
24 | 
25 |     Applies a colormap to a label image.
26 | 
27 |     Args:
28 |         labels: Label image. Typically Array[h, w].
29 |         colors: Either 'rand' or one of ['Pastel1', 'Pastel2', 'Paired', 'Accent', 'Dark2', 'Set1', 'Set2', 'Set3',
30 |             'tab10', 'tab20', 'tab20b', 'tab20c'] (see matplotlib's qualitative colormaps) or Array[n, c].
31 |         zero_val: Special color for the zero label (usually background).
32 |         rgba: Whether to add an alpha channel to rgb colors.
33 |         alpha: Specific alpha value.
34 |         reduce_axis: Axis to be reduced with alpha compositing (e.g. channel axis).
35 |         neg_alpha_factor: Factor for alpha channel of colorized negative labels.
36 | 
37 |     Returns:
38 |         Mapped labels. E.g. rgba mapping Array[h, w] -> Array[h, w, 4].
39 |     """
40 |     assert issubclass(labels.dtype.type, np.integer), 'Pass labels as an integer array.'
41 |     if isinstance(colors, str):
42 |         if colors == 'rand':
43 |             lim = np.max(labels)
44 |             # extra colors for negative labels
45 |             neg = labels[labels < 0]
46 |             if neg.size:
47 |                 lim += np.abs(neg.min())
48 |             colors = random_colors_hsv(max(1, min(9999, lim)))
49 |         else:
50 |             colors = plt.get_cmap(colors).colors
51 |     if isinstance(colors, (tuple, list)):
52 |         colors = np.array(colors)
53 |     if rgba and colors.shape[1] == 3:
54 |         colors = np.concatenate((colors, np.ones((len(colors), 1))), -1)
55 |     if alpha is not None and colors.shape[1] == 4:
56 |         colors[:, -1] = alpha
57 |     neg_m = labels < 0
58 |     if zero_val is not None:
59 |         labels = np.copy(labels)
60 |         m = labels != 0
61 |         labels[m] = labels[m] % len(colors) + 1
62 |         colors = np.concatenate((np.ones_like(colors[:1]) * zero_val, colors))
63 |     if ubyte:
64 |         from skimage import img_as_ubyte
65 |         colors = img_as_ubyte(colors)
66 |     res = colors[labels]
67 |     if rgba and neg_m.any():
68 |         res[neg_m, -1] *= neg_alpha_factor
69 | 
70 |     if reduce_axis is not None and labels.ndim > 2:  # reduce label channels by alpha weighted average
71 |         weight = res[..., -1]
72 |         weight = weight / (weight.sum(reduce_axis, keepdims=True).astype('float32') + 1e-12)
73 |         if reduce_axis < 0:
74 |             reduce_axis = labels.ndim + reduce_axis
75 |         res = (weight[..., None] * res).sum(reduce_axis).astype(res.dtype)  # todo: consider checking bounds
76 | 
77 |     return res
78 | 


--------------------------------------------------------------------------------
/celldetection_scripts/__init__.py:
--------------------------------------------------------------------------------
1 | from .cpn_inference import cpn_inference
2 | 


--------------------------------------------------------------------------------
/demos/demo-binary.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "metadata": {},
  6 |    "source": [
  7 |     "# Contour Proposal Networks － How to detect objects"
  8 |    ]
  9 |   },
 10 |   {
 11 |    "cell_type": "code",
 12 |    "execution_count": null,
 13 |    "metadata": {},
 14 |    "outputs": [],
 15 |    "source": [
 16 |     "import torch\n",
 17 |     "if torch.cuda.is_available():\n",
 18 |     "    torch.backends.cudnn.benchmark = True\n",
 19 |     "import celldetection as cd\n",
 20 |     "from celldetection import models, toydata\n",
 21 |     "import numpy as np\n",
 22 |     "import os\n",
 23 |     "from collections import OrderedDict\n",
 24 |     "import torch.nn as nn\n",
 25 |     "import torch.nn.functional as F\n",
 26 |     "from torch.utils.data import Dataset, DataLoader\n",
 27 |     "from tqdm import tqdm\n",
 28 |     "from matplotlib import pyplot as plt\n",
 29 |     "from torch.cuda.amp import GradScaler, autocast"
 30 |    ]
 31 |   },
 32 |   {
 33 |    "cell_type": "markdown",
 34 |    "metadata": {},
 35 |    "source": [
 36 |     "# 1. Config\n",
 37 |     "You can save a config `conf.to_json(filename)` and load it with `cd.Config.from_json(filename)`."
 38 |    ]
 39 |   },
 40 |   {
 41 |    "cell_type": "code",
 42 |    "execution_count": null,
 43 |    "metadata": {},
 44 |    "outputs": [],
 45 |    "source": [
 46 |     "conf = cd.Config(\n",
 47 |     "    # data\n",
 48 |     "    in_channels=3,\n",
 49 |     "    classes=2,\n",
 50 |     "    shuffle=True,\n",
 51 |     "    bg_fg_dists=(0.8, 0.85),\n",
 52 |     "    \n",
 53 |     "    # augmentation (schema: <class_name>:<kwargs>)\n",
 54 |     "    augmentation=OrderedDict({\n",
 55 |     "        'Transpose': {'p': 0.5},  # see: https://albumentations.ai/docs/\n",
 56 |     "        'RandomRotate90': {'p': 0.5},\n",
 57 |     "    }),\n",
 58 |     "    \n",
 59 |     "    # cpn\n",
 60 |     "    cpn='CpnU22',  # see https://git.io/JOnWX for alternatives\n",
 61 |     "    score_thresh=.9,\n",
 62 |     "    nms_thresh=.5,\n",
 63 |     "    contour_head_stride=8,\n",
 64 |     "    order=7,  # the higher, the more complex shapes can be detected\n",
 65 |     "    samples=128,  # number of coordinates per contour\n",
 66 |     "    refinement_iterations=3,\n",
 67 |     "    refinement_buckets=6,\n",
 68 |     "    inputs_mean=.5,\n",
 69 |     "    inputs_std=.5,\n",
 70 |     "    tweaks={\n",
 71 |     "        'BatchNorm2d': {'momentum': 0.05}\n",
 72 |     "    },\n",
 73 |     "    \n",
 74 |     "    # optimizer\n",
 75 |     "    optimizer={'Adadelta': {'lr': 1., 'rho': 0.9}},\n",
 76 |     "    scheduler={'StepLR': {'step_size': 5, 'gamma': .99}},\n",
 77 |     "    \n",
 78 |     "    # training\n",
 79 |     "    epochs=100,\n",
 80 |     "    steps_per_epoch=8 * 512,\n",
 81 |     "    batch_size=8,\n",
 82 |     "    amp=torch.cuda.is_available(),\n",
 83 |     "    \n",
 84 |     "    # misc\n",
 85 |     "    num_workers=0,\n",
 86 |     "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
 87 |     ")\n",
 88 |     "print(conf)"
 89 |    ]
 90 |   },
 91 |   {
 92 |    "cell_type": "markdown",
 93 |    "metadata": {},
 94 |    "source": [
 95 |     "# 2. Toy Dataset\n",
 96 |     "Each item of the dataset consists of an image and a label image.\n",
 97 |     "Just put your label image into the `CPNTargetGenerator` via `gen.feed(labels=labels)` and it will generate training targets.\n",
 98 |     "\n",
 99 |     "If you do not have a label image, you might have one of those:\n",
100 |     "- A list of masks, each mask shows a single object. Then you can use:\n",
101 |     "```\n",
102 |     "labels = cd.data.unary_masks2labels([mask1, mask2, ..., maskN])\n",
103 |     "```\n",
104 |     "- A list of masks, each mask may show multiple objects, but touching objects were assigned different numbers. Then you can use:\n",
105 |     "```\n",
106 |     "labels = cd.data.masks2labels([mask1, mask2, ..., maskN])\n",
107 |     "```\n"
108 |    ]
109 |   },
110 |   {
111 |    "cell_type": "code",
112 |    "execution_count": null,
113 |    "metadata": {},
114 |    "outputs": [],
115 |    "source": [
116 |     "class Dataset:\n",
117 |     "    def __init__(self, samples, order, max_bg_dist, min_fg_dist, transforms=None, items=2**12):\n",
118 |     "        self.gen = cd.data.CPNTargetGenerator(\n",
119 |     "            samples=samples,\n",
120 |     "            order=order,\n",
121 |     "            min_fg_dist=min_fg_dist,\n",
122 |     "            max_bg_dist=max_bg_dist,\n",
123 |     "        )\n",
124 |     "        self._items = items\n",
125 |     "        self.transforms = transforms\n",
126 |     "    \n",
127 |     "    def __len__(self):\n",
128 |     "        return self._items\n",
129 |     "    \n",
130 |     "    @staticmethod\n",
131 |     "    def map(image):\n",
132 |     "        image = image / 255\n",
133 |     "        return image\n",
134 |     "    \n",
135 |     "    @staticmethod\n",
136 |     "    def unmap(image):\n",
137 |     "        image = image * 255\n",
138 |     "        image = np.clip(image, 0, 255).astype('uint8')\n",
139 |     "        return image\n",
140 |     "    \n",
141 |     "    def __getitem__(self, item):\n",
142 |     "        img, _, labels, _ = cd.toydata.random_geometric_objects()\n",
143 |     "        \n",
144 |     "        if self.transforms is not None:\n",
145 |     "            r = self.transforms(image=img, mask=labels)\n",
146 |     "            img, labels = r['image'], r['mask']\n",
147 |     "        \n",
148 |     "        gen = self.gen\n",
149 |     "        gen.feed(labels=labels)\n",
150 |     "        \n",
151 |     "        image = self.map(img)\n",
152 |     "        return OrderedDict({\n",
153 |     "            'inputs': image.astype('float32'),\n",
154 |     "            'labels': gen.reduced_labels,\n",
155 |     "            'fourier': (gen.fourier.astype('float32'),),\n",
156 |     "            'locations': (gen.locations.astype('float32'),),\n",
157 |     "            'sampled_contours': (gen.sampled_contours.astype('float32'),),\n",
158 |     "            'sampling': (gen.sampling.astype('float32'),),\n",
159 |     "        })"
160 |    ]
161 |   },
162 |   {
163 |    "cell_type": "code",
164 |    "execution_count": null,
165 |    "metadata": {},
166 |    "outputs": [],
167 |    "source": [
168 |     "transforms = cd.conf2augmentation(conf.augmentation)\n",
169 |     "train_data = Dataset(conf.samples, conf.order, *conf.bg_fg_dists, transforms=transforms, items=conf.steps_per_epoch)\n",
170 |     "test_data = Dataset(conf.samples, conf.order, *conf.bg_fg_dists, items=2)\n",
171 |     "train_loader = DataLoader(train_data, batch_size=conf.batch_size, num_workers=conf.num_workers,\n",
172 |     "                          collate_fn=cd.universal_dict_collate_fn, shuffle=conf.shuffle)\n",
173 |     "test_loader = DataLoader(test_data, batch_size=2, num_workers=0, collate_fn=cd.universal_dict_collate_fn)\n",
174 |     "\n",
175 |     "# Plot example\n",
176 |     "example = train_data[0]\n",
177 |     "contours = example['sampled_contours'][0]\n",
178 |     "image = Dataset.unmap(example['inputs'])\n",
179 |     "cd.vis.show_detection(image, contours=contours, contour_line_width=5, figsize=(11, 11))"
180 |    ]
181 |   },
182 |   {
183 |    "cell_type": "markdown",
184 |    "metadata": {},
185 |    "source": [
186 |     "# 3. CPN Model"
187 |    ]
188 |   },
189 |   {
190 |    "cell_type": "code",
191 |    "execution_count": null,
192 |    "metadata": {},
193 |    "outputs": [],
194 |    "source": [
195 |     "model = getattr(models, conf.cpn)(in_channels=conf.in_channels, order=conf.order, samples=conf.samples,\n",
196 |     "                                  refinement_iterations=conf.refinement_iterations, nms_thresh=conf.nms_thresh,\n",
197 |     "                                  score_thresh=conf.score_thresh, contour_head_stride=conf.contour_head_stride,\n",
198 |     "                                  classes=conf.classes, refinement_buckets=conf.refinement_buckets,\n",
199 |     "                                  backbone_kwargs=dict(inputs_mean=conf.inputs_mean, inputs_std=conf.inputs_std))\n",
200 |     "cd.conf2tweaks_(conf.tweaks, model)\n",
201 |     "model.to(conf.device)\n",
202 |     "optimizer = cd.conf2optimizer(conf.optimizer, model.parameters())\n",
203 |     "scheduler = cd.conf2scheduler(conf.scheduler, optimizer)\n",
204 |     "scaler = GradScaler() if conf.amp else None  # set None to disable"
205 |    ]
206 |   },
207 |   {
208 |    "cell_type": "markdown",
209 |    "metadata": {},
210 |    "source": [
211 |     "# 4. Training\n",
212 |     "## 4.1 Training Functions"
213 |    ]
214 |   },
215 |   {
216 |    "cell_type": "code",
217 |    "execution_count": null,
218 |    "metadata": {},
219 |    "outputs": [],
220 |    "source": [
221 |     "def train_epoch(model, train_loder, device, optimizer, epoch, scaler=None, scheduler=None):\n",
222 |     "    model.train()\n",
223 |     "    for batch_idx, batch in enumerate(tqdm(train_loader, desc=\"Epoch %d\" % epoch)):\n",
224 |     "        batch = cd.to_device(batch, device)\n",
225 |     "        optimizer.zero_grad()\n",
226 |     "        with autocast(scaler is not None):\n",
227 |     "            outputs = model(batch['inputs'], targets=batch)\n",
228 |     "        loss = outputs['loss']\n",
229 |     "        if scaler is None:\n",
230 |     "            loss.backward()\n",
231 |     "            optimizer.step()\n",
232 |     "        else:\n",
233 |     "            scaler.scale(loss).backward()\n",
234 |     "            scaler.step(optimizer)\n",
235 |     "            scaler.update()\n",
236 |     "    if scheduler is not None:\n",
237 |     "        scheduler.step()\n",
238 |     "\n",
239 |     "def show_results(model, test_loader, device):\n",
240 |     "    model.eval()\n",
241 |     "    batch = cd.to_device(next(iter(test_loader)), device)\n",
242 |     "    with torch.no_grad():\n",
243 |     "        outputs = model(batch['inputs'])\n",
244 |     "    o = cd.asnumpy(outputs)\n",
245 |     "    num = len(o['contours'])\n",
246 |     "    plt.figure(None, (13 * num, 13))\n",
247 |     "    for idx in range(num):\n",
248 |     "        image = cd.asnumpy(batch['inputs'][idx])\n",
249 |     "        plt.subplot(1, num, idx + 1)\n",
250 |     "        cd.vis.show_detection(Dataset.unmap(image.transpose(1, 2, 0)), contours=o['contours'][idx],\n",
251 |     "                              contour_line_width=5, scores=o['scores'][idx])\n",
252 |     "    plt.show()"
253 |    ]
254 |   },
255 |   {
256 |    "cell_type": "markdown",
257 |    "metadata": {},
258 |    "source": [
259 |     "## 4.2 Training"
260 |    ]
261 |   },
262 |   {
263 |    "cell_type": "code",
264 |    "execution_count": null,
265 |    "metadata": {},
266 |    "outputs": [],
267 |    "source": [
268 |     "for epoch in range(1, conf.epochs):\n",
269 |     "    train_epoch(model, train_loader, conf.device, optimizer, epoch, scaler, scheduler)\n",
270 |     "    if epoch % 1 == 0:\n",
271 |     "        show_results(model, test_loader, conf.device)"
272 |    ]
273 |   }
274 |  ],
275 |  "metadata": {
276 |   "kernelspec": {
277 |    "display_name": "Python 3 (ipykernel)",
278 |    "language": "python",
279 |    "name": "python3"
280 |   },
281 |   "language_info": {
282 |    "codemirror_mode": {
283 |     "name": "ipython",
284 |     "version": 3
285 |    },
286 |    "file_extension": ".py",
287 |    "mimetype": "text/x-python",
288 |    "name": "python",
289 |    "nbconvert_exporter": "python",
290 |    "pygments_lexer": "ipython3",
291 |    "version": "3.9.16"
292 |   }
293 |  },
294 |  "nbformat": 4,
295 |  "nbformat_minor": 4
296 | }
297 | 


--------------------------------------------------------------------------------
/demos/demo-multiclass.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "metadata": {},
  6 |    "source": [
  7 |     "# Contour Proposal Networks － How to detect objects"
  8 |    ]
  9 |   },
 10 |   {
 11 |    "cell_type": "code",
 12 |    "execution_count": null,
 13 |    "metadata": {},
 14 |    "outputs": [],
 15 |    "source": [
 16 |     "import torch\n",
 17 |     "if torch.cuda.is_available():\n",
 18 |     "    torch.backends.cudnn.benchmark = True\n",
 19 |     "import celldetection as cd\n",
 20 |     "from celldetection import models, toydata\n",
 21 |     "import numpy as np\n",
 22 |     "import os\n",
 23 |     "from collections import OrderedDict\n",
 24 |     "import torch.nn as nn\n",
 25 |     "import torch.nn.functional as F\n",
 26 |     "from torch.utils.data import Dataset, DataLoader\n",
 27 |     "from tqdm import tqdm\n",
 28 |     "from matplotlib import pyplot as plt\n",
 29 |     "from torch.cuda.amp import GradScaler, autocast"
 30 |    ]
 31 |   },
 32 |   {
 33 |    "cell_type": "markdown",
 34 |    "metadata": {},
 35 |    "source": [
 36 |     "# 1. Config\n",
 37 |     "You can save a config `conf.to_json(filename)` and load it with `cd.Config.from_json(filename)`."
 38 |    ]
 39 |   },
 40 |   {
 41 |    "cell_type": "code",
 42 |    "execution_count": null,
 43 |    "metadata": {},
 44 |    "outputs": [],
 45 |    "source": [
 46 |     "conf = cd.Config(\n",
 47 |     "    # data\n",
 48 |     "    in_channels=3,\n",
 49 |     "    classes=4,\n",
 50 |     "    shuffle=True,\n",
 51 |     "    bg_fg_dists=(0.8, 0.85),\n",
 52 |     "    \n",
 53 |     "    # augmentation (schema: <class_name>:<kwargs>)\n",
 54 |     "    augmentation=OrderedDict({\n",
 55 |     "        'Transpose': {'p': 0.5},  # see: https://albumentations.ai/docs/\n",
 56 |     "        'RandomRotate90': {'p': 0.5},\n",
 57 |     "    }),\n",
 58 |     "    \n",
 59 |     "    # cpn\n",
 60 |     "    cpn='CpnU22',  # see https://git.io/JOnWX for alternatives\n",
 61 |     "    nms_thresh=.5,\n",
 62 |     "    contour_head_stride=8,\n",
 63 |     "    order=7,  # the higher, the more complex shapes can be detected\n",
 64 |     "    samples=128,  # number of coordinates per contour\n",
 65 |     "    refinement_iterations=3,\n",
 66 |     "    refinement_buckets=6,\n",
 67 |     "    inputs_mean=.5,\n",
 68 |     "    inputs_std=.5,\n",
 69 |     "    tweaks={\n",
 70 |     "        'BatchNorm2d': {'momentum': 0.05}\n",
 71 |     "    },\n",
 72 |     "    \n",
 73 |     "    # optimizer\n",
 74 |     "    optimizer={'Adadelta': {'lr': 1., 'rho': 0.9}},\n",
 75 |     "    scheduler={'StepLR': {'step_size': 5, 'gamma': .99}},\n",
 76 |     "    \n",
 77 |     "    # training\n",
 78 |     "    epochs=100,\n",
 79 |     "    steps_per_epoch=8 * 512,\n",
 80 |     "    batch_size=8,\n",
 81 |     "    amp=torch.cuda.is_available(),\n",
 82 |     "    \n",
 83 |     "    # misc\n",
 84 |     "    num_workers=0,\n",
 85 |     "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
 86 |     ")\n",
 87 |     "print(conf)"
 88 |    ]
 89 |   },
 90 |   {
 91 |    "cell_type": "markdown",
 92 |    "metadata": {},
 93 |    "source": [
 94 |     "# 2. Toy Dataset\n",
 95 |     "Each item of the dataset consists of an image and a label image.\n",
 96 |     "Just put your label image into the `CPNTargetGenerator` via `gen.feed(labels=labels)` and it will generate training targets.\n",
 97 |     "\n",
 98 |     "If you do not have a label image, you might have one of those:\n",
 99 |     "- A list of masks, each mask shows a single object. Then you can use:\n",
100 |     "```\n",
101 |     "labels = cd.data.unary_masks2labels([mask1, mask2, ..., maskN])\n",
102 |     "```\n",
103 |     "- A list of masks, each mask may show multiple objects, but touching objects were assigned different numbers. Then you can use:\n",
104 |     "```\n",
105 |     "labels = cd.data.masks2labels([mask1, mask2, ..., maskN])\n",
106 |     "```\n"
107 |    ]
108 |   },
109 |   {
110 |    "cell_type": "code",
111 |    "execution_count": null,
112 |    "metadata": {},
113 |    "outputs": [],
114 |    "source": [
115 |     "class Dataset:\n",
116 |     "    def __init__(self, samples, order, max_bg_dist, min_fg_dist, transforms=None, items=2**12):\n",
117 |     "        self.gen = cd.data.CPNTargetGenerator(\n",
118 |     "            samples=samples,\n",
119 |     "            order=order,\n",
120 |     "            min_fg_dist=min_fg_dist,\n",
121 |     "            max_bg_dist=max_bg_dist,\n",
122 |     "        )\n",
123 |     "        self._items = items\n",
124 |     "        self.transforms = transforms\n",
125 |     "    \n",
126 |     "    def __len__(self):\n",
127 |     "        return self._items\n",
128 |     "    \n",
129 |     "    @staticmethod\n",
130 |     "    def map(image):\n",
131 |     "        image = image / 255\n",
132 |     "        return image\n",
133 |     "    \n",
134 |     "    @staticmethod\n",
135 |     "    def unmap(image):\n",
136 |     "        image = image * 255\n",
137 |     "        image = np.clip(image, 0, 255).astype('uint8')\n",
138 |     "        return image\n",
139 |     "    \n",
140 |     "    def __getitem__(self, item):\n",
141 |     "        img, _, labels, classes = cd.toydata.random_geometric_objects()\n",
142 |     "        \n",
143 |     "        if self.transforms is not None:\n",
144 |     "            r = self.transforms(image=img, mask=labels)\n",
145 |     "            img, labels = r['image'], r['mask']\n",
146 |     "        \n",
147 |     "        gen = self.gen\n",
148 |     "        gen.feed(labels=labels)\n",
149 |     "        \n",
150 |     "        image = self.map(img)\n",
151 |     "        return OrderedDict({\n",
152 |     "            'inputs': image.astype('float32'),\n",
153 |     "            'labels': gen.reduced_labels,\n",
154 |     "            'fourier': (gen.fourier.astype('float32'),),\n",
155 |     "            'locations': (gen.locations.astype('float32'),),\n",
156 |     "            'sampled_contours': (gen.sampled_contours.astype('float32'),),\n",
157 |     "            'sampling': (gen.sampling.astype('float32'),),\n",
158 |     "            'classes': (classes,)  # just comment out this line if you only have two classes\n",
159 |     "        })"
160 |    ]
161 |   },
162 |   {
163 |    "cell_type": "code",
164 |    "execution_count": null,
165 |    "metadata": {},
166 |    "outputs": [],
167 |    "source": [
168 |     "transforms = cd.conf2augmentation(conf.augmentation)\n",
169 |     "train_data = Dataset(conf.samples, conf.order, *conf.bg_fg_dists, transforms=transforms, items=conf.steps_per_epoch)\n",
170 |     "test_data = Dataset(conf.samples, conf.order, *conf.bg_fg_dists, items=2)\n",
171 |     "train_loader = DataLoader(train_data, batch_size=conf.batch_size, num_workers=conf.num_workers,\n",
172 |     "                          collate_fn=cd.universal_dict_collate_fn, shuffle=conf.shuffle)\n",
173 |     "test_loader = DataLoader(test_data, batch_size=2, num_workers=0, collate_fn=cd.universal_dict_collate_fn)\n",
174 |     "\n",
175 |     "# Plot example\n",
176 |     "example = train_data[0]\n",
177 |     "class_d = cd.toydata.CLASS_NAMES_GEOMETRIC\n",
178 |     "cd.vis.show_detection(image=Dataset.unmap(example['inputs']), contours=example['sampled_contours'][0],\n",
179 |     "                      class_name=[class_d[i] for i in example['classes'][0]], scores=np.ones(len(example['classes'][0])),\n",
180 |     "                      contour_line_width=5, figsize=(11, 11))"
181 |    ]
182 |   },
183 |   {
184 |    "cell_type": "markdown",
185 |    "metadata": {},
186 |    "source": [
187 |     "# 3. CPN Model"
188 |    ]
189 |   },
190 |   {
191 |    "cell_type": "code",
192 |    "execution_count": null,
193 |    "metadata": {},
194 |    "outputs": [],
195 |    "source": [
196 |     "model = getattr(models, conf.cpn)(in_channels=conf.in_channels, order=conf.order, samples=conf.samples,\n",
197 |     "                                  refinement_iterations=conf.refinement_iterations, nms_thresh=conf.nms_thresh,\n",
198 |     "                                  contour_head_stride=conf.contour_head_stride, classes=conf.classes,\n",
199 |     "                                  refinement_buckets=conf.refinement_buckets,\n",
200 |     "                                  backbone_kwargs=dict(inputs_mean=conf.inputs_mean, inputs_std=conf.inputs_std))\n",
201 |     "cd.conf2tweaks_(conf.tweaks, model)\n",
202 |     "model.to(conf.device)\n",
203 |     "optimizer = cd.conf2optimizer(conf.optimizer, model.parameters())\n",
204 |     "scheduler = cd.conf2scheduler(conf.scheduler, optimizer)\n",
205 |     "scaler = GradScaler() if conf.amp else None  # set None to disable"
206 |    ]
207 |   },
208 |   {
209 |    "cell_type": "markdown",
210 |    "metadata": {},
211 |    "source": [
212 |     "# 4. Training\n",
213 |     "## 4.1 Training Functions"
214 |    ]
215 |   },
216 |   {
217 |    "cell_type": "code",
218 |    "execution_count": null,
219 |    "metadata": {},
220 |    "outputs": [],
221 |    "source": [
222 |     "def train_epoch(model, train_loder, device, optimizer, epoch, scaler=None, scheduler=None):\n",
223 |     "    model.train()\n",
224 |     "    for batch_idx, batch in enumerate(tqdm(train_loader, desc=\"Epoch %d\" % epoch)):\n",
225 |     "        batch = cd.to_device(batch, device)\n",
226 |     "        optimizer.zero_grad()\n",
227 |     "        with autocast(scaler is not None):\n",
228 |     "            outputs = model(batch['inputs'], targets=batch)\n",
229 |     "        loss = outputs['loss']\n",
230 |     "        if scaler is None:\n",
231 |     "            loss.backward()\n",
232 |     "            optimizer.step()\n",
233 |     "        else:\n",
234 |     "            scaler.scale(loss).backward()\n",
235 |     "            scaler.step(optimizer)\n",
236 |     "            scaler.update()\n",
237 |     "    if scheduler is not None:\n",
238 |     "        scheduler.step()\n",
239 |     "        \n",
240 |     "def show_results(model, test_loader, device):\n",
241 |     "    model.eval()\n",
242 |     "    batch = cd.to_device(next(iter(test_loader)), device)\n",
243 |     "    with torch.no_grad():\n",
244 |     "        outputs = model(batch['inputs'])\n",
245 |     "    o = cd.asnumpy(outputs)\n",
246 |     "    num = len(o['contours'])\n",
247 |     "    plt.figure(None, (13 * num, 13))\n",
248 |     "    for idx in range(num):\n",
249 |     "        image = cd.asnumpy(batch['inputs'][idx])\n",
250 |     "        plt.subplot(1, num, idx + 1)\n",
251 |     "        class_names = [cd.toydata.CLASS_NAMES_GEOMETRIC[j] for j in o['classes'][idx]]\n",
252 |     "        cd.vis.show_detection(Dataset.unmap(image.transpose(1, 2, 0)), contours=o['contours'][idx],\n",
253 |     "                              contour_line_width=5, scores=o['scores'][idx], class_name=class_names)\n",
254 |     "    plt.show()"
255 |    ]
256 |   },
257 |   {
258 |    "cell_type": "markdown",
259 |    "metadata": {},
260 |    "source": [
261 |     "## 4.2 Training"
262 |    ]
263 |   },
264 |   {
265 |    "cell_type": "code",
266 |    "execution_count": null,
267 |    "metadata": {},
268 |    "outputs": [],
269 |    "source": [
270 |     "for epoch in range(1, conf.epochs):\n",
271 |     "    train_epoch(model, train_loader, conf.device, optimizer, epoch, scaler, scheduler)\n",
272 |     "    if epoch % 1 == 0:\n",
273 |     "        show_results(model, test_loader, conf.device)"
274 |    ]
275 |   }
276 |  ],
277 |  "metadata": {
278 |   "kernelspec": {
279 |    "display_name": "Python 3 (ipykernel)",
280 |    "language": "python",
281 |    "name": "python3"
282 |   },
283 |   "language_info": {
284 |    "codemirror_mode": {
285 |     "name": "ipython",
286 |     "version": 3
287 |    },
288 |    "file_extension": ".py",
289 |    "mimetype": "text/x-python",
290 |    "name": "python",
291 |    "nbconvert_exporter": "python",
292 |    "pygments_lexer": "ipython3",
293 |    "version": "3.9.16"
294 |   }
295 |  },
296 |  "nbformat": 4,
297 |  "nbformat_minor": 4
298 | }
299 | 


--------------------------------------------------------------------------------
/docker/Dockerfile:
--------------------------------------------------------------------------------
 1 | # syntax=docker/dockerfile:1
 2 | 
 3 | FROM python:3.10-buster
 4 | 
 5 | WORKDIR /app
 6 | 
 7 | RUN pip install --upgrade pip
 8 | RUN pip install safetensors==0.2.8
 9 | 
10 | RUN apt-get update && apt-get install libgl1 mpich -y
11 | RUN pip install mpi4py
12 | 
13 | COPY ./celldetection ./celldetection
14 | COPY ./setup.py ./
15 | COPY ./requirements.txt ./
16 | COPY ./MANIFEST.in ./
17 | COPY ./README.md ./
18 | COPY ./LICENSE ./
19 | COPY ./celldetection_scripts ./celldetection_scripts
20 | 
21 | RUN pip install imagecodecs; exit 0  # optional due to unresolved issues on ARM
22 | RUN pip install .
23 | RUN pip install psutil
24 | 
25 | RUN mkdir /inputs
26 | RUN mkdir /outputs
27 | RUN mkdir /models
28 | 


--------------------------------------------------------------------------------
/docker/DockerfileNvidia:
--------------------------------------------------------------------------------
 1 | # syntax=docker/dockerfile:1
 2 | 
 3 | FROM nvcr.io/nvidia/pytorch:23.10-py3
 4 | 
 5 | WORKDIR /app
 6 | 
 7 | RUN pip install --upgrade pip
 8 | RUN pip install safetensors==0.2.8
 9 | 
10 | RUN apt-get update && apt-get install libgl1 -y
11 | RUN pip install mpi4py
12 | 
13 | COPY ./celldetection ./celldetection
14 | COPY ./setup.py ./
15 | COPY ./requirements.txt ./
16 | COPY ./MANIFEST.in ./
17 | COPY ./README.md ./
18 | COPY ./LICENSE ./
19 | COPY ./celldetection_scripts ./celldetection_scripts
20 | 
21 | RUN pip install imagecodecs; exit 0  # optional due to unresolved issues on ARM
22 | RUN pip install .
23 | RUN pip install psutil
24 | 
25 | RUN mkdir /inputs
26 | RUN mkdir /outputs
27 | RUN mkdir /models
28 | 


--------------------------------------------------------------------------------
/docker/README.md:
--------------------------------------------------------------------------------
 1 | # Docker
 2 | 
 3 | ## Build
 4 | Build with context `./` (project root) and dockerfile `./docker/Dockerfile`.
 5 | ```
 6 | docker build --tag "celldetection:latest" -f docker/Dockerfile .
 7 | ```
 8 | 
 9 | For Mac Silicon:
10 | ```
11 | docker build --platform linux/arm64 --tag "celldetection:latest" -f docker/Dockerfile 
12 | ```
13 | 
14 | With Multi-Platform Support:
15 | ```
16 | docker buildx create --use
17 | docker buildx build --platform linux/amd64,linux/arm64 --tag "celldetection:latest" -f docker/Dockerfile 
18 | ```
19 | 
20 | ## Export
21 | ```
22 | docker save celldetection | gzip -1 -c > celldetection.tar.gz
23 | ```
24 | 
25 | ## Load
26 | ```
27 | docker load -i celldetection.tar.gz
28 | ```
29 | 
30 | 
31 | ## Run
32 | ```
33 | docker container run celldetection:latest /bin/bash -c "python -c 'print(2)'"
34 | ```
35 | 


--------------------------------------------------------------------------------
/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/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/mk.sh:
--------------------------------------------------------------------------------
1 | make clean && make html
2 | 


--------------------------------------------------------------------------------
/docs/requirements.txt:
--------------------------------------------------------------------------------
1 | sphinx
2 | sphinx-rtd-theme


--------------------------------------------------------------------------------
/docs/source/_static/img/magnify.svg:
--------------------------------------------------------------------------------
1 | <?xml version="1.0" encoding="UTF-8"?><!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
2 |         "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
3 | <svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" version="1.1" width="24" height="24"
4 |      viewBox="0 0 24 24">
5 |     <path fill="#0f1429" d="M9.5,3A6.5,6.5 0 0,1 16,9.5C16,11.11 15.41,12.59 14.44,13.73L14.71,14H15.5L20.5,19L19,20.5L14,15.5V14.71L13.73,14.44C12.59,15.41 11.11,16 9.5,16A6.5,6.5 0 0,1 3,9.5A6.5,6.5 0 0,1 9.5,3M9.5,5C7,5 5,7 5,9.5C5,12 7,14 9.5,14C12,14 14,12 14,9.5C14,7 12,5 9.5,5Z"/>
6 | </svg>


--------------------------------------------------------------------------------
/docs/source/_templates/layout.html:
--------------------------------------------------------------------------------
  1 | {# TEMPLATE VAR SETTINGS #}
  2 | {%- set url_root = pathto('', 1) %}
  3 | {%- if url_root == '#' %}{% set url_root = '' %}{% endif %}
  4 | {%- if not embedded and docstitle %}
  5 |   {%- set titlesuffix = " &mdash; "|safe + docstitle|e %}
  6 | {%- else %}
  7 |   {%- set titlesuffix = "" %}
  8 | {%- endif %}
  9 | {%- set lang_attr = 'en' if language == None else (language | replace('_', '-')) %}
 10 | {%- set sphinx_writer = 'writer-html5' if html5_doctype else 'writer-html4' -%}
 11 | 
 12 | {# Build sphinx_version_info tuple from sphinx_version string in pure Jinja #}
 13 | {%- set (_ver_major, _ver_minor, _ver_bugfix) = sphinx_version.split('.') | map('int') -%}
 14 | {%- set sphinx_version_info = (_ver_major, _ver_minor, _ver_bugfix) -%}
 15 | 
 16 | <!DOCTYPE html>
 17 | <html class="{{ sphinx_writer }}" lang="{{ lang_attr }}" >
 18 | <head>
 19 |   <meta charset="utf-8" />
 20 |   {{- metatags }}
 21 |   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
 22 |   {%- block htmltitle %}
 23 |   <title>{{ title|striptags|e }}{{ titlesuffix }}</title>
 24 |   {%- endblock -%}
 25 | 
 26 |   {#- CSS #}
 27 |   {%- if sphinx_version_info < (4, 0) -%}
 28 |     <link rel="stylesheet" href="{{ pathto('_static/' + style, 1) }}" type="text/css" />
 29 |     <link rel="stylesheet" href="{{ pathto('_static/pygments.css', 1) }}" type="text/css" />
 30 |   {%- endif %}
 31 |   {%- for css in css_files %}
 32 |     {%- if css|attr("rel") %}
 33 |       <link rel="{{ css.rel }}" href="{{ pathto(css.filename, 1) }}" type="text/css"{% if css.title is not none %} title="{{ css.title }}"{% endif %} />
 34 |     {%- else %}
 35 |       <link rel="stylesheet" href="{{ pathto(css, 1) }}" type="text/css" />
 36 |     {%- endif %}
 37 |   {%- endfor %}
 38 | 
 39 |   {%- for cssfile in extra_css_files %}
 40 |     <link rel="stylesheet" href="{{ pathto(cssfile, 1) }}" type="text/css" />
 41 |   {%- endfor -%}
 42 | 
 43 |   {#- FAVICON #}
 44 |   {%- if favicon %}
 45 |     {%- if sphinx_version_info < (4, 0) -%}
 46 |     <link rel="shortcut icon" href="{{ pathto('_static/' + favicon, 1) }}"/>
 47 |     {%- else %}
 48 |     <link rel="shortcut icon" href="{{ favicon_url }}"/>
 49 |     {%- endif %}
 50 |   {%- endif -%}
 51 | 
 52 |   {#- CANONICAL URL (deprecated) #}
 53 |   {%- if theme_canonical_url and not pageurl %}
 54 |     <link rel="canonical" href="{{ theme_canonical_url }}{{ pagename }}.html"/>
 55 |   {%- endif -%}
 56 | 
 57 |   {#- CANONICAL URL #}
 58 |   {%- if pageurl %}
 59 |     <link rel="canonical" href="{{ pageurl|e }}" />
 60 |   {%- endif -%}
 61 | 
 62 |   {#- JAVASCRIPTS #}
 63 |   {%- block scripts %}
 64 |   <!--[if lt IE 9]>
 65 |     <script src="{{ pathto('_static/js/html5shiv.min.js', 1) }}"></script>
 66 |   <![endif]-->
 67 |   {%- if not embedded %}
 68 |   {# XXX Sphinx 1.8.0 made this an external js-file, quick fix until we refactor the template to inherert more blocks directly from sphinx #}
 69 |     {%- if sphinx_version_info >= (1, 8) -%}
 70 |       {%- if sphinx_version_info < (4, 0) -%}
 71 |       <script id="documentation_options" data-url_root="{{ url_root }}" src="{{ pathto('_static/documentation_options.js', 1) }}"></script>
 72 |       {%- endif -%}
 73 |       {%- for scriptfile in script_files %}
 74 |         {{ js_tag(scriptfile) }}
 75 |       {%- endfor %}
 76 |     {%- else %}
 77 |       <script>
 78 |           var DOCUMENTATION_OPTIONS = {
 79 |               URL_ROOT:'{{ url_root }}',
 80 |               VERSION:'{{ release|e }}',
 81 |               LANGUAGE:'{{ language }}',
 82 |               COLLAPSE_INDEX:false,
 83 |               FILE_SUFFIX:'{{ '' if no_search_suffix else file_suffix }}',
 84 |               HAS_SOURCE:  {{ has_source|lower }},
 85 |               SOURCELINK_SUFFIX: '{{ sourcelink_suffix }}'
 86 |           };
 87 |       </script>
 88 |       {%- for scriptfile in script_files %}
 89 |         <script src="{{ pathto(scriptfile, 1) }}"></script>
 90 |       {%- endfor %}
 91 |     {%- endif %}
 92 |     <script src="{{ pathto('_static/js/theme.js', 1) }}"></script>
 93 | 
 94 |     {#- OPENSEARCH #}
 95 |     {%- if use_opensearch %}
 96 |     <link rel="search" type="application/opensearchdescription+xml"
 97 |           title="{% trans docstitle=docstitle|e %}Search within {{ docstitle }}{% endtrans %}"
 98 |           href="{{ pathto('_static/opensearch.xml', 1) }}"/>
 99 |     {%- endif %}
100 |   {%- endif %}
101 |   {%- endblock %}
102 | 
103 |   {%- block linktags %}
104 |     {%- if hasdoc('about') %}
105 |     <link rel="author" title="{{ _('About these documents') }}" href="{{ pathto('about') }}" />
106 |     {%- endif %}
107 |     {%- if hasdoc('genindex') %}
108 |     <link rel="index" title="{{ _('Index') }}" href="{{ pathto('genindex') }}" />
109 |     {%- endif %}
110 |     {%- if hasdoc('search') %}
111 |     <link rel="search" title="{{ _('Search') }}" href="{{ pathto('search') }}" />
112 |     {%- endif %}
113 |     {%- if hasdoc('copyright') %}
114 |     <link rel="copyright" title="{{ _('Copyright') }}" href="{{ pathto('copyright') }}" />
115 |     {%- endif %}
116 |     {%- if next %}
117 |     <link rel="next" title="{{ next.title|striptags|e }}" href="{{ next.link|e }}" />
118 |     {%- endif %}
119 |     {%- if prev %}
120 |     <link rel="prev" title="{{ prev.title|striptags|e }}" href="{{ prev.link|e }}" />
121 |     {%- endif %}
122 |   {%- endblock %}
123 |   {%- block extrahead %} {% endblock %}
124 | 
125 |     <style>
126 |         .class>dt .sig-name>span::before,
127 |         .function>dt .sig-name>span::before {
128 |             content: "{{ title }}.";
129 |             color: var(--module-prefix-color);
130 |         }
131 |     </style>
132 | 
133 | </head>
134 | 
135 | <body class="wy-body-for-nav">
136 | 
137 |   {%- block extrabody %} {% endblock %}
138 |   <div class="wy-grid-for-nav">
139 |     {#- SIDE NAV, TOGGLES ON MOBILE #}
140 |     <nav data-toggle="wy-nav-shift" class="wy-nav-side">
141 |       <div class="wy-side-scroll">
142 |         <div class="wy-side-nav-search" {% if theme_style_nav_header_background %} style="background: {{theme_style_nav_header_background}}" {% endif %}>
143 |           {%- block sidebartitle %}
144 | 
145 |           {%- if logo and theme_logo_only %}
146 |             <a href="{{ pathto(master_doc) }}">
147 |           {%- else %}
148 |             <a href="{{ pathto(master_doc) }}" class="icon icon-home"> {{ project }}
149 |           {%- endif %}
150 | 
151 |           {%- if logo %}
152 |             {#- Not strictly valid HTML, but it's the only way to display/scale
153 |                 it properly, without weird scripting or heaps of work
154 |             #}
155 |             {%- if sphinx_version_info < (4, 0) -%}
156 |             <img src="{{ pathto('_static/' + logo, 1) }}" class="logo" alt="{{ _('Logo') }}"/>
157 |             {%- else %}
158 |             <img src="{{ logo_url }}" class="logo" alt="{{ _('Logo') }}"/>
159 |             {%- endif %}
160 |           {%- endif %}
161 |           </a>
162 | 
163 | <!--          {%- if theme_display_version %}-->
164 | <!--            {%- set nav_version = version %}-->
165 | <!--            {%- if READTHEDOCS and current_version %}-->
166 | <!--              {%- set nav_version = current_version %}-->
167 | <!--            {%- endif %}-->
168 | <!--            {%- if nav_version %}-->
169 | <!--              <div class="version">-->
170 | <!--                {{ nav_version }}-->
171 | <!--              </div>-->
172 | <!--            {%- endif %}-->
173 | <!--          {%- endif %}-->
174 | 
175 |           {%- include "searchbox.html" %}
176 | 
177 |           {%- endblock %}
178 |         </div>
179 | 
180 |         {%- block navigation %}
181 |         {#- Translators: This is an ARIA section label for the main navigation menu -#}
182 |         <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="{{ _('Navigation menu') }}">
183 | <!--          <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>-->
184 |           {%- block menu %}
185 |             {%- set toctree = toctree(maxdepth=theme_navigation_depth|int,
186 |                                       collapse=theme_collapse_navigation|tobool,
187 |                                       includehidden=theme_includehidden|tobool,
188 |                                       titles_only=theme_titles_only|tobool) %}
189 |             {%- if toctree %}
190 |               {{ toctree }}
191 |             {%- else %}
192 |               <!-- Local TOC -->
193 |               <div class="local-toc">{{ toc }}</div>
194 |             {%- endif %}
195 |           {%- endblock %}
196 |         </div>
197 |         {%- endblock %}
198 |       </div>
199 |     </nav>
200 | 
201 |     <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
202 | 
203 |       {#- MOBILE NAV, TRIGGLES SIDE NAV ON TOGGLE #}
204 |       {#- Translators: This is an ARIA section label for the navigation menu that is visible when viewing the page on mobile devices -#}
205 |       <nav class="wy-nav-top" aria-label="{{ _('Mobile navigation menu') }}" {% if theme_style_nav_header_background %} style="background: {{theme_style_nav_header_background}}" {% endif %}>
206 |         {%- block mobile_nav %}
207 |           <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
208 |           <a href="{{ pathto(master_doc) }}">{{ project }}</a>
209 |         {%- endblock %}
210 |       </nav>
211 | 
212 |       <div class="wy-nav-content">
213 |       {%- block content %}
214 |         {%- if theme_style_external_links|tobool %}
215 |         <div class="rst-content style-external-links">
216 |         {%- else %}
217 |         <div class="rst-content">
218 |         {%- endif %}
219 |           {% include "breadcrumbs.html" %}
220 |           <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
221 |           {%- block document %}
222 |            <div itemprop="articleBody">
223 |              {% block body %}{% endblock %}
224 |            </div>
225 |            {%- if self.comments()|trim %}
226 |              <div class="articleComments">
227 |                {%- block comments %}{% endblock %}
228 |              </div>
229 |            {%- endif%}
230 |           </div>
231 |           {%- endblock %}
232 |           {% include "footer.html" %}
233 |         </div>
234 |       {%- endblock %}
235 |       </div>
236 |     </section>
237 |   </div>
238 |   {% include "versions.html" -%}
239 | 
240 |   <script>
241 |       jQuery(function () {
242 |           SphinxRtdTheme.Navigation.enable({{ 'true' if theme_sticky_navigation|tobool else 'false' }});
243 |       });
244 |   </script>
245 | 
246 |   {#- Do not conflict with RTD insertion of analytics script #}
247 |   {%- if not READTHEDOCS %}
248 |     {%- if theme_analytics_id %}
249 |     <!-- Theme Analytics -->
250 |     <script async src="https://www.googletagmanager.com/gtag/js?id={{ theme_analytics_id }}"></script>
251 |     <script>
252 |       window.dataLayer = window.dataLayer || [];
253 |       function gtag(){dataLayer.push(arguments);}
254 |       gtag('js', new Date());
255 | 
256 |       gtag('config', '{{ theme_analytics_id }}', {
257 |           'anonymize_ip': {{ 'true' if theme_analytics_anonymize_ip|tobool else 'false' }},
258 |       });
259 |     </script>
260 | 
261 |     {%- endif %}
262 |   {%- endif %}
263 | 
264 |   {%- block footer %} {% endblock %}
265 | 
266 | </body>
267 | </html>
268 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.callbacks.rst:
--------------------------------------------------------------------------------
1 | cd.callbacks
2 | ============
3 | 
4 | .. automodule:: celldetection.callbacks.dropout
5 |    :members:
6 |    :undoc-members:
7 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.data.rst:
--------------------------------------------------------------------------------
 1 | cd.data
 2 | ==================
 3 | 
 4 | .. automodule:: celldetection.data
 5 |    :members:
 6 |    :undoc-members:
 7 | 
 8 | Datasets
 9 | --------
10 | 
11 | .. automodule:: celldetection.data.datasets
12 |    :members:
13 |    :undoc-members:
14 | 
15 | 
16 | BBBC039
17 | *******
18 | 
19 | .. automodule:: celldetection.data.datasets.bbbc039
20 |    :members:
21 |    :undoc-members:
22 | 
23 | BBBC041
24 | *******
25 | 
26 | .. automodule:: celldetection.data.datasets.bbbc041
27 |    :members:
28 |    :undoc-members:
29 | 
30 | Synth
31 | *****
32 | 
33 | .. automodule:: celldetection.data.datasets.synth
34 |    :members:
35 |    :undoc-members:
36 | 
37 | Data Operations
38 | ---------------
39 | 
40 | .. automodule:: celldetection.data.misc
41 |    :members:
42 |    :undoc-members:
43 | 
44 | .. automodule:: celldetection.data.segmentation
45 |    :members:
46 |    :undoc-members:
47 | 
48 | CPN Operations
49 | --------------
50 | 
51 | .. automodule:: celldetection.data.cpn
52 |    :members:
53 |    :undoc-members:
54 | 
55 | Eval Operations
56 | ---------------
57 | 
58 | .. automodule:: celldetection.data.instance_eval
59 |    :members:
60 |    :undoc-members:
61 | 
62 | Toy Data
63 | --------
64 | 
65 | .. automodule:: celldetection.data.toydata
66 |    :members:
67 |    :undoc-members:
68 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.models.rst:
--------------------------------------------------------------------------------
 1 | cd.models
 2 | ====================
 3 | 
 4 | .. automodule:: celldetection.models
 5 |    :members:
 6 |    :undoc-members:
 7 | 
 8 | U-Net
 9 | -----
10 | 
11 | .. automodule:: celldetection.models.unet
12 |    :members:
13 |    :undoc-members:
14 | 
15 | Multi-Scale Attention Network
16 | -----------------------------
17 | 
18 | .. automodule:: celldetection.models.manet
19 |    :members:
20 |    :undoc-members:
21 | 
22 | Contour Proposal Network
23 | ------------------------
24 | 
25 | .. automodule:: celldetection.models.cpn
26 |    :members:
27 |    :undoc-members:
28 | 
29 | ResNet
30 | ------
31 | 
32 | .. automodule:: celldetection.models.resnet
33 |    :members:
34 |    :undoc-members:
35 | 
36 | Feature Pyramid Network
37 | -----------------------
38 | 
39 | .. automodule:: celldetection.models.fpn
40 |    :members:
41 |    :undoc-members:
42 | 
43 | Timm Encoder
44 | ------------
45 | 
46 | .. automodule:: celldetection.models.timmodels
47 |    :members:
48 |    :undoc-members:
49 | 
50 | Smp Encoder
51 | -----------
52 | 
53 | .. automodule:: celldetection.models.smp
54 |    :members:
55 |    :undoc-members:
56 | 
57 | ConvNeXt
58 | --------
59 | 
60 | .. automodule:: celldetection.models.convnext
61 |    :members:
62 |    :undoc-members:
63 | 
64 | MobileNet
65 | ---------
66 | 
67 | .. automodule:: celldetection.models.mobilenetv3
68 |    :members:
69 |    :undoc-members:
70 | 
71 | Common Modules
72 | --------------
73 | 
74 | .. automodule:: celldetection.models.commons
75 |    :members:
76 |    :undoc-members:
77 | 
78 | Filters
79 | -------
80 | 
81 | .. automodule:: celldetection.models.filters
82 |    :members:
83 |    :undoc-members:
84 | 
85 | Normalization
86 | -------------
87 | 
88 | .. automodule:: celldetection.models.normalization
89 |    :members:
90 |    :undoc-members:
91 | 
92 | Losses
93 | ------
94 | 
95 | .. automodule:: celldetection.models.loss
96 |    :members:
97 |    :undoc-members:
98 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.mpi.rst:
--------------------------------------------------------------------------------
 1 | cd.mpi
 2 | =================
 3 | 
 4 | 
 5 | .. automodule:: celldetection.mpi
 6 |    :members:
 7 |    :undoc-members:
 8 | 
 9 | .. automodule:: celldetection.mpi.mpi
10 |    :members:
11 |    :undoc-members:
12 | 
13 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.ops.rst:
--------------------------------------------------------------------------------
 1 | cd.ops
 2 | ======
 3 | 
 4 | .. automodule:: celldetection.ops
 5 |    :members:
 6 |    :undoc-members:
 7 | 
 8 | Drawing Operations
 9 | ------------------
10 | 
11 | .. automodule:: celldetection.ops.draw
12 |    :members:
13 |    :undoc-members:
14 | 
15 | Box Operations
16 | --------------
17 | 
18 | .. automodule:: celldetection.ops.boxes
19 |    :members:
20 |    :undoc-members:
21 | 
22 | Losses
23 | ------
24 | 
25 | .. automodule:: celldetection.ops.loss
26 |    :members:
27 |    :undoc-members:
28 | 
29 | Normalization
30 | -------------
31 | 
32 | .. automodule:: celldetection.ops.normalization
33 |    :members:
34 |    :undoc-members:
35 | 
36 | Common Operations
37 | -----------------
38 | 
39 | .. automodule:: celldetection.ops.commons
40 |    :members:
41 |    :undoc-members:
42 | 
43 | CPN Operations
44 | -----------------
45 | 
46 | .. automodule:: celldetection.ops.cpn
47 |    :members:
48 |    :undoc-members:
49 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.optim.rst:
--------------------------------------------------------------------------------
1 | cd.optim
2 | ========
3 | 
4 | .. automodule:: celldetection.optim.lr_scheduler
5 |    :members:
6 |    :undoc-members:
7 | 
8 | 


--------------------------------------------------------------------------------
/docs/source/celldetection.rst:
--------------------------------------------------------------------------------
 1 | cd
 2 | ==
 3 | 
 4 | .. automodule:: celldetection
 5 |    :members:
 6 |    :undoc-members:
 7 | 
 8 | Visualization
 9 | -------------
10 | 
11 | .. automodule:: celldetection.visualization.cmaps
12 |    :members:
13 |    :undoc-members:
14 | 
15 | .. automodule:: celldetection.visualization.images
16 |    :members:
17 |    :undoc-members:
18 | 
19 | Config
20 | ------
21 |    
22 | .. automodule:: celldetection.util.schedule
23 |    :members:
24 |    :undoc-members:
25 | 
26 | Timing
27 | ------
28 | 
29 | .. automodule:: celldetection.util.timer
30 |    :members:
31 |    :undoc-members:
32 | 
33 | Util
34 | ------
35 | 
36 | .. automodule:: celldetection.util.util
37 |    :members:
38 |    :undoc-members:
39 | 


--------------------------------------------------------------------------------
/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 | # https://www.sphinx-doc.org/en/master/usage/configuration.html
 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 | import os
14 | import sys
15 | sys.path.insert(0, os.path.abspath('../..'))
16 | 
17 | 
18 | # -- Project information -----------------------------------------------------
19 | 
20 | project = 'CellDetection'
21 | copyright = '2022, Eric Upschulte'
22 | author = 'Eric Upschulte'
23 | 
24 | # The full version, including alpha/beta/rc tags
25 | release = 'release'
26 | 
27 | 
28 | # -- General configuration ---------------------------------------------------
29 | 
30 | # Add any Sphinx extension module names here, as strings. They can be
31 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
32 | # ones.
33 | extensions = ['sphinx.ext.autodoc', 'sphinx.ext.autosummary', 'sphinx.ext.napoleon', 'sphinx.ext.imgmath']
34 | 
35 | # Add any paths that contain templates here, relative to this directory.
36 | templates_path = ['_templates']
37 | 
38 | # List of patterns, relative to source directory, that match files and
39 | # directories to ignore when looking for source files.
40 | # This pattern also affects html_static_path and html_extra_path.
41 | exclude_patterns = []
42 | 
43 | 
44 | # -- Options for HTML output -------------------------------------------------
45 | 
46 | # The theme to use for HTML and HTML Help pages.  See the documentation for
47 | # a list of builtin themes.
48 | #
49 | html_theme = 'sphinx_rtd_theme'
50 | 
51 | # Add any paths that contain custom static files (such as style sheets) here,
52 | # relative to this directory. They are copied after the builtin static files,
53 | # so a file named "default.css" will overwrite the builtin "default.css".
54 | html_static_path = ['_static']
55 | 
56 | 
57 | autoclass_content = 'both'
58 | autoapi_python_use_implicit_namespaces = False
59 | add_module_names = False
60 | autodoc_inherit_docstrings = False
61 | theme_navigation_depth = 2
62 | html_theme_options = {'navigation_depth': 2}
63 | 


--------------------------------------------------------------------------------
/docs/source/index.rst:
--------------------------------------------------------------------------------
 1 | Cell Detection
 2 | ==============
 3 | 
 4 | A Python package for Cell Detection with PyTorch.
 5 | 
 6 | Showcase
 7 | --------
 8 | 
 9 | Nuclei of U2OS cells in a chemical screen
10 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
11 | 
12 | .. image:: ../../assets/bbbc039-cpn-u22-demo.png
13 |   :alt: bbbc039
14 | 
15 | https://bbbc.broadinstitute.org/BBBC039 (CC0)
16 | 
17 | P. vivax (malaria) infected human blood
18 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
19 | 
20 | .. image:: ../../assets/bbbc041-cpn-u22-demo.png
21 |   :alt: bbbc041
22 | 
23 | https://bbbc.broadinstitute.org/BBBC041 (CC BY-NC-SA 3.0)
24 | 
25 | Content
26 | -------
27 | 
28 | .. toctree::
29 |    :maxdepth: 1
30 | 
31 |    install
32 |    inference
33 | 
34 | API Reference
35 | -------------
36 | 
37 | .. toctree::
38 |    :maxdepth: 1
39 | 
40 |    celldetection
41 |    celldetection.data
42 |    celldetection.models
43 |    celldetection.mpi
44 |    celldetection.ops
45 |    celldetection.optim
46 |    celldetection.callbacks
47 | 


--------------------------------------------------------------------------------
/docs/source/install.rst:
--------------------------------------------------------------------------------
  1 | Installation Guide
  2 | ==================
  3 | 
  4 | This guide provides detailed instructions for installing the `celldetection` package. Depending on your environment and requirements, you can choose from several installation methods.
  5 | 
  6 | Docker and Apptainer Installation
  7 | ---------------------------------
  8 | 
  9 | For users who prefer using Docker or Apptainer, installation of PyTorch or setting up virtual environments is not required, as the Docker image comes with all necessary dependencies.
 10 | 
 11 | **Docker Installation**
 12 | 
 13 | Pull the latest `celldetection` Docker image:
 14 | 
 15 | .. code-block:: bash
 16 | 
 17 |     docker pull ericup/celldetection:latest
 18 | 
 19 | **Verification for Docker**
 20 | 
 21 | After pulling the image, you can run a Docker container to verify the installation:
 22 | 
 23 | .. code-block:: bash
 24 | 
 25 |     docker run --rm --gpus="device=0" ericup/celldetection:latest python -c "import celldetection; print(celldetection.__version__)"
 26 | 
 27 | You may remove ``--gpus="device=0"`` if you do not have GPUs on your system.
 28 | 
 29 | **Apptainer Installation**
 30 | 
 31 | In HPC environments where Apptainer is preferred:
 32 | 
 33 | .. code-block:: bash
 34 | 
 35 |     apptainer pull --dir . --disable-cache docker://ericup/celldetection:latest
 36 | 
 37 | If your system allows caching, you may remove ``--disable-cache``.
 38 | On some systems you may need to `specify a custom cache directory <https://apptainer.org/docs/user/latest/build_env.html>`_ with sufficient disk space.
 39 | 
 40 | **Verification for Apptainer**
 41 | 
 42 | To verify the installation in Apptainer, run the following command using the downloaded `.sif` file:
 43 | 
 44 | .. code-block:: bash
 45 | 
 46 |     apptainer exec --nv celldetection_latest.sif python -c "import celldetection; print(celldetection.__version__)"
 47 | 
 48 | You may remove ``--nv`` if you do not have GPUs on your system.
 49 | 
 50 | Installation in Python Environment
 51 | ----------------------------------
 52 | 
 53 | For users who wish to install `celldetection` directly in their Python environment, follow the steps below. Ensure that you have PyTorch installed as it is a critical dependency for the package. Visit the `PyTorch Installation Guide <https://pytorch.org/get-started/locally/>`_ for instructions.
 54 | 
 55 | **Virtual Environment Setup**
 56 | 
 57 | It's recommended to install `celldetection` in a virtual environment.
 58 | 
 59 | Using venv:
 60 | 
 61 | 1. Create a virtual environment:
 62 | 
 63 |    .. code-block:: bash
 64 | 
 65 |        python -m venv celldetection_env
 66 | 
 67 | 2. Activate the virtual environment:
 68 | 
 69 |    On Windows:
 70 | 
 71 |    .. code-block:: bash
 72 | 
 73 |        celldetection_env\Scripts\activate
 74 | 
 75 |    On macOS and Linux:
 76 | 
 77 |    .. code-block:: bash
 78 | 
 79 |        source celldetection_env/bin/activate
 80 | 
 81 | Using Conda:
 82 | 
 83 | 1. Create a Conda environment:
 84 | 
 85 |    .. code-block:: bash
 86 | 
 87 |        conda create -n celldetection_env python=3.x
 88 | 
 89 |    Replace `3.x` with the specific Python version you want to use.
 90 | 
 91 | 2. Activate the Conda environment:
 92 | 
 93 |    .. code-block:: bash
 94 | 
 95 |        conda activate celldetection_env
 96 | 
 97 | **PyPI Installation**
 98 | 
 99 | Install the latest stable release from PyPI:
100 | 
101 | .. code-block:: bash
102 | 
103 |     pip install -U celldetection
104 | 
105 | **GitHub Installation**
106 | 
107 | For the latest development version from GitHub:
108 | 
109 | .. code-block:: bash
110 | 
111 |     pip install git+https://github.com/FZJ-INM1-BDA/celldetection.git
112 | 
113 | Post-Installation
114 | ------------------
115 | 
116 | After installation, you can start using the `celldetection` package for your image processing tasks. If installed in a Python environment, remember to keep your virtual environment active. To exit the virtual environment, use the `deactivate` command.
117 | 


--------------------------------------------------------------------------------
/hubconf.py:
--------------------------------------------------------------------------------
 1 | """
 2 | PyTorch Hub models
 3 | 
 4 | Examples:
 5 |     import torch
 6 |     model = torch.hub.load('FZJ-INM1-BDA/celldetection', 'ginoro')
 7 | """
 8 | 
 9 | import celldetection as cd
10 | from torch.nn import Module
11 | from torch import device as _device
12 | from typing import Optional, Union
13 | 
14 | 
15 | def ginoro(
16 |         pretrained: bool = True,
17 |         pretrained_strict=True,
18 |         device: Optional[Union[_device, str, None]] = None,
19 |         **kwargs
20 | ) -> 'Module':
21 |     """
22 |     Ginoro: CPN + UNet + ResNeXt101
23 | 
24 |     References:
25 |         https://proceedings.mlr.press/v212/upschulte23a/upschulte23a.pdf
26 | 
27 |     Args:
28 |         pretrained: Whether to load the `state_dict` of a pretrained model.
29 |         pretrained_strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by
30 |             the loaded module’s state_dict() function.
31 |         device: Device to map the model to.
32 |         kwargs: Keyword arguments. Allows to override any setting of the model's constructor.
33 |     """
34 |     return cd.fetch_model('ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c', pretrained=pretrained,
35 |                           load_state_dict_kwargs=dict(check_hash=True), map_location=device,
36 |                           pretrained_strict=pretrained_strict, **kwargs)
37 | 


--------------------------------------------------------------------------------
/pytest.ini:
--------------------------------------------------------------------------------
1 | # pytest.ini
2 | [pytest]
3 | minversion = 6.0
4 | addopts = -ra -q
5 | testpaths =
6 |     tests
7 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | numpy
 2 | matplotlib
 3 | scipy
 4 | scikit-image
 5 | opencv-python
 6 | torch>=2.0.0
 7 | torchvision
 8 | tensorboard
 9 | seaborn
10 | tqdm
11 | h5py
12 | pillow
13 | nvidia-ml-py
14 | albumentations>=1.3.0
15 | timm
16 | segmentation-models-pytorch
17 | pandas
18 | pytorch_lightning
19 | pyyaml


--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
 1 | from setuptools import setup
 2 | from os.path import join, dirname, abspath
 3 | 
 4 | 
 5 | def read_utf8(*args):
 6 |     with open(join(*args), encoding="utf-8") as f:
 7 |         return f.read()
 8 | 
 9 | 
10 | directory, m = dirname(abspath(__file__)), {}
11 | exec(read_utf8(directory, 'celldetection', '__meta__.py'), m)
12 | requirements = read_utf8(directory, 'requirements.txt').strip().split("\n")
13 | long_description = read_utf8(directory, 'README.md')
14 | 
15 | setup(
16 |     author=m['__author__'],
17 |     author_email=m['__email__'],
18 |     name=m['__title__'],
19 |     version=m['__version__'],
20 |     description=m['__summary__'],
21 |     long_description=long_description,
22 |     long_description_content_type='text/markdown',
23 |     url=m['__url__'],
24 |     packages=['celldetection', 'celldetection.data', 'celldetection.callbacks', 'celldetection.optim',
25 |               'celldetection.data.datasets', 'celldetection.models', 'celldetection.mpi', 'celldetection.ops',
26 |               'celldetection.util', 'celldetection.visualization', 'celldetection_scripts'],
27 |     package_data={'': ['LICENSE', 'requirements.txt', 'README.md']},
28 |     include_package_data=True,
29 |     install_requires=requirements,
30 |     license=m['__license__'],
31 |     keywords=['cell', 'detection', 'object', 'segmentation', 'pytorch', 'cpn', 'contour', 'proposal', 'network', 'deep',
32 |               'learning', 'unet', 'fzj', 'julich', 'juelich', 'ai'],
33 |     classifiers=[
34 |         'Topic :: Scientific/Engineering :: Artificial Intelligence',
35 |         'Topic :: Scientific/Engineering :: Medical Science Apps.',
36 |         'License :: OSI Approved :: Apache Software License',
37 |         'Operating System :: OS Independent'
38 |     ],
39 |     entry_points={
40 |         'console_scripts': [
41 |             'cd-inference-cpn=celldetection_scripts.cpn_inference:main',
42 |             # 'cd-train=celldetection_scripts.train:main'
43 |          ]
44 |     }
45 | )
46 | 


--------------------------------------------------------------------------------
/tests/test_cpn_inference.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import celldetection as cd
 3 | import celldetection_scripts as cs
 4 | from skimage.data import coins
 5 | import pytest
 6 | 
 7 | 
 8 | def test_inference():
 9 |     accelerator = 'cpu'
10 |     model = 'ginoro_CpnResNeXt101UNet-fbe875f1a3e5ce2c'
11 |     img = coins()
12 |     print('Image:', img.shape)
13 |     cs.cpn_inference(img, model, accelerator=accelerator)
14 | 


--------------------------------------------------------------------------------