From fdea03e22e61244f5257dbc439dd31ef318652b4 Mon Sep 17 00:00:00 2001 From: Alister Burt Date: Fri, 5 Dec 2025 15:46:28 -0800 Subject: [PATCH 1/4] Migrate to PEP 517/621 with src-layout - Restructure package to use src/easymode/ layout - Replace setup.py with modern pyproject.toml - Add hatchling build backend with hatch-vcs versioning - Add comprehensive CI/CD via GitHub Actions - Add pre-commit hooks (ruff, mypy, validate-pyproject) - Add pytest testing infrastructure - Preserve GPL v3 license and all functionality - Version: v0.0.2 (managed via git tags) BREAKING CHANGES: - Minimum Python version now 3.10 (previously implicit 3.7+) - Package structure changed to src-layout (import paths unchanged) --- .github/ISSUE_TEMPLATE.md | 15 ++ .github/TEST_FAIL_TEMPLATE.md | 12 ++ .github/dependabot.yml | 10 + .github/workflows/ci.yml | 103 ++++++++++ .gitignore | 111 +++++++++++ .pre-commit-config.yaml | 42 ++++ CHANGELOG.md | 28 +++ CONTRIBUTING.md | 118 ++++++++++++ README.md | 24 ++- pyproject.toml | 182 ++++++++++++++++++ setup.py | 37 ---- {easymode => src/easymode}/__init__.py | 0 {easymode => src/easymode}/core/__init__.py | 0 .../easymode}/core/ais_wrapper.py | 0 {easymode => src/easymode}/core/config.py | 0 .../easymode}/core/distribution.py | 0 {easymode => src/easymode}/core/packaging.py | 0 {easymode => src/easymode}/core/settings.txt | 0 .../easymode}/core/warp_wrapper.py | 0 {easymode => src/easymode}/ddw/__init__.py | 0 {easymode => src/easymode}/ddw/inference.py | 0 {easymode => src/easymode}/ddw/loss.py | 0 {easymode => src/easymode}/ddw/model.py | 0 {easymode => src/easymode}/ddw/train.py | 0 {easymode => src/easymode}/main.py | 0 {easymode => src/easymode}/n2n/__init__.py | 0 {easymode => src/easymode}/n2n/inference.py | 0 {easymode => src/easymode}/n2n/model.py | 0 {easymode => src/easymode}/n2n/train.py | 0 .../__init__.py => src/easymode/py.typed | 0 .../easymode/segmentation}/__init__.py | 0 .../easymode}/segmentation/augmentations.py | 0 .../easymode}/segmentation/inference.py | 0 .../__init__.py | 0 .../filter_utils.py | 0 .../fourier_augmentations.py | 0 .../transforms.py | 0 .../easymode}/segmentation/model.py | 0 .../easymode}/segmentation/train.py | 0 tests/test_import.py | 45 +++++ 40 files changed, 689 insertions(+), 38 deletions(-) create mode 100644 .github/ISSUE_TEMPLATE.md create mode 100644 .github/TEST_FAIL_TEMPLATE.md create mode 100644 .github/dependabot.yml create mode 100644 .github/workflows/ci.yml create mode 100644 .gitignore create mode 100644 .pre-commit-config.yaml create mode 100644 CHANGELOG.md create mode 100644 CONTRIBUTING.md create mode 100644 pyproject.toml delete mode 100644 setup.py rename {easymode => src/easymode}/__init__.py (100%) rename {easymode => src/easymode}/core/__init__.py (100%) rename {easymode => src/easymode}/core/ais_wrapper.py (100%) rename {easymode => src/easymode}/core/config.py (100%) rename {easymode => src/easymode}/core/distribution.py (100%) rename {easymode => src/easymode}/core/packaging.py (100%) rename {easymode => src/easymode}/core/settings.txt (100%) rename {easymode => src/easymode}/core/warp_wrapper.py (100%) rename {easymode => src/easymode}/ddw/__init__.py (100%) rename {easymode => src/easymode}/ddw/inference.py (100%) rename {easymode => src/easymode}/ddw/loss.py (100%) rename {easymode => src/easymode}/ddw/model.py (100%) rename {easymode => src/easymode}/ddw/train.py (100%) rename {easymode => src/easymode}/main.py (100%) rename {easymode => src/easymode}/n2n/__init__.py (100%) rename {easymode => src/easymode}/n2n/inference.py (100%) rename {easymode => src/easymode}/n2n/model.py (100%) rename {easymode => src/easymode}/n2n/train.py (100%) rename easymode/segmentation/__init__.py => src/easymode/py.typed (100%) rename {easymode/segmentation/membrain_fourier_augmentations => src/easymode/segmentation}/__init__.py (100%) rename {easymode => src/easymode}/segmentation/augmentations.py (100%) rename {easymode => src/easymode}/segmentation/inference.py (100%) create mode 100644 src/easymode/segmentation/membrain_fourier_augmentations/__init__.py rename {easymode => src/easymode}/segmentation/membrain_fourier_augmentations/filter_utils.py (100%) rename {easymode => src/easymode}/segmentation/membrain_fourier_augmentations/fourier_augmentations.py (100%) rename {easymode => src/easymode}/segmentation/membrain_fourier_augmentations/transforms.py (100%) rename {easymode => src/easymode}/segmentation/model.py (100%) rename {easymode => src/easymode}/segmentation/train.py (100%) create mode 100644 tests/test_import.py diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000..09971ea --- /dev/null +++ b/.github/ISSUE_TEMPLATE.md @@ -0,0 +1,15 @@ +* easymode version: +* Python version: +* Operating System: + +### Description + +Describe what you were trying to get done. +Tell us what happened, what went wrong, and what you expected to happen. + +### What I Did + +``` +Paste the command(s) you ran and the output. +If there was a crash, please include the traceback here. +``` diff --git a/.github/TEST_FAIL_TEMPLATE.md b/.github/TEST_FAIL_TEMPLATE.md new file mode 100644 index 0000000..3512972 --- /dev/null +++ b/.github/TEST_FAIL_TEMPLATE.md @@ -0,0 +1,12 @@ +--- +title: "{{ env.TITLE }}" +labels: [bug] +--- +The {{ workflow }} workflow failed on {{ date | date("YYYY-MM-DD HH:mm") }} UTC + +The most recent failing test was on {{ env.PLATFORM }} py{{ env.PYTHON }} +with commit: {{ sha }} + +Full run: https://github.com/{{ repo }}/actions/runs/{{ env.RUN_ID }} + +(This post will be updated if another test fails, as long as this issue remains open.) diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..96505a9 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,10 @@ +# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates + +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + commit-message: + prefix: "ci(dependabot):" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..1c082c9 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,103 @@ +name: CI + +on: + push: + branches: [master] + tags: [v*] + pull_request: + workflow_dispatch: + schedule: + # run every week (for --pre release tests) + - cron: "0 0 * * 0" + +# cancel in-progress runs that use the same workflow and branch +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + check-manifest: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + - run: pipx run check-manifest + + test: + name: ${{ matrix.platform }} (${{ matrix.python-version }}) + runs-on: ${{ matrix.platform }} + env: + UV_PRERELEASE: ${{ github.event_name == 'schedule' && 'allow' || 'if-necessary-or-explicit' }} + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + platform: [ubuntu-latest, macos-latest, windows-latest] + + steps: + - uses: actions/checkout@v5 + + - name: 🐍 Set up Python ${{ matrix.python-version }} + uses: astral-sh/setup-uv@v6 + with: + python-version: ${{ matrix.python-version }} + enable-cache: true + + - name: Install Dependencies + run: uv sync --no-dev --group test + + - name: πŸ§ͺ Run Tests + run: uv run pytest --cov --cov-report=xml --cov-report=term-missing + + # If something goes wrong with --pre tests, we can open an issue in the repo + - name: πŸ“ Report --pre Failures + if: failure() && github.event_name == 'schedule' + uses: JasonEtco/create-an-issue@v2 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PLATFORM: ${{ matrix.platform }} + PYTHON: ${{ matrix.python-version }} + RUN_ID: ${{ github.run_id }} + TITLE: "[test-bot] pip install --pre is failing" + with: + filename: .github/TEST_FAIL_TEMPLATE.md + update_existing: true + + - name: Coverage + uses: codecov/codecov-action@v5 + # with: + # token: ${{ secrets.CODECOV_TOKEN }} + + build-and-inspect-package: + name: Build & inspect package. + needs: test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + with: + fetch-depth: 0 + - uses: hynek/build-and-inspect-python-package@v2 + + upload-to-pypi: + name: Upload package to PyPI + needs: build-and-inspect-package + if: success() && startsWith(github.ref, 'refs/tags/') && github.event_name != 'schedule' + runs-on: ubuntu-latest + permissions: + # IMPORTANT: this permission is mandatory for trusted publishing on PyPi + # see https://docs.pypi.org/trusted-publishers/ + id-token: write + # This permission allows writing releases + contents: write + + steps: + - name: Download built artifact to dist/ + uses: actions/download-artifact@v5 + with: + name: Packages + path: dist + - name: 🚒 Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + - uses: softprops/action-gh-release@v2 + with: + generate_release_notes: true + files: './dist/*' diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..13afad0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,111 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +.DS_Store + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# dotenv +.env + +# virtualenv +.venv +venv/ +ENV/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ + +# ruff +.ruff_cache/ + +# IDE settings +.vscode/ +.idea/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..36d14b1 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,42 @@ +# enable pre-commit.ci at https://pre-commit.ci/ +# it adds: +# 1. auto fixing pull requests +# 2. auto updating the pre-commit configuration +ci: + autoupdate_schedule: monthly + autofix_commit_msg: "style(pre-commit.ci): auto fixes [...]" + autoupdate_commit_msg: "ci(pre-commit.ci): autoupdate" + +repos: + - repo: https://github.com/abravalheri/validate-pyproject + rev: v0.24.1 + hooks: + - id: validate-pyproject + + - repo: https://github.com/rhysd/actionlint + rev: v1.7.7 + hooks: + - id: actionlint + + - repo: https://github.com/adhtruong/mirrors-typos + rev: v1.36.2 + hooks: + - id: typos + args: [--force-exclude] # omitting --write-changes + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.13.0 + hooks: + - id: ruff-check + args: [--fix] # may also add '--unsafe-fixes' + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.18.1 + hooks: + - id: mypy + files: "^src/" + # # you have to add the things you want to type check against here + # additional_dependencies: + # - numpy + diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..179fda8 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,28 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [0.0.2] - 2025-12-05 + +### Changed +- Migrated from setup.py to pyproject.toml (PEP 517/621) +- Restructured to src-layout for better packaging practices +- Added modern development tooling (ruff, mypy, pre-commit) +- Added comprehensive CI/CD via GitHub Actions +- Implemented git-based versioning via hatch-vcs + +### Infrastructure +- Python 3.10+ required (dropped 3.7-3.9 support) +- Added pytest-based testing infrastructure +- Added pre-commit hooks for code quality +- Added multi-platform CI (Ubuntu, macOS, Windows) +- All tests run on Python 3.10-3.13 + +### Maintained +- All existing functionality preserved +- GPL v3 license maintained +- All dependencies kept as-is +- Package imports unchanged (`import easymode.X`) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..a3e3ee3 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,118 @@ +# Contributing to easymode + +## Development Setup + +1. Clone the repository: +```bash +git clone https://github.com/mgflast/easymode.git +cd easymode +``` + +2. Create a virtual environment and install development dependencies: +```bash +python -m venv .venv +source .venv/bin/activate # On Windows: .venv\Scripts\activate +pip install -e ".[dev]" +``` + +3. Install pre-commit hooks: +```bash +pre-commit install +``` + +## Running Tests + +Run all tests: +```bash +pytest +``` + +With coverage: +```bash +pytest --cov --cov-report=html +``` + +View coverage report: +```bash +open htmlcov/index.html # On macOS +xdg-open htmlcov/index.html # On Linux +``` + +## Code Quality + +This project uses: +- **ruff** for linting and formatting +- **mypy** for type checking (gradual adoption) +- **pre-commit** for automated checks + +Before committing, run: +```bash +pre-commit run --all-files +``` + +Or rely on the pre-commit hooks to run automatically on `git commit`. + +## Code Style + +- Line length: 88 characters (Black/Ruff default) +- Docstring convention: NumPy style +- Import sorting: isort via ruff + +## Making Changes + +1. Create a new branch for your feature or bugfix: +```bash +git checkout -b feature/your-feature-name +``` + +2. Make your changes and commit them: +```bash +git add . +git commit -m "Description of changes" +``` + +3. Push to your fork and create a pull request: +```bash +git push origin feature/your-feature-name +``` + +## Making a Release + +Releases are managed via git tags and automated through GitHub Actions: + +1. Update CHANGELOG.md with the new version +2. Commit changes: +```bash +git add CHANGELOG.md +git commit -m "Prepare release vX.Y.Z" +``` + +3. Create and push a tag: +```bash +git tag -a vX.Y.Z -m "Release vX.Y.Z" +git push origin vX.Y.Z +``` + +4. GitHub Actions will automatically build and publish to PyPI + +## Project Structure + +``` +easymode/ +β”œβ”€β”€ src/easymode/ # Source code (src-layout) +β”‚ β”œβ”€β”€ main.py # CLI entry point +β”‚ β”œβ”€β”€ core/ # Core utilities +β”‚ β”œβ”€β”€ segmentation/ # 3D UNet segmentation +β”‚ β”œβ”€β”€ ddw/ # DeepDeWedge denoising +β”‚ └── n2n/ # Noise2Noise denoising +β”œβ”€β”€ tests/ # Test suite +β”œβ”€β”€ .github/workflows/ # CI/CD configuration +└── pyproject.toml # Project configuration + +``` + +## Questions or Issues? + +- Open an issue on GitHub +- Check existing issues and pull requests first +- Provide minimal reproducible examples for bugs diff --git a/README.md b/README.md index b0798a5..db5a0de 100644 --- a/README.md +++ b/README.md @@ -8,10 +8,32 @@ Easymode is a collection of general pretrained neural networks for cellular cryo-electron tomography (cryoET). The goal is to offer single command line interface functions to handle one of the most time-consuming and tedious steps of the cryoET workflow: feature detection. **Inspired and based upon [Membrain-seg](https://github.com/teamtomo/membrain-seg), easymode provides general pretrained networks for segmentation of cellular cryoET data.** -All networks were trained on a dataset of over 2000 tilt series from 50 different sources, covering many prokaryotic, archaeal, and eukaryotic species, different sample preparation techniques (cryo-FIB, lift-out, purified organelles, purified proteins, intact virus particles, whole cells), different hardware configurations (e.g., K2/K3/Falcon4i, 200/300 kV, Volta phase plate), a range of acquisition pixel sizes, electron doses, and defocus values, and lots of different original applications including various subtomogram averaging and benchmarking projects. +All networks were trained on a dataset of over 2000 tilt series from 50 different sources, covering many prokaryotic, archaeal, and eukaryotic species, different sample preparation techniques (cryo-FIB, lift-out, purified organelles, purified proteins, intact virus particles, whole cells), different hardware configurations (e.g., K2/K3/Falcon4i, 200/300 kV, Volta phase plate), a range of acquisition pixel sizes, electron doses, and defocus values, and lots of different original applications including various subtomogram averaging and benchmarking projects. The dataset is biased towards eukaryotic and specifically human cell types. For each tilt series in the collection, we reconstructed tomograms in five 'flavours' using [WarpTools](https://warpem.github.io/), [AreTomo3](https://github.com/czimaginginstitute/AreTomo3), [cryoCARE](https://github.com/juglab/cryoCARE_pip), and [DeepDeWedge](https://github.com/MLI-lab/DeepDeWedge), all at a common voxel size of 10.0 Γ…ngstrom. Training was performed using linear combinations of different tomogram flavours as the input; we use almost exactly the same 3D UNet architecture and many of the same augmentations as used in [Membrain-seg](https://github.com/teamtomo/membrain-seg), including their missing wedge augmentation. Labels were generated by applying [Ais](github.com/bionanopatterning/Ais) 2D UNets to selected subtomograms and manually curating the results. +### installation + +**From PyPI (recommended):** +```bash +pip install easymode +``` + +**From source:** +```bash +git clone https://github.com/mgflast/easymode.git +cd easymode +pip install -e . +``` + +**For development:** +```bash +git clone https://github.com/mgflast/easymode.git +cd easymode +pip install -e ".[dev]" +pre-commit install +``` + ### validation We are currently testing the accuracy of the feature detection networks by subtomogram averaging (STA). For a start the goal is to enable decent-resolution STA of the most common particle types, so that these can be used as references for tilt-series alignment and to navigate datasets better. **Please note that these are all preliminary results.** We use three moderately large eukaryotic datasets for the validation: S. cerevisiae (500 tilt series, by the MPI Dortmund team (see table below)), C. reinhardtii (500 tilt series, EMPIAR-11830), and H. sapiens (HeLa) (620 tilt series, our own) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..cde481f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,182 @@ +# https://peps.python.org/pep-0517/ +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + +# https://hatch.pypa.io/latest/config/metadata/ +[tool.hatch.version] +source = "vcs" + +# read more about configuring hatch at: +# https://hatch.pypa.io/latest/config/build/ +[tool.hatch.build.targets.wheel] +only-include = ["src"] +sources = ["src"] + +# Include non-Python files as package data +[tool.hatch.build] +artifacts = [ + "*.txt", + "*.png", + "*.glsl", + "*.pdf", +] + +# https://peps.python.org/pep-0621/ +[project] +name = "easymode" +dynamic = ["version"] +description = "Easymode - a collection of pretrained general networks for segmenting common eukaryotic features in cryoET" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "GPL-3.0-or-later" } +authors = [ + { name = "Mart Last", email = "mgflast@gmail.com" } +] +keywords = [ + "cryo-electron tomography", + "cryoET", + "deep learning", + "segmentation", + "neural networks", + "structural biology", +] +# https://pypi.org/classifiers/ +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Bio-Informatics", + "Topic :: Scientific/Engineering :: Image Recognition", +] + +dependencies = [ + "tensorflow>=2.8.0,<3.0.0", + "mrcfile", + "numpy", + "scipy", + "huggingface_hub", + "requests", + "tifffile", + "psutil", + "starfile", +] + +[project.urls] +homepage = "https://github.com/mgflast/easymode" +repository = "https://github.com/mgflast/easymode" + +# Entry points +# https://peps.python.org/pep-0621/#entry-points +[project.scripts] +easymode = "easymode.main:main" + +# https://peps.python.org/pep-0735/ +# Development dependencies using PEP 735 dependency groups +[dependency-groups] +test = ["pytest>=7.0", "pytest-cov"] +dev = [ + { include-group = "test" }, + "ipython", + "mypy", + "pre-commit-uv", + "ruff", +] + +# https://docs.astral.sh/ruff +[tool.ruff] +line-length = 88 +target-version = "py310" +src = ["src"] +fix = true + +[tool.ruff.lint] +pydocstyle = { convention = "numpy" } +select = [ + "E", # style errors + "W", # style warnings + "F", # flakes + "I", # isort + "UP", # pyupgrade + "C4", # flake8-comprehensions + "B", # flake8-bugbear + "A001", # flake8-builtins + "RUF", # ruff-specific rules +] +ignore = [ + "D", # Ignore all pydocstyle for now (existing code not documented) +] + +[tool.ruff.lint.per-file-ignores] +"tests/*.py" = ["D", "S"] + +# https://docs.astral.sh/ruff/formatter/ +[tool.ruff.format] +docstring-code-format = true +skip-magic-trailing-comma = false + +# https://docs.pytest.org/ +[tool.pytest.ini_options] +minversion = "7.0" +addopts = ["--color=yes"] +testpaths = ["tests"] +# Don't treat warnings as errors due to TensorFlow warnings +filterwarnings = [ + "ignore::DeprecationWarning", + "ignore::UserWarning", +] + +# https://mypy.readthedocs.io/en/stable/config_file.html +[tool.mypy] +files = "src/**/" +# Start with less strict settings for ML code +warn_return_any = false +warn_unused_configs = true +disallow_untyped_defs = false + +# Ignore missing imports for scientific packages +[[tool.mypy.overrides]] +module = [ + "tensorflow.*", + "mrcfile.*", + "scipy.*", + "huggingface_hub.*", + "tifffile.*", + "starfile.*", +] +ignore_missing_imports = true + +# https://coverage.readthedocs.io/ +[tool.coverage.report] +show_missing = true +exclude_lines = [ + "pragma: no cover", + "if TYPE_CHECKING:", + "@overload", + "except ImportError", + "raise AssertionError", + "\\.\\.\\.", + "raise NotImplementedError()", + "pass", + "if __name__ == .__main__.:", +] + +[tool.coverage.run] +source = ["easymode"] + +# https://github.com/mgedmin/check-manifest#configuration +[tool.check-manifest] +ignore = [ + ".pre-commit-config.yaml", + ".ruff_cache/**/*", + "tests/**/*", + ".github/**/*", + "*.yml", + "*.yaml", +] diff --git a/setup.py b/setup.py deleted file mode 100644 index d017bdd..0000000 --- a/setup.py +++ /dev/null @@ -1,37 +0,0 @@ -from setuptools import setup, find_packages -import shutil -# how to release: -# UPDATE VERSION IN 3 PLACES: Ais/core/config.py, setup.py, docs/conf.py - -# push to pypi: -# python setup.py sdist -# twine upload dist/* - -VERSION = '0.0.2' - -setup( - name='easymode', - version=VERSION, - packages=find_packages(), - entry_points={'console_scripts': ['easymode=easymode.main:main']}, - url='', - license='GPL v3', - author='mgflast', - author_email='mgflast@gmail.com', - description='Easymode - a collection of pretrained general networks for segmenting common eukaryotic features in cryoET', - package_data={'': ['*.png', '*.glsl', '*.pdf', '*.txt']}, - include_package_data=False, - install_requires=[ - "tensorflow>=2.8.0,<3.0.0", - "mrcfile", - "numpy", - "scipy", - "huggingface_hub", - "requests", - "tifffile", - "psutil", - "starfile" - #"Ais-cryoET" - ] -) - diff --git a/easymode/__init__.py b/src/easymode/__init__.py similarity index 100% rename from easymode/__init__.py rename to src/easymode/__init__.py diff --git a/easymode/core/__init__.py b/src/easymode/core/__init__.py similarity index 100% rename from easymode/core/__init__.py rename to src/easymode/core/__init__.py diff --git a/easymode/core/ais_wrapper.py b/src/easymode/core/ais_wrapper.py similarity index 100% rename from easymode/core/ais_wrapper.py rename to src/easymode/core/ais_wrapper.py diff --git a/easymode/core/config.py b/src/easymode/core/config.py similarity index 100% rename from easymode/core/config.py rename to src/easymode/core/config.py diff --git a/easymode/core/distribution.py b/src/easymode/core/distribution.py similarity index 100% rename from easymode/core/distribution.py rename to src/easymode/core/distribution.py diff --git a/easymode/core/packaging.py b/src/easymode/core/packaging.py similarity index 100% rename from easymode/core/packaging.py rename to src/easymode/core/packaging.py diff --git a/easymode/core/settings.txt b/src/easymode/core/settings.txt similarity index 100% rename from easymode/core/settings.txt rename to src/easymode/core/settings.txt diff --git a/easymode/core/warp_wrapper.py b/src/easymode/core/warp_wrapper.py similarity index 100% rename from easymode/core/warp_wrapper.py rename to src/easymode/core/warp_wrapper.py diff --git a/easymode/ddw/__init__.py b/src/easymode/ddw/__init__.py similarity index 100% rename from easymode/ddw/__init__.py rename to src/easymode/ddw/__init__.py diff --git a/easymode/ddw/inference.py b/src/easymode/ddw/inference.py similarity index 100% rename from easymode/ddw/inference.py rename to src/easymode/ddw/inference.py diff --git a/easymode/ddw/loss.py b/src/easymode/ddw/loss.py similarity index 100% rename from easymode/ddw/loss.py rename to src/easymode/ddw/loss.py diff --git a/easymode/ddw/model.py b/src/easymode/ddw/model.py similarity index 100% rename from easymode/ddw/model.py rename to src/easymode/ddw/model.py diff --git a/easymode/ddw/train.py b/src/easymode/ddw/train.py similarity index 100% rename from easymode/ddw/train.py rename to src/easymode/ddw/train.py diff --git a/easymode/main.py b/src/easymode/main.py similarity index 100% rename from easymode/main.py rename to src/easymode/main.py diff --git a/easymode/n2n/__init__.py b/src/easymode/n2n/__init__.py similarity index 100% rename from easymode/n2n/__init__.py rename to src/easymode/n2n/__init__.py diff --git a/easymode/n2n/inference.py b/src/easymode/n2n/inference.py similarity index 100% rename from easymode/n2n/inference.py rename to src/easymode/n2n/inference.py diff --git a/easymode/n2n/model.py b/src/easymode/n2n/model.py similarity index 100% rename from easymode/n2n/model.py rename to src/easymode/n2n/model.py diff --git a/easymode/n2n/train.py b/src/easymode/n2n/train.py similarity index 100% rename from easymode/n2n/train.py rename to src/easymode/n2n/train.py diff --git a/easymode/segmentation/__init__.py b/src/easymode/py.typed similarity index 100% rename from easymode/segmentation/__init__.py rename to src/easymode/py.typed diff --git a/easymode/segmentation/membrain_fourier_augmentations/__init__.py b/src/easymode/segmentation/__init__.py similarity index 100% rename from easymode/segmentation/membrain_fourier_augmentations/__init__.py rename to src/easymode/segmentation/__init__.py diff --git a/easymode/segmentation/augmentations.py b/src/easymode/segmentation/augmentations.py similarity index 100% rename from easymode/segmentation/augmentations.py rename to src/easymode/segmentation/augmentations.py diff --git a/easymode/segmentation/inference.py b/src/easymode/segmentation/inference.py similarity index 100% rename from easymode/segmentation/inference.py rename to src/easymode/segmentation/inference.py diff --git a/src/easymode/segmentation/membrain_fourier_augmentations/__init__.py b/src/easymode/segmentation/membrain_fourier_augmentations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/easymode/segmentation/membrain_fourier_augmentations/filter_utils.py b/src/easymode/segmentation/membrain_fourier_augmentations/filter_utils.py similarity index 100% rename from easymode/segmentation/membrain_fourier_augmentations/filter_utils.py rename to src/easymode/segmentation/membrain_fourier_augmentations/filter_utils.py diff --git a/easymode/segmentation/membrain_fourier_augmentations/fourier_augmentations.py b/src/easymode/segmentation/membrain_fourier_augmentations/fourier_augmentations.py similarity index 100% rename from easymode/segmentation/membrain_fourier_augmentations/fourier_augmentations.py rename to src/easymode/segmentation/membrain_fourier_augmentations/fourier_augmentations.py diff --git a/easymode/segmentation/membrain_fourier_augmentations/transforms.py b/src/easymode/segmentation/membrain_fourier_augmentations/transforms.py similarity index 100% rename from easymode/segmentation/membrain_fourier_augmentations/transforms.py rename to src/easymode/segmentation/membrain_fourier_augmentations/transforms.py diff --git a/easymode/segmentation/model.py b/src/easymode/segmentation/model.py similarity index 100% rename from easymode/segmentation/model.py rename to src/easymode/segmentation/model.py diff --git a/easymode/segmentation/train.py b/src/easymode/segmentation/train.py similarity index 100% rename from easymode/segmentation/train.py rename to src/easymode/segmentation/train.py diff --git a/tests/test_import.py b/tests/test_import.py new file mode 100644 index 0000000..1fe096e --- /dev/null +++ b/tests/test_import.py @@ -0,0 +1,45 @@ +"""Basic import tests to verify package installation.""" + + +def test_import_main(): + """Test that main module can be imported.""" + import easymode.main + + assert hasattr(easymode.main, "main") + + +def test_import_core(): + """Test that core modules can be imported.""" + import easymode.core.config + + assert hasattr(easymode.core.config, "settings") + + +def test_import_segmentation(): + """Test that segmentation modules can be imported.""" + import easymode.segmentation.inference + + # Verify module loads + assert easymode.segmentation.inference is not None + + +def test_import_ddw(): + """Test that ddw modules can be imported.""" + import easymode.ddw.inference + + assert easymode.ddw.inference is not None + + +def test_import_n2n(): + """Test that n2n modules can be imported.""" + import easymode.n2n.inference + + assert easymode.n2n.inference is not None + + +def test_cli_entry_point(): + """Test that CLI entry point exists.""" + from easymode.main import main + + # Entry point exists + assert callable(main) From 0c6bce404251e26f3a027a96b7a52f2da145162f Mon Sep 17 00:00:00 2001 From: Alister Burt Date: Fri, 5 Dec 2025 16:05:33 -0800 Subject: [PATCH 2/4] update install/run instructions --- README.md | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index db5a0de..825cda1 100644 --- a/README.md +++ b/README.md @@ -14,19 +14,33 @@ The dataset is biased towards eukaryotic and specifically human cell types. For ### installation -**From PyPI (recommended):** -```bash -pip install easymode +#### User + +Install [*uv*](https://docs.astral.sh/uv/getting-started/installation/) then run all +*easymode* commands with *uvx*. + +```shell +uvx easymode --help ``` -**From source:** +*uv* will manage a Python environment for you behind the scenes, easy! 🀠 + +#### Multi User Installation + +If installing for multiple users, please install easymode as you would any other Python +package. + ```bash git clone https://github.com/mgflast/easymode.git -cd easymode -pip install -e . + +pip install easymode ``` -**For development:** + +#### Developer + +Assuming understanding of Python development workflows: + ```bash git clone https://github.com/mgflast/easymode.git cd easymode From 913411665c8c0650075556b509e003aa8fddfc68 Mon Sep 17 00:00:00 2001 From: Alister Burt Date: Fri, 5 Dec 2025 16:06:38 -0800 Subject: [PATCH 3/4] exclude mac and windows tests --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1c082c9..6bf10c9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,7 +31,7 @@ jobs: fail-fast: false matrix: python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - platform: [ubuntu-latest, macos-latest, windows-latest] + platform: [ubuntu-latest, ] # macos-latest, windows-latest] steps: - uses: actions/checkout@v5 From 1187ec09f005629e5f6caeba08d7c9763dee10e8 Mon Sep 17 00:00:00 2001 From: Alister Burt Date: Sun, 7 Dec 2025 21:49:56 -0800 Subject: [PATCH 4/4] torch for inference? --- pyproject.toml | 4 + .../segmentation/torch_version/__init__.py | 0 .../torch_version/convert_weights.py | 685 ++++++++++++++++++ .../segmentation/torch_version/model_torch.py | 337 +++++++++ .../segmentation/torch_version/torch_utils.py | 304 ++++++++ .../torch_version/validate_equivalence.py | 315 ++++++++ 6 files changed, 1645 insertions(+) create mode 100644 src/easymode/segmentation/torch_version/__init__.py create mode 100644 src/easymode/segmentation/torch_version/convert_weights.py create mode 100644 src/easymode/segmentation/torch_version/model_torch.py create mode 100644 src/easymode/segmentation/torch_version/torch_utils.py create mode 100644 src/easymode/segmentation/torch_version/validate_equivalence.py diff --git a/pyproject.toml b/pyproject.toml index cde481f..3ca5938 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,6 +87,8 @@ dev = [ "mypy", "pre-commit-uv", "ruff", + "torch>=2.0.0", + "h5py", ] # https://docs.astral.sh/ruff @@ -144,6 +146,8 @@ disallow_untyped_defs = false [[tool.mypy.overrides]] module = [ "tensorflow.*", + "torch.*", + "h5py.*", "mrcfile.*", "scipy.*", "huggingface_hub.*", diff --git a/src/easymode/segmentation/torch_version/__init__.py b/src/easymode/segmentation/torch_version/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/easymode/segmentation/torch_version/convert_weights.py b/src/easymode/segmentation/torch_version/convert_weights.py new file mode 100644 index 0000000..79bb401 --- /dev/null +++ b/src/easymode/segmentation/torch_version/convert_weights.py @@ -0,0 +1,685 @@ +"""Convert TensorFlow/Keras HDF5 weights to PyTorch format. + +This script converts weights from the TensorFlow 3D U-Net model to the PyTorch +implementation, handling tensor format conversion and layer name mapping. + +The main challenges are: +1. Conv3D weights: TensorFlow (D, H, W, Cin, Cout) β†’ PyTorch (Cout, Cin, D, H, W) +2. BatchNorm params: gammaβ†’weight, betaβ†’bias, moving_meanβ†’running_mean, etc. +3. Layer names: TensorFlow 'encoder_0/...' β†’ PyTorch 'encoders.0....' + +Strategy: +We load both the TF model and PyTorch model, then copy weights in order based on +the model structure rather than trying to parse auto-generated names. +""" + +import argparse +from collections import OrderedDict +from pathlib import Path + +import h5py +import numpy as np +import torch +import tensorflow as tf + + +def load_tf_weights_from_h5(h5_path): + """Load all TensorFlow weights from HDF5 file. + + Parameters + ---------- + h5_path : str or Path + Path to .h5 weights file + + Returns + ------- + dict + Mapping of weight names to numpy arrays + """ + weights = {} + + with h5py.File(h5_path, 'r') as f: + # Recursively extract all datasets + def extract_weights(name, obj): + if isinstance(obj, h5py.Dataset): + weights[name] = np.array(obj) + + f.visititems(extract_weights) + + return weights + + +def convert_conv3d_weight(tf_weight): + """Convert Conv3D weight from TensorFlow to PyTorch format. + + Parameters + ---------- + tf_weight : np.ndarray + TensorFlow weight with shape (D, H, W, Cin, Cout) + + Returns + ------- + torch.Tensor + PyTorch weight with shape (Cout, Cin, D, H, W) + """ + # Transpose from (D, H, W, Cin, Cout) to (Cout, Cin, D, H, W) + pytorch_weight = np.transpose(tf_weight, (4, 3, 0, 1, 2)) + return torch.from_numpy(pytorch_weight).float() + + +def convert_batchnorm_weights(gamma, beta, moving_mean, moving_var): + """Convert BatchNorm weights from TensorFlow to PyTorch. + + Parameters + ---------- + gamma : np.ndarray + TensorFlow gamma (scale) + beta : np.ndarray + TensorFlow beta (shift) + moving_mean : np.ndarray + TensorFlow moving mean + moving_var : np.ndarray + TensorFlow moving variance + + Returns + ------- + dict + Dictionary with PyTorch BatchNorm parameters + """ + return { + 'weight': torch.from_numpy(gamma).float(), + 'bias': torch.from_numpy(beta).float(), + 'running_mean': torch.from_numpy(moving_mean).float(), + 'running_var': torch.from_numpy(moving_var).float(), + } + + +def map_tf_to_pytorch_names(tf_weights): + """Map TensorFlow weight names to PyTorch state_dict keys. + + This function implements the mapping between TensorFlow layer naming + and PyTorch module naming conventions. + + Parameters + ---------- + tf_weights : dict + TensorFlow weights loaded from HDF5 + + Returns + ------- + OrderedDict + PyTorch state_dict with converted weights + """ + pytorch_state = OrderedDict() + + # Group weights by block (encoder_N, decoder_N, output) + blocks = {} + for name, weight in tf_weights.items(): + # Skip optimizer states and other metadata + if 'optimizer' in name.lower() or 'iteration' in name.lower(): + continue + + # Extract block name (encoder_0, decoder_1, output, etc.) + parts = name.split('/') + if len(parts) < 1: + continue + + block_name = parts[0] + if block_name not in blocks: + blocks[block_name] = {} + blocks[block_name][name] = weight + + # Process each block + for block_name, block_weights in sorted(blocks.items()): + if block_name.startswith('encoder_'): + idx = int(block_name.split('_')[1]) + convert_encoder_block(block_weights, idx, pytorch_state) + elif block_name.startswith('decoder_'): + idx = int(block_name.split('_')[1]) + convert_decoder_block(block_weights, idx, pytorch_state) + elif block_name == 'output': + convert_output_layer(block_weights, pytorch_state) + + return pytorch_state + + +def convert_encoder_block(block_weights, idx, pytorch_state): + """Convert an encoder block's weights. + + Parameters + ---------- + block_weights : dict + All weights for this encoder block + idx : int + Encoder index (0-5) + pytorch_state : OrderedDict + PyTorch state dict to populate + """ + prefix = f'encoders.{idx}' + + # Separate downsample and res_block weights + downsample_weights = {} + resblock_weights = {} + + for name, weight in block_weights.items(): + # Remove redundant nesting: encoder_N/u_net/encoder_N/... + if 'res_block' in name: + resblock_weights[name] = weight + else: + downsample_weights[name] = weight + + # Process downsample (exists for idx > 0) + if idx > 0: + # Find downsample conv + for name, weight in downsample_weights.items(): + if 'kernel' in name and 'res_block' not in name: + # This is the downsample conv + # Note: TFSameConv3d has a nested .conv module + conv_weight = convert_conv3d_weight(weight) + pytorch_state[f'{prefix}.downsample.conv.weight'] = conv_weight + break + + # Find downsample BN - collect all params at once + bn_groups = {} + for name, weight in downsample_weights.items(): + if 'res_block' in name: + continue + if any(x in name for x in ['gamma', 'beta', 'moving_mean', 'moving_variance']): + bn_path = '/'.join(name.split('/')[:-1]) + if bn_path not in bn_groups: + bn_groups[bn_path] = {} + param = name.split('/')[-1].replace(':0', '') + bn_groups[bn_path][param] = weight + + # Assign the first complete BN group as downsample_bn + for bn_path, bn_params in bn_groups.items(): + if len(bn_params) == 4: + converted = convert_batchnorm_weights( + bn_params['gamma'], + bn_params['beta'], + bn_params['moving_mean'], + bn_params['moving_variance'] + ) + for key, value in converted.items(): + pytorch_state[f'{prefix}.downsample_bn.{key}'] = value + break + + # Process res_block + convert_resblock(resblock_weights, f'{prefix}.res_block', pytorch_state) + + +def convert_decoder_block(block_weights, idx, pytorch_state): + """Convert a decoder block's weights. + + Parameters + ---------- + block_weights : dict + All weights for this decoder block + idx : int + Decoder index (0-4) + pytorch_state : OrderedDict + PyTorch state dict to populate + """ + prefix = f'decoders.{idx}' + + # Separate upsample and res_block weights + upsample_weights = {} + resblock_weights = {} + + for name, weight in block_weights.items(): + if 'res_block' in name: + resblock_weights[name] = weight + else: + upsample_weights[name] = weight + + # Process upsample conv + for name, weight in upsample_weights.items(): + if 'conv3d_transpose' in name and 'kernel' in name: + conv_weight = convert_conv3d_weight(weight) + pytorch_state[f'{prefix}.upsample.weight'] = conv_weight + break + + # Find upsample BN - collect all params at once + bn_groups = {} + for name, weight in upsample_weights.items(): + if 'res_block' in name: + continue + if any(x in name for x in ['gamma', 'beta', 'moving_mean', 'moving_variance']): + bn_path = '/'.join(name.split('/')[:-1]) + if bn_path not in bn_groups: + bn_groups[bn_path] = {} + param = name.split('/')[-1].replace(':0', '') + bn_groups[bn_path][param] = weight + + # Assign the first complete BN group as upsample_bn + for bn_path, bn_params in bn_groups.items(): + if len(bn_params) == 4: + converted = convert_batchnorm_weights( + bn_params['gamma'], + bn_params['beta'], + bn_params['moving_mean'], + bn_params['moving_variance'] + ) + for key, value in converted.items(): + pytorch_state[f'{prefix}.upsample_bn.{key}'] = value + break + + # Process res_block + convert_resblock(resblock_weights, f'{prefix}.res_block', pytorch_state) + + +def convert_resblock(block_weights, prefix, pytorch_state): + """Convert a ResBlock's weights. + + ResBlocks have: + - conv1 (3x3x3) + bn1 + - conv2 (3x3x3) + bn2 + - skip_conv (1x1x1) + skip_bn (optional, when in_channels != out_channels) + + Parameters + ---------- + block_weights : dict + All weights for this res_block + prefix : str + PyTorch prefix (e.g., 'encoders.0.res_block') + pytorch_state : OrderedDict + PyTorch state dict to populate + """ + # Collect all conv kernels and BN params + conv_kernels = [] + bn_groups = [] + + # Group BN parameters by their path + bn_paths = {} + for name, weight in block_weights.items(): + if 'kernel' in name: + conv_kernels.append((name, weight)) + elif 'gamma' in name: + bn_path = '/'.join(name.split('/')[:-1]) + if bn_path not in bn_paths: + bn_paths[bn_path] = {} + param = name.split('/')[-1].replace(':0', '') + bn_paths[bn_path][param] = weight + elif 'beta' in name or 'moving_mean' in name or 'moving_variance' in name: + bn_path = '/'.join(name.split('/')[:-1]) + if bn_path not in bn_paths: + bn_paths[bn_path] = {} + param = name.split('/')[-1].replace(':0', '') + bn_paths[bn_path][param] = weight + + # Sort conv kernels by numeric suffix in name (e.g., conv3d_9, conv3d_10) + # Extract number from path like '/encoder_3/.../conv3d_9/kernel:0' + import re + def get_layer_number(name): + # Get the layer name (e.g., 'conv3d_9' from '.../conv3d_9/kernel:0') + layer_name = name.split('/')[-2] + # Extract the trailing number (e.g., '9' from 'conv3d_9') + match = re.search(r'_(\d+)$', layer_name) + return int(match.group(1)) if match else 0 + conv_kernels.sort(key=lambda x: get_layer_number(x[0])) + + # Identify convs by kernel size + conv1_weight = None + conv2_weight = None + skip_conv_weight = None + + for name, kernel in conv_kernels: + shape = kernel.shape + if shape[0] == 1 and shape[1] == 1 and shape[2] == 1: + # 1x1x1 kernel = skip_conv + skip_conv_weight = kernel + elif conv1_weight is None: + # First 3x3x3 kernel = conv1 + conv1_weight = kernel + else: + # Second 3x3x3 kernel = conv2 + conv2_weight = kernel + + # Convert and assign conv weights + if conv1_weight is not None: + pytorch_state[f'{prefix}.conv1.weight'] = convert_conv3d_weight(conv1_weight) + if conv2_weight is not None: + pytorch_state[f'{prefix}.conv2.weight'] = convert_conv3d_weight(conv2_weight) + if skip_conv_weight is not None: + pytorch_state[f'{prefix}.skip_conv.weight'] = convert_conv3d_weight(skip_conv_weight) + + # Process BN groups + # We have 2 or 3 BN groups (bn1, bn2, and optionally skip_bn) + bn_groups_list = [] + for bn_path, params in bn_paths.items(): + if len(params) == 4: # Complete BN + bn_groups_list.append((bn_path, params)) + + # Sort by numeric suffix in path name (e.g., batch_normalization_9, batch_normalization_10) + def get_bn_number(path): + bn_name = path.split('/')[-1] + # Extract trailing number (e.g., '9' from 'batch_normalization_9') + match = re.search(r'_(\d+)$', bn_name) + return int(match.group(1)) if match else 0 + bn_groups_list.sort(key=lambda x: get_bn_number(x[0])) + + # Assign BN groups + # The order should be: bn1, bn2, skip_bn (if present) + # But in the HDF5, they might be in different order + # Use the number of BN groups to determine + if len(bn_groups_list) == 2: + # No skip connection + convert_and_assign_bn(bn_groups_list[0][1], f'{prefix}.bn1', pytorch_state) + convert_and_assign_bn(bn_groups_list[1][1], f'{prefix}.bn2', pytorch_state) + elif len(bn_groups_list) == 3: + # Has skip connection + # The skip_bn usually has a higher number in the auto-generated name + # So the last one is likely skip_bn + convert_and_assign_bn(bn_groups_list[0][1], f'{prefix}.bn1', pytorch_state) + convert_and_assign_bn(bn_groups_list[1][1], f'{prefix}.bn2', pytorch_state) + convert_and_assign_bn(bn_groups_list[2][1], f'{prefix}.skip_bn', pytorch_state) + + +def convert_and_assign_bn(bn_params, prefix, pytorch_state): + """Convert and assign BatchNorm parameters.""" + converted = convert_batchnorm_weights( + bn_params['gamma'], + bn_params['beta'], + bn_params['moving_mean'], + bn_params['moving_variance'] + ) + for key, value in converted.items(): + pytorch_state[f'{prefix}.{key}'] = value + + +def convert_output_layer(block_weights, pytorch_state): + """Convert the output layer weights. + + Parameters + ---------- + block_weights : dict + All weights for the output layer + pytorch_state : OrderedDict + PyTorch state dict to populate + """ + for name, weight in block_weights.items(): + if 'kernel' in name: + conv_weight = convert_conv3d_weight(weight) + pytorch_state['final_conv.weight'] = conv_weight + elif 'bias' in name: + bias_weight = torch.from_numpy(weight).float() + pytorch_state['final_conv.bias'] = bias_weight + + +def convert_weights_direct(h5_path, output_path, verbose=True): + """Convert weights by loading both TF and PyTorch models. + + This is a more robust approach that avoids fragile name mapping. + + Parameters + ---------- + h5_path : str or Path + Path to TensorFlow .h5 weights file + output_path : str or Path + Path to save PyTorch .pth weights file + verbose : bool, optional + Whether to print conversion details (default: True) + + Returns + ------- + OrderedDict + The converted PyTorch state_dict + """ + from easymode.segmentation.model import create as create_tf + from easymode.segmentation.torch_version.model_torch import create as create_torch + + h5_path = Path(h5_path) + output_path = Path(output_path) + + if not h5_path.exists(): + raise FileNotFoundError(f"Input file not found: {h5_path}") + + if verbose: + print(f"Loading TensorFlow model and weights from: {h5_path}") + + # Create TensorFlow model and load weights + tf_model = create_tf() + # Build the model with a dummy forward pass + dummy_tf = tf.zeros((1, 64, 64, 64, 1)) + _ = tf_model(dummy_tf) + tf_model.load_weights(h5_path) + + if verbose: + print(f"Loaded TensorFlow model with {len(tf_model.trainable_variables)} trainable vars") + + # Create PyTorch model + torch_model = create_torch() + + if verbose: + print(f"Created PyTorch model") + + # Convert weights by iterating through both models in parallel + pytorch_state = OrderedDict() + + # Get TF variables + tf_vars = tf_model.trainable_variables + + # Get PyTorch parameters + torch_params = list(torch_model.named_parameters()) + + if verbose: + print(f"TensorFlow trainable vars: {len(tf_vars)}") + print(f"PyTorch trainable params: {len(torch_params)}") + + # Manual mapping based on model structure + # This requires understanding the exact correspondence + var_idx = 0 + + def copy_conv_weight(tf_var, torch_name): + """Copy and convert Conv3D weight.""" + nonlocal var_idx + tf_weight = tf_var.numpy() + # TF: (D, H, W, Cin, Cout) -> PyTorch: (Cout, Cin, D, H, W) + torch_weight = np.transpose(tf_weight, (4, 3, 0, 1, 2)) + pytorch_state[torch_name] = torch.from_numpy(torch_weight).float() + if verbose and var_idx < 5: + print(f" Conv: {tf_var.name} -> {torch_name}") + print(f" TF shape: {tf_weight.shape} -> PyTorch shape: {torch_weight.shape}") + var_idx += 1 + + def copy_bn_params(gamma, beta, mean, var, torch_prefix): + """Copy BatchNorm parameters.""" + nonlocal var_idx + pytorch_state[f'{torch_prefix}.weight'] = torch.from_numpy(gamma.numpy()).float() + pytorch_state[f'{torch_prefix}.bias'] = torch.from_numpy(beta.numpy()).float() + pytorch_state[f'{torch_prefix}.running_mean'] = torch.from_numpy(mean.numpy()).float() + pytorch_state[f'{torch_prefix}.running_var'] = torch.from_numpy(var.numpy()).float() + # PyTorch also has num_batches_tracked but it's not in TF, will be initialized + if verbose and var_idx < 5: + print(f" BN: {gamma.name} -> {torch_prefix}") + var_idx += 4 + + # Iterate through model structure + # Encoders + for i, (tf_enc, torch_enc) in enumerate(zip(tf_model.encoders, torch_model.encoders)): + enc_prefix = f'encoders.{i}' + + # Downsample (if stride > 1) + if hasattr(tf_enc, 'downsample') and tf_enc.downsample is not None: + # Find downsample vars in TF + for var in tf_enc.trainable_variables: + if 'downsample' in var.name or (i > 0 and 'conv3d' in var.name and 'res_block' not in var.name): + if 'kernel' in var.name: + copy_conv_weight(var, f'{enc_prefix}.downsample.weight') + break + + # Downsample BN + for j, var in enumerate(tf_enc.trainable_variables): + if ('batch_normalization' in var.name or 'downsample_bn' in var.name) and 'res_block' not in var.name: + # Collect BN params + bn_vars = [v for v in tf_enc.trainable_variables if 'batch_normalization' in v.name and 'res_block' not in var.name] + if len(bn_vars) >= 4: + gamma = [v for v in bn_vars if 'gamma' in v.name][0] + beta = [v for v in bn_vars if 'beta' in v.name][0] + + # Get non-trainable (moving stats) + all_vars = tf_enc.variables + bn_all = [v for v in all_vars if 'batch_normalization' in v.name and 'res_block' not in v.name] + mean = [v for v in bn_all if 'moving_mean' in v.name][0] + var_v = [v for v in bn_all if 'moving_variance' in v.name][0] + + copy_bn_params(gamma, beta, mean, var_v, f'{enc_prefix}.downsample_bn') + break + + # ResBlock - this is more complex, skip for now and use name-based fallback + + if verbose: + print(f"\n⚠ Direct mapping is complex due to TF's auto-generated names") + print("Falling back to name-based conversion...") + + # Fall back to the name-based approach + return convert_weights(h5_path, output_path, verbose) + + +def convert_weights(h5_path, output_path, verbose=True): + """Full conversion pipeline from TensorFlow HDF5 to PyTorch .pth. + + Parameters + ---------- + h5_path : str or Path + Path to TensorFlow .h5 weights file + output_path : str or Path + Path to save PyTorch .pth weights file + verbose : bool, optional + Whether to print conversion details (default: True) + + Returns + ------- + OrderedDict + The converted PyTorch state_dict + """ + h5_path = Path(h5_path) + output_path = Path(output_path) + + if not h5_path.exists(): + raise FileNotFoundError(f"Input file not found: {h5_path}") + + if verbose: + print(f"Loading TensorFlow weights from: {h5_path}") + + # Load TensorFlow weights + tf_weights = load_tf_weights_from_h5(h5_path) + + if verbose: + print(f"Loaded {len(tf_weights)} TensorFlow parameters") + + # Convert to PyTorch format + pytorch_state = map_tf_to_pytorch_names(tf_weights) + + if verbose: + print(f"Converted to {len(pytorch_state)} PyTorch parameters") + + # Save PyTorch weights + output_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(pytorch_state, output_path) + + if verbose: + print(f"Saved PyTorch weights to: {output_path}") + print("\nSample converted weights:") + for i, (key, tensor) in enumerate(list(pytorch_state.items())[:5]): + print(f" {key}: {tuple(tensor.shape)}") + if len(pytorch_state) > 5: + print(f" ... and {len(pytorch_state) - 5} more") + + return pytorch_state + + +def verify_conversion(pytorch_state, expected_param_count=128897505): + """Verify that weight conversion was successful. + + Parameters + ---------- + pytorch_state : OrderedDict + Converted PyTorch state dict + expected_param_count : int, optional + Expected number of trainable parameters + + Returns + ------- + bool + True if verification passed + """ + # Count parameters + total_params = sum(p.numel() for p in pytorch_state.values()) + + print(f"\nVerification:") + print(f" Total parameters: {total_params:,}") + print(f" Expected parameters: {expected_param_count:,}") + + # Check for NaN or Inf + has_nan = any(torch.isnan(p).any() for p in pytorch_state.values()) + has_inf = any(torch.isinf(p).any() for p in pytorch_state.values()) + + if has_nan: + print(" βœ— WARNING: Found NaN values in weights!") + return False + if has_inf: + print(" βœ— WARNING: Found Inf values in weights!") + return False + + print(" βœ“ No NaN or Inf values") + + # Check parameter count (allow some tolerance for non-trainable params) + if abs(total_params - expected_param_count) < 100000: # Within 100k params + print(" βœ“ Parameter count matches expected") + return True + else: + print(f" βœ— WARNING: Parameter count mismatch (diff: {abs(total_params - expected_param_count):,})") + return False + + +def main(): + """Command-line interface for weight conversion.""" + parser = argparse.ArgumentParser( + description='Convert TensorFlow .h5 weights to PyTorch .pth format' + ) + parser.add_argument( + 'h5_path', + type=str, + help='Input .h5 file (TensorFlow weights)' + ) + parser.add_argument( + 'output_path', + type=str, + help='Output .pth file (PyTorch weights)' + ) + parser.add_argument( + '--verify', + action='store_true', + help='Verify conversion after completing' + ) + parser.add_argument( + '--quiet', + action='store_true', + help='Reduce verbosity' + ) + + args = parser.parse_args() + + # Convert weights + pytorch_state = convert_weights( + args.h5_path, + args.output_path, + verbose=not args.quiet + ) + + # Verify if requested + if args.verify: + success = verify_conversion(pytorch_state) + if not success: + print("\n⚠ Warning: Verification checks failed!") + return 1 + else: + print("\nβœ“ Conversion successful!") + return 0 + + return 0 + + +if __name__ == '__main__': + exit(main()) diff --git a/src/easymode/segmentation/torch_version/model_torch.py b/src/easymode/segmentation/torch_version/model_torch.py new file mode 100644 index 0000000..05eb171 --- /dev/null +++ b/src/easymode/segmentation/torch_version/model_torch.py @@ -0,0 +1,337 @@ +"""PyTorch implementation of 3D U-Net segmentation model. + +This module provides a PyTorch implementation of the 3D U-Net architecture that is +numerically equivalent to the TensorFlow/Keras implementation in model.py. + +Key architectural details: +- Encoder: 6 blocks with filters [32, 64, 128, 256, 512, 1024] +- Decoder: 5 blocks with skip connections +- ResBlock3D: Two 3x3x3 convolutions + BatchNorm + ReLU + residual connection +- BatchNorm parameters match TensorFlow: eps=1e-3, momentum=0.01 (PyTorch definition) +- Input format: NCDHW (channels-first) +- Output: Single channel probability map [0-1] +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .torch_utils import TFSameConv3d + + +class ResBlock3D(nn.Module): + """3D Residual block with batch normalization and ReLU activation. + + This block implements a residual connection with two 3D convolutions, + each followed by batch normalization. The first conv-BN pair is followed + by ReLU, and a final ReLU is applied after the residual addition. + + Parameters + ---------- + in_channels : int + Number of input channels + out_channels : int + Number of output channels + """ + + def __init__(self, in_channels, out_channels): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + # First convolution path + self.conv1 = nn.Conv3d( + in_channels, out_channels, kernel_size=3, padding=1, bias=False + ) + # BatchNorm parameters matching TensorFlow: + # TF momentum=0.99 (exponential moving average weight for running stats) + # PyTorch momentum=0.01 (1 - TF_momentum, weight for current batch) + # TF epsilon=0.001 (default) + self.bn1 = nn.BatchNorm3d(out_channels, eps=1e-3, momentum=0.01) + self.relu1 = nn.ReLU(inplace=False) # inplace=False for easier debugging + + # Second convolution path + self.conv2 = nn.Conv3d( + out_channels, out_channels, kernel_size=3, padding=1, bias=False + ) + self.bn2 = nn.BatchNorm3d(out_channels, eps=1e-3, momentum=0.01) + + # Skip connection adjustment if input/output channels differ + if in_channels != out_channels: + self.skip_conv = nn.Conv3d( + in_channels, out_channels, kernel_size=1, padding=0, bias=False + ) + self.skip_bn = nn.BatchNorm3d(out_channels, eps=1e-3, momentum=0.01) + else: + self.skip_conv = None + self.skip_bn = None + + def forward(self, x): + """Forward pass through residual block. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (N, C_in, D, H, W) + + Returns + ------- + torch.Tensor + Output tensor of shape (N, C_out, D, H, W) + """ + # Main path + out = self.conv1(x) + out = self.bn1(out) + out = self.relu1(out) + + out = self.conv2(out) + out = self.bn2(out) + + # Skip connection + if self.skip_conv is not None: + skip = self.skip_conv(x) + skip = self.skip_bn(skip) + else: + skip = x + + # Residual addition and final activation + out = out + skip + out = F.relu(out) + + return out + + +class EncoderBlock(nn.Module): + """Encoder block with optional downsampling followed by residual block. + + Parameters + ---------- + in_channels : int + Number of input channels + out_channels : int + Number of output channels + stride : int, optional + Stride for downsampling (default: 1, no downsampling) + """ + + def __init__(self, in_channels, out_channels, stride=1): + super().__init__() + self.stride = stride + + # Optional downsampling + if stride > 1: + # Use TFSameConv3d to match TensorFlow's asymmetric 'same' padding + self.downsample = TFSameConv3d( + in_channels, out_channels, + kernel_size=3, stride=stride, bias=False + ) + self.downsample_bn = nn.BatchNorm3d(out_channels, eps=1e-3, momentum=0.01) + self.downsample_relu = nn.ReLU(inplace=False) + res_in_channels = out_channels + else: + self.downsample = None + res_in_channels = in_channels + + # Residual block + self.res_block = ResBlock3D(res_in_channels, out_channels) + + def forward(self, x): + """Forward pass through encoder block. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (N, C_in, D, H, W) + + Returns + ------- + torch.Tensor + Output tensor of shape (N, C_out, D', H', W') + where D', H', W' = D/stride, H/stride, W/stride if stride > 1 + """ + if self.downsample is not None: + x = self.downsample(x) + x = self.downsample_bn(x) + x = self.downsample_relu(x) + + x = self.res_block(x) + return x + + +class DecoderBlock(nn.Module): + """Decoder block with optional upsampling, skip connection, and residual block. + + Parameters + ---------- + in_channels : int + Number of input channels from previous decoder layer + out_channels : int + Number of output channels + upsample_kernel_size : int, optional + Kernel size for upsampling (default: 2) + If 1, no upsampling is performed + """ + + def __init__(self, in_channels, out_channels, upsample_kernel_size=2): + super().__init__() + self.upsample_kernel_size = upsample_kernel_size + + # Optional upsampling + if upsample_kernel_size > 1: + self.upsample = nn.ConvTranspose3d( + in_channels, out_channels, + kernel_size=upsample_kernel_size, + stride=upsample_kernel_size, + padding=0, # TF uses padding='same' but with stride=kernel_size, this is equiv to padding=0 + bias=False + ) + self.upsample_bn = nn.BatchNorm3d(out_channels, eps=1e-3, momentum=0.01) + self.upsample_relu = nn.ReLU(inplace=False) + # After concatenating with skip connection, input channels = out_channels * 2 + concat_in_channels = out_channels * 2 + else: + self.upsample = None + # After concatenating with skip connection, input channels = in_channels * 2 + concat_in_channels = in_channels * 2 + + # Residual block after concatenation + self.res_block = ResBlock3D(concat_in_channels, out_channels) + + def forward(self, x, skip_connection=None): + """Forward pass through decoder block. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (N, C_in, D, H, W) + skip_connection : torch.Tensor, optional + Skip connection from encoder of shape (N, C_skip, D', H', W') + + Returns + ------- + torch.Tensor + Output tensor of shape (N, C_out, D', H', W') + """ + # Upsample if needed + if self.upsample is not None: + x = self.upsample(x) + x = self.upsample_bn(x) + x = self.upsample_relu(x) + + # Concatenate with skip connection + # PyTorch uses NCDHW format, so concatenate on dim=1 (channel dimension) + if skip_connection is not None: + x = torch.cat([x, skip_connection], dim=1) + + # Apply residual block + x = self.res_block(x) + return x + + +class UNet(nn.Module): + """3D U-Net architecture for volumetric segmentation. + + This implements a 3D U-Net with residual blocks, matching the TensorFlow + implementation. The network consists of: + - 6 encoder blocks with progressively increasing filters + - 5 decoder blocks with skip connections + - Final 1x1 convolution with sigmoid activation + + The architecture uses filters: [32, 64, 128, 256, 512, 1024] + with downsampling strides: [1, 2, 2, 2, 2, 2] + and upsampling kernel sizes: [1, 2, 2, 2, 2, 2] + + Input shape: (N, 1, D, H, W) - single channel volumes + Output shape: (N, 1, D, H, W) - probability maps [0-1] + """ + + def __init__(self): + super().__init__() + + # Architecture configuration (must match TensorFlow exactly) + filters = [32, 64, 128, 256, 512, 1024] + strides = [1, 2, 2, 2, 2, 2] + upsample_kernel_sizes = [1, 2, 2, 2, 2, 2] + + # Encoder blocks + self.encoders = nn.ModuleList() + in_ch = 1 # Single channel input + for i, (f, s) in enumerate(zip(filters, strides)): + self.encoders.append(EncoderBlock(in_ch, f, stride=s)) + in_ch = f + + # Decoder blocks + self.decoders = nn.ModuleList() + decoder_filters = filters[:-1][::-1] # [512, 256, 128, 64, 32] + decoder_upsample = upsample_kernel_sizes[1:][::-1] # [2, 2, 2, 2, 1] + + in_ch = filters[-1] # 1024 from bottleneck + for i, (f, us) in enumerate(zip(decoder_filters, decoder_upsample)): + self.decoders.append(DecoderBlock(in_ch, f, upsample_kernel_size=us)) + in_ch = f + + # Final output layer + self.final_conv = nn.Conv3d(decoder_filters[-1], 1, kernel_size=1) + + def forward(self, x): + """Forward pass through U-Net. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (N, 1, D, H, W) + + Returns + ------- + torch.Tensor + Output probability map of shape (N, 1, D, H, W) + """ + # Encoder path - collect outputs for skip connections + encoder_outputs = [] + for encoder in self.encoders: + x = encoder(x) + encoder_outputs.append(x) + + # Decoder path with skip connections + # Skip connections are from all encoder outputs except the bottleneck, + # in reverse order + skip_connections = encoder_outputs[:-1][::-1] # Reverse, exclude bottleneck + x = encoder_outputs[-1] # Start from bottleneck + + for i, decoder in enumerate(self.decoders): + skip = skip_connections[i] if i < len(skip_connections) else None + x = decoder(x, skip_connection=skip) + + # Final output with sigmoid activation + output = torch.sigmoid(self.final_conv(x)) + + return output + + def count_parameters(self): + """Count total and trainable parameters. + + Returns + ------- + tuple[int, int] + (total_parameters, trainable_parameters) + """ + total = sum(p.numel() for p in self.parameters()) + trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) + return total, trainable + + +def create(): + """Create and return a UNet model instance. + + Returns + ------- + UNet + Initialized U-Net model + + Notes + ----- + This function provides a factory method matching the TensorFlow implementation's + create() function for API compatibility. + """ + model = UNet() + return model diff --git a/src/easymode/segmentation/torch_version/torch_utils.py b/src/easymode/segmentation/torch_version/torch_utils.py new file mode 100644 index 0000000..c9e8ea0 --- /dev/null +++ b/src/easymode/segmentation/torch_version/torch_utils.py @@ -0,0 +1,304 @@ +"""Utility functions for PyTorch segmentation model. + +This module provides helper functions for converting between TensorFlow and PyTorch +tensor formats, as well as other utilities specific to the PyTorch implementation. +""" + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def tf_to_torch_format(volume): + """Convert TensorFlow NDHWC format to PyTorch NCDHW format. + + Parameters + ---------- + volume : numpy.ndarray + Input volume in TensorFlow format: + - 4D: (D, H, W, C) - single volume with channels + - 5D: (N, D, H, W, C) - batch of volumes with channels + + Returns + ------- + torch.Tensor + Output tensor in PyTorch format: + - 4D: (C, D, H, W) - single volume with channels + - 5D: (N, C, D, H, W) - batch of volumes with channels + + Notes + ----- + TensorFlow uses channels-last format (NDHWC), while PyTorch uses + channels-first format (NCDHW). This function handles the conversion. + """ + if volume.ndim == 4: # (D, H, W, C) -> (C, D, H, W) + volume = np.transpose(volume, (3, 0, 1, 2)) + elif volume.ndim == 5: # (N, D, H, W, C) -> (N, C, D, H, W) + volume = np.transpose(volume, (0, 4, 1, 2, 3)) + elif volume.ndim == 3: # (D, H, W) -> (1, D, H, W) + # Single channel volume without explicit channel dimension + volume = np.expand_dims(volume, axis=0) + else: + raise ValueError( + f"Expected 3D, 4D, or 5D array, got shape {volume.shape}" + ) + + return torch.from_numpy(volume).float() + + +def torch_to_tf_format(tensor): + """Convert PyTorch NCDHW format to TensorFlow NDHWC format. + + Parameters + ---------- + tensor : torch.Tensor + Input tensor in PyTorch format: + - 4D: (C, D, H, W) - single volume with channels + - 5D: (N, C, D, H, W) - batch of volumes with channels + + Returns + ------- + numpy.ndarray + Output array in TensorFlow format: + - 4D: (D, H, W, C) - single volume with channels + - 5D: (N, D, H, W, C) - batch of volumes with channels + + Notes + ----- + PyTorch uses channels-first format (NCDHW), while TensorFlow uses + channels-last format (NDHWC). This function handles the conversion. + """ + # Convert to numpy first + if isinstance(tensor, torch.Tensor): + array = tensor.detach().cpu().numpy() + else: + array = tensor + + if array.ndim == 4: # (C, D, H, W) -> (D, H, W, C) + array = np.transpose(array, (1, 2, 3, 0)) + elif array.ndim == 5: # (N, C, D, H, W) -> (N, D, H, W, C) + array = np.transpose(array, (0, 2, 3, 4, 1)) + else: + raise ValueError( + f"Expected 4D or 5D array, got shape {array.shape}" + ) + + return array + + +def prepare_model_for_inference(model, device='cuda'): + """Prepare PyTorch model for inference. + + Parameters + ---------- + model : torch.nn.Module + PyTorch model to prepare + device : str, optional + Device to move model to ('cuda' or 'cpu'), by default 'cuda' + + Returns + ------- + torch.nn.Module + Model in evaluation mode on the specified device + + Notes + ----- + This function: + 1. Sets model to evaluation mode (important for BatchNorm and Dropout) + 2. Moves model to the specified device + """ + model.eval() + + # Check if CUDA is available if requested + if device == 'cuda' and not torch.cuda.is_available(): + print("Warning: CUDA requested but not available, using CPU") + device = 'cpu' + + model = model.to(device) + return model + + +def get_device(gpu_id=None): + """Get PyTorch device for computation. + + Parameters + ---------- + gpu_id : int, optional + GPU device ID to use. If None, uses CUDA if available, else CPU. + + Returns + ------- + torch.device + PyTorch device object + """ + if gpu_id is not None: + if torch.cuda.is_available(): + return torch.device(f'cuda:{gpu_id}') + else: + print(f"Warning: GPU {gpu_id} requested but CUDA not available, using CPU") + return torch.device('cpu') + elif torch.cuda.is_available(): + return torch.device('cuda') + else: + return torch.device('cpu') + + +def clear_gpu_memory(): + """Clear PyTorch GPU memory cache. + + This is useful when running multiple inferences or handling OOM errors. + """ + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def print_gpu_memory_stats(): + """Print current GPU memory usage statistics. + + Useful for debugging memory issues. + """ + if torch.cuda.is_available(): + allocated = torch.cuda.memory_allocated() / 1024**3 + reserved = torch.cuda.memory_reserved() / 1024**3 + print(f"GPU Memory: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved") + else: + print("CUDA not available") + + +def calculate_tf_same_padding_3d(input_shape, kernel_size, stride): + """Calculate asymmetric padding for TensorFlow-style 'same' padding. + + TensorFlow's 'same' padding with stride > 1 uses asymmetric padding, + adding extra padding on the right/bottom when needed. This function + computes the padding tuple needed for F.pad to replicate this behavior. + + Parameters + ---------- + input_shape : tuple of int + Spatial dimensions (D, H, W) of the input tensor + kernel_size : int or tuple of int + Kernel size for the convolution. Can be a single int (applied to all dims) + or a tuple of 3 ints (kD, kH, kW) + stride : int or tuple of int + Stride for the convolution. Can be a single int (applied to all dims) + or a tuple of 3 ints (sD, sH, sW) + + Returns + ------- + tuple of int + Padding tuple for F.pad in the format: + (W_left, W_right, H_left, H_right, D_left, D_right) + + Notes + ----- + TensorFlow 'same' padding formula: + - output_size = ceil(input_size / stride) + - total_padding = max(0, (output_size - 1) * stride + kernel_size - input_size) + - pad_before = total_padding // 2 + - pad_after = total_padding - pad_before + + F.pad expects padding in reverse order: (last_dim, ..., first_dim) + """ + # Ensure kernel_size and stride are tuples + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size, kernel_size) + if isinstance(stride, int): + stride = (stride, stride, stride) + + # Calculate padding for each dimension + paddings = [] + for dim_size, k_size, s_size in zip(input_shape, kernel_size, stride): + # TensorFlow 'same' padding calculation + output_size = (dim_size + s_size - 1) // s_size # Ceiling division + total_padding = max(0, (output_size - 1) * s_size + k_size - dim_size) + pad_before = total_padding // 2 + pad_after = total_padding - pad_before + paddings.append((pad_before, pad_after)) + + # F.pad expects (W_left, W_right, H_left, H_right, D_left, D_right) + # Our paddings are in order (D, H, W), so reverse them + return ( + paddings[2][0], paddings[2][1], # W_left, W_right + paddings[1][0], paddings[1][1], # H_left, H_right + paddings[0][0], paddings[0][1], # D_left, D_right + ) + + +class TFSameConv3d(nn.Module): + """Conv3d with TensorFlow-style 'same' padding. + + This module wraps nn.Conv3d to provide padding behavior that matches + TensorFlow's 'same' padding mode when stride > 1. PyTorch's built-in + padding uses symmetric padding, while TensorFlow uses asymmetric padding + that may add extra padding on the right/bottom sides. + + Parameters + ---------- + in_channels : int + Number of input channels + out_channels : int + Number of output channels + kernel_size : int or tuple of int + Size of the convolving kernel + stride : int or tuple of int, optional + Stride of the convolution (default: 1) + bias : bool, optional + If True, adds a learnable bias to the output (default: False) + + Notes + ----- + This is critical for converting TensorFlow models to PyTorch when stride > 1, + as the different padding strategies can cause significant numerical differences. + + For stride=1, this behaves identically to regular Conv3d with padding='same'. + + Examples + -------- + >>> conv = TFSameConv3d(32, 64, kernel_size=3, stride=2) + >>> x = torch.randn(1, 32, 64, 64, 64) + >>> y = conv(x) + >>> y.shape + torch.Size([1, 64, 32, 32, 32]) + """ + + def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=False): + super().__init__() + self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride) + self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size, kernel_size) + + # Create conv with no padding - we'll add it manually + self.conv = nn.Conv3d( + in_channels, out_channels, + kernel_size=kernel_size, + stride=stride, + padding=0, # No built-in padding + bias=bias + ) + + def forward(self, x): + """Forward pass with TF-style 'same' padding. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (N, C, D, H, W) + + Returns + ------- + torch.Tensor + Output tensor after convolution + """ + # Get spatial dimensions (D, H, W) + D, H, W = x.shape[2:] + + # Calculate padding based on input shape + padding = calculate_tf_same_padding_3d( + (D, H, W), self.kernel_size, self.stride + ) + + # Apply asymmetric padding + x = F.pad(x, padding) + + # Apply convolution (with padding=0 since we already padded) + return self.conv(x) diff --git a/src/easymode/segmentation/torch_version/validate_equivalence.py b/src/easymode/segmentation/torch_version/validate_equivalence.py new file mode 100644 index 0000000..9b4f33b --- /dev/null +++ b/src/easymode/segmentation/torch_version/validate_equivalence.py @@ -0,0 +1,315 @@ +"""Numerical validation of PyTorch vs TensorFlow model equivalence. + +This script validates that the PyTorch implementation produces numerically +equivalent results to the TensorFlow implementation. +""" + +import argparse +from pathlib import Path + +import numpy as np +import tensorflow as tf +import torch + +from easymode.segmentation.model import create as create_tf +from easymode.segmentation.torch_version.model_torch import create as create_torch +from easymode.segmentation.torch_version.torch_utils import tf_to_torch_format, torch_to_tf_format + + +class NumericalValidator: + """Validate numerical equivalence between TensorFlow and PyTorch models. + + Parameters + ---------- + tf_weights_path : str or Path + Path to TensorFlow .h5 weights file + torch_weights_path : str or Path + Path to PyTorch .pth weights file + device : str, optional + PyTorch device ('cuda' or 'cpu'), default 'cpu' + """ + + def __init__(self, tf_weights_path, torch_weights_path, device='cpu'): + self.device = device + + # Load TensorFlow model + self.tf_model = create_tf() + dummy_tf = tf.zeros((1, 64, 64, 64, 1)) + _ = self.tf_model(dummy_tf) + self.tf_model.load_weights(tf_weights_path) + + # Load PyTorch model + self.torch_model = create_torch() + try: + state_dict = torch.load(torch_weights_path, map_location=device) + missing, unexpected = self.torch_model.load_state_dict(state_dict, strict=False) + + if missing: + print(f"⚠ Warning: {len(missing)} missing keys in PyTorch state_dict") + for key in missing[:5]: + print(f" - {key}") + if len(missing) > 5: + print(f" ... and {len(missing) - 5} more") + + if unexpected: + print(f"⚠ Warning: {len(unexpected)} unexpected keys in PyTorch state_dict") + for key in unexpected[:5]: + print(f" - {key}") + if len(unexpected) > 5: + print(f" ... and {len(unexpected) - 5} more") + + except Exception as e: + print(f"Error loading PyTorch weights: {e}") + raise + + self.torch_model.eval() + self.torch_model = self.torch_model.to(device) + + def validate_weights(self): + """Validate that weights were loaded correctly. + + Returns + ------- + bool + True if validation passed + """ + print("\n" + "="*60) + print("WEIGHT VALIDATION") + print("="*60) + + # Compare parameter counts + tf_trainable = sum([np.prod(v.shape) for v in self.tf_model.trainable_variables]) + tf_total = sum([np.prod(v.shape) for v in self.tf_model.variables]) + + torch_trainable = sum(p.numel() for p in self.torch_model.parameters() if p.requires_grad) + torch_total = sum(p.numel() for p in self.torch_model.parameters()) + + print(f"TensorFlow parameters:") + print(f" Trainable: {tf_trainable:,}") + print(f" Total: {tf_total:,}") + + print(f"\nPyTorch parameters:") + print(f" Trainable: {torch_trainable:,}") + print(f" Total: {torch_total:,}") + + if tf_trainable != torch_trainable: + print("\nβœ— FAIL: Trainable parameter count mismatch!") + return False + + print("\nβœ“ PASS: Parameter counts match") + return True + + def validate_forward_pass(self, test_input=None, rtol=1e-4, atol=1e-5): + """Validate forward pass produces same outputs. + + Parameters + ---------- + test_input : np.ndarray, optional + Test input in TF format (D, H, W, 1). If None, random input is used. + rtol : float, optional + Relative tolerance for np.allclose + atol : float, optional + Absolute tolerance for np.allclose + + Returns + ------- + bool + True if validation passed + """ + print("\n" + "="*60) + print("FORWARD PASS VALIDATION") + print("="*60) + + # Create test input + if test_input is None: + # Use small random input for speed + test_input = np.random.randn(64, 128, 128, 1).astype(np.float32) + + print(f"Test input shape (TF format): {test_input.shape}") + + # TensorFlow inference + tf_input = np.expand_dims(test_input, 0) # Add batch dim: (1, D, H, W, 1) + tf_output = self.tf_model(tf_input, training=False).numpy() + + print(f"TensorFlow output shape: {tf_output.shape}") + print(f"TensorFlow output range: [{tf_output.min():.6f}, {tf_output.max():.6f}]") + + # PyTorch inference + torch_input = tf_to_torch_format(test_input) # (1, D, H, W) + if torch_input.ndim == 4: + torch_input = torch_input.unsqueeze(0) # Add batch: (1, 1, D, H, W) + torch_input = torch_input.to(self.device) + + with torch.no_grad(): + torch_output = self.torch_model(torch_input) + + torch_output_np = torch_to_tf_format(torch_output.cpu()) # Back to TF format + + print(f"PyTorch output shape: {torch_output_np.shape}") + print(f"PyTorch output range: [{torch_output_np.min():.6f}, {torch_output_np.max():.6f}]") + + # Compare outputs + abs_diff = np.abs(tf_output - torch_output_np) + rel_diff = abs_diff / (np.abs(tf_output) + 1e-8) + + max_abs_diff = np.max(abs_diff) + mean_abs_diff = np.mean(abs_diff) + max_rel_diff = np.max(rel_diff) + mean_rel_diff = np.mean(rel_diff) + + print(f"\nNumerical Differences:") + print(f" Max absolute difference: {max_abs_diff:.2e}") + print(f" Mean absolute difference: {mean_abs_diff:.2e}") + print(f" Max relative difference: {max_rel_diff:.2e}") + print(f" Mean relative difference: {mean_rel_diff:.2e}") + + # Check tolerance + passed = np.allclose(tf_output, torch_output_np, rtol=rtol, atol=atol) + + if passed: + print(f"\nβœ“ PASS: Outputs match within tolerance (rtol={rtol}, atol={atol})") + else: + print(f"\nβœ— FAIL: Outputs exceed tolerance!") + + # Find worst mismatches + flat_idx = np.argmax(abs_diff) + idx = np.unravel_index(flat_idx, abs_diff.shape) + print(f"\n Worst mismatch at index {idx}:") + print(f" TensorFlow: {tf_output[idx]:.8f}") + print(f" PyTorch: {torch_output_np[idx]:.8f}") + print(f" Difference: {abs_diff[idx]:.8f}") + + # Show distribution of differences + percentiles = [50, 90, 95, 99] + print(f"\n Absolute difference percentiles:") + for p in percentiles: + val = np.percentile(abs_diff, p) + print(f" {p}th: {val:.2e}") + + return passed + + def validate_multiple_inputs(self, num_tests=5, rtol=1e-4, atol=1e-5): + """Validate on multiple random inputs. + + Parameters + ---------- + num_tests : int, optional + Number of random inputs to test + rtol : float, optional + Relative tolerance + atol : float, optional + Absolute tolerance + + Returns + ------- + bool + True if all tests passed + """ + print("\n" + "="*60) + print(f"MULTIPLE INPUT VALIDATION ({num_tests} tests)") + print("="*60) + + all_passed = True + for i in range(num_tests): + print(f"\nTest {i+1}/{num_tests}:") + + # Different sizes to test + sizes = [(64, 64, 64, 1), (128, 128, 128, 1)] + size = sizes[i % len(sizes)] + + test_input = np.random.randn(*size).astype(np.float32) + passed = self.validate_forward_pass(test_input, rtol=rtol, atol=atol) + + if not passed: + all_passed = False + + if all_passed: + print(f"\nβœ“ ALL TESTS PASSED ({num_tests}/{num_tests})") + else: + print(f"\nβœ— SOME TESTS FAILED") + + return all_passed + + +def run_validation(tf_weights, torch_weights, num_tests=5): + """Main validation entry point. + + Parameters + ---------- + tf_weights : str or Path + Path to TensorFlow .h5 weights file + torch_weights : str or Path + Path to PyTorch .pth weights file + num_tests : int, optional + Number of random input tests to run + + Returns + ------- + bool + True if all validations passed + """ + print("="*60) + print("PyTorch vs TensorFlow Model Equivalence Validation") + print("="*60) + + validator = NumericalValidator(tf_weights, torch_weights, device='cpu') + + # 1. Validate weights + weights_ok = validator.validate_weights() + + # 2. Validate multiple inputs + inputs_ok = validator.validate_multiple_inputs(num_tests=num_tests) + + # Summary + print("\n" + "="*60) + print("VALIDATION SUMMARY") + print("="*60) + print(f"Weight validation: {'βœ“ PASSED' if weights_ok else 'βœ— FAILED'}") + print(f"Forward pass tests: {'βœ“ PASSED' if inputs_ok else 'βœ— FAILED'}") + + all_ok = weights_ok and inputs_ok + if all_ok: + print("\nπŸŽ‰ ALL VALIDATIONS PASSED!") + else: + print("\n⚠ VALIDATION FAILED - weights may need adjustment") + + return all_ok + + +def main(): + """Command-line interface for validation.""" + parser = argparse.ArgumentParser( + description='Validate TensorFlow vs PyTorch model equivalence' + ) + parser.add_argument( + '--tf_weights', + type=str, + required=True, + help='TensorFlow .h5 weights file' + ) + parser.add_argument( + '--torch_weights', + type=str, + required=True, + help='PyTorch .pth weights file' + ) + parser.add_argument( + '--num_tests', + type=int, + default=5, + help='Number of random input tests (default: 5)' + ) + + args = parser.parse_args() + + success = run_validation( + args.tf_weights, + args.torch_weights, + num_tests=args.num_tests + ) + + return 0 if success else 1 + + +if __name__ == '__main__': + exit(main())