omniem 0.1.0__tar.gz

omniem-0.1.0/.gitignore

@@ -0,0 +1,33 @@
+/weights/
+/configs/
+/out/
+
+*.pyc
+*.txt
+# requirements-dev.txt is tracked dev tooling, NOT a runtime data file — it must NOT
+# match the broad `*.txt` rule above (otherwise a fresh clone has no dev-deps file).
+!requirements-dev.txt
+*.pth
+*~
+
+build/
+dist/
+*.egg-info/
+**/__pycache__/
+
+# Generated API docs — pdoc HTML built by scripts/build_docs.sh. A build
+# artifact regenerated from the docstrings on demand; never committed.
+/docs/api/
+
+# Codex-review fix #5: local agent / tooling state. These dirs contain user-/
+# machine-specific settings (Claude Code permissions, Codex profile) and must not
+# be committed.
+.claude/
+.codex/
+
+# Editor / pytest / ruff caches
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.vscode/
+.idea/

omniem-0.1.0/CHANGELOG.md

@@ -0,0 +1,48 @@
+# Changelog
+
+
+## [0.1.0]
+
+
+### Added
+
+- **Encoder.** `EMEncoder` — `load(arch, weights)` (raw `vit.*` checkpoint, loaded
+  directly), single-shot `forward` → CLS / patch / inner-block features,
+  `apply_input` (split-out input transform), and `name_parameter_group`. The
+  owner-frozen encoder arch catalog: `list_encoders` / `arch_info` (`emdinov1`).
+- **Model.** `OmniEM` (EM-DINO encoder + STAdapter z-fusion + UNETR decoder) —
+  `load` / `from_config` (optional, separable weight loading: merged, or
+  encoder-/head-only, or none → random init), `predict` (single-shot forward →
+  **pure logits** at the caller's shape), `apply_input`, the `task_type`-gated output
+  stage `apply_output` (`image2image` → sigmoid+uint image; `image2label` → argmax
+  label map), `save_weights` (merged or backbone+head split), and `prepare_train`
+  (training handoff: unfreeze, optionally fix the encoder backbone). The owner-frozen
+  model arch catalog: `list_models` / `model_arch_info` (`omniemv1`).
+- **Shared-encoder borrow.** `OmniEM.load` / `from_config` accept `encoder=` (a
+  pre-built `EMEncoder`) to share one ViT backbone by reference across many heads
+  (memory-efficient). Head-only load; borrowed models are read-only (whole-model
+  mutators rejected) so the shared encoder is never mutated.
+- **Input conform round-trip.** `predict` / `apply_input` accept
+  `conform={'strict','pad','resize'}` so non-square / non-stride-multiple XY is handled
+  gracefully and the output is round-tripped to the original shape.
+- **Output-size control (CLI super-resolution).** `omniem infer --output-scale F`
+  bicubic-resizes the input XY by `F` before inference; since the model is
+  shape-preserving, the output lands at the scaled size (`F>1` super-resolution,
+  `F<1` quick-inference). XY only — Z is never resized; 3D (`zyx`) inputs warn
+  (anisotropy / no Z alignment) and still run. Orthogonal to `--conform`; CLI-only
+  (the Python API resizes the input directly — see `docs/api.md`). The resize is
+  recorded in the infer sidecar (`output_scale: {factor, input_yx, scaled_yx}`).
+- **CLI** (`omniem`): `list-encoders`, `list-models`, `features` (single-shot encoder
+  feature extraction), `infer` (single-shot model inference with `--weights` merged or
+  `--backbone`/`--head` split, `--conform`, `--output-scale`,
+  `--scale`/`--unit-range`/`--norm`/`--mean`/`--std`, `--out-dtype`, `--save-logits`),
+  and `split` / `merge` (weight-file utilities: split a merged `.pt` into a
+  `--backbone` + `--head` pair, or merge a pair back — the boundary is the net's
+  derived encoder prefix, not a hardcoded `vit.`). Each `features`/`infer` run writes a
+  store + a JSON reproducibility sidecar.
+- **Typed error taxonomy** under `OmniEMError`: `ConfigError`, `WeightFormatError`,
+  `MissingExtraError`, `InputContractError`, `OOMError`.
+- **Config.** `omniem.config.ModelConfig` (+ `BaseConfig`) with YAML I/O and a
+  schema-version policy.
+
+

omniem-0.1.0/LICENSE

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

omniem-0.1.0/PKG-INFO

@@ -0,0 +1,307 @@
+Metadata-Version: 2.4
+Name: omniem
+Version: 0.1.0
+Summary: Inference and utilities for EM-specific encoders and OmniEM models.
+Author-email: Liuyuan He <liyhe@pku.edu.cn>
+License: MIT
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: monai<2,>=1.2
+Requires-Dist: numpy>=1.24
+Requires-Dist: pydantic>=2.0
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: tifffile>=2023.1.0
+Requires-Dist: torch>=2.3
+Description-Content-Type: text/markdown
+
+# omniem
+
+`omniem` is a GUI-free Python package for electron microscopy (EM) image
+workflows, introduced from [EM-SSL project](https://github.com/pku-maleilab/EM-SSL-project). It provides two main capabilities:
+
+- **Run OmniEM models** for single-shot segmentation or restoration.
+- **Run EM-DINO encoders** to extract CLS, patch, or inner-block features from
+  EM images.
+
+Downstream tools build on the same public API: the
+[`omniem-train`](https://github.com/pku-maleilab/omniem-train) training pipeline
+and the [napari-omniem](https://github.com/pku-maleilab/napari-omniem) GUI plugin.
+
+## Contents
+
+- [Install](#install)
+- [Main Features](#main-features)
+- [Model Config YAML](#model-config-yaml)
+- [First Commands](#first-commands)
+- [Full Guides](#full-guides)
+- [Related Projects](#related-projects)
+- [Future Features](#future-features)
+- [License](#license)
+
+## Install
+
+`omniem` requires **Python >= 3.10**.
+
+For inference and feature extraction, CUDA is recommended when you have a
+supported NVIDIA GPU. Install the PyTorch build that matches your CUDA driver /
+runtime first; use the selector in the
+[PyTorch install guide](https://pytorch.org/get-started/locally/) for the exact
+command for your machine.
+
+Then install `omniem` from PyPI:
+
+```bash
+pip install omniem
+```
+
+Or clone the package repository and install it locally:
+
+```bash
+git clone https://github.com/pku-maleilab/omniem-package.git
+cd omniem-package
+pip install .
+```
+
+Core runtime dependencies include
+[PyTorch](https://pytorch.org/), [NumPy](https://numpy.org/),
+[tifffile](https://github.com/cgohlke/tifffile),
+[Pydantic](https://docs.pydantic.dev/), [PyYAML](https://pyyaml.org/), and
+[MONAI](https://monai.io/).
+
+## Main Features
+
+| Feature | Use it when | Main CLI | Main Python API |
+|---|---|---|---|
+| Model inference | you have a model config plus model weights and want segmentation, restoration, or raw logits | `omniem infer` | `OmniEM.load(...)`, `model.predict(...)`, `model.apply_output(...)` |
+| Encoder features | you only need EM-DINO backbone features, without a model head | `omniem features` | `EMEncoder.load(...)`, `enc(...)` |
+
+### Common Concepts
+
+#### Model = Config + Weights
+
+An OmniEM model is fully specified by a model config YAML plus model weights.
+The config describes how to build the head and interpret its output: model
+architecture, encoder architecture, 2D/3D shape, output channels, `task_type`,
+and the fixed training `mean`/`std` in `[0, 1]` image space.
+
+Weights are plain PyTorch `state_dict` files. They may be split into a shared
+EM-DINO backbone file plus a head file, or stored as one merged whole-model file.
+Split weights are useful when several heads share one encoder backbone. Merged
+weights are convenient when you want one standalone model file.
+
+### Available Models
+
+Model files are distributed outside the Python wheel. Download config YAML files
+from [here](https://drive.google.com/drive/folders/1cFPBmozY5VAh8ZgSe16U7ydX9RMmvbzu?usp=drive_link). Download backbone and head weight files from [here](https://drive.google.com/drive/folders/1vpzVk6vDui8Aj34FdTMfJpXbt5wlMsx_?usp=drive_link).
+
+#### Encoder
+
+Use an encoder when you only need the EM-DINO backbone output, without an
+OmniEM head or model config. The encoder converts an EM image into feature
+tensors that downstream code can reuse:
+
+- `cls`: one global feature vector for the image;
+- `patch`: a grid of local patch features;
+- `inner`: optional intermediate block features.
+
+For a 2D image, the encoder extracts features from that single XY tile. For a
+3D volume, each XY slice is encoded with the same backbone, and the resulting
+features are kept alongside the z-axis so downstream code can relate features
+back to their original slices.
+
+Available encoder models:
+
+| Encoder arch | Description | Default norm | Input stride | Weights |
+|---|---|---|---|---|
+| `emdinov1` | EM-DINOv2 ViT-L/14, EM-domain pretrained encoder | mean `0.595446`, std `0.211906` in `[0, 1]` image space | 14 | `backbone_emdino_v1.pt` (bare `vit.*` checkpoint) |
+
+#### OmniEM
+
+Use an OmniEM model when you have a config YAML, model weights, and a 2D or 3D
+EM image. The model returns raw logits internally; the config controls whether
+`omniem` also applies a canonical output transform.
+
+Available OmniEM models:
+
+| Model | Purpose | Training on | Input | Weights | Config YAML |
+|---|---|---|---|---|---|
+| `mito-seg-ViT-L-2D` | mitochondria segmentation (2D) | MitoLab dataset | 2D EM tile | `backbone_emdino_v1.pt` + `head_mito-seg-ViT-L-2D.pt` | `model_mito-seg-ViT-L-2D.yaml` |
+| `mito-seg-ViT-L-3D` | mitochondria segmentation (3D) | MitoEM-R | 3D subvolume (z >= 16) | `backbone_emdino_v1.pt` + `head_mito-seg-ViT-L-3D.pt` | `model_mito-seg-ViT-L-3D.yaml` |
+| `denoise-emdiffuse-l` | image denoise | Low-level denoise EMDiffuse | 2D EM tile | `backbone_emdino_v1.pt` + `head_denoise-emdiffuse-l.pt` | `model_denoise-emdiffuse-l.yaml` |
+| `superreso-emdiffuse-l` | image super-resolution | Low-level superresolution EMDiffuse | 2D EM tile | `backbone_emdino_v1.pt` + `head_superreso-emdiffuse-l.pt` | `model_superreso-emdiffuse-l.yaml` |
+
+## Model Config YAML
+
+A model config tells `OmniEM` how to build the model head and how to interpret
+outputs.
+
+```yaml
+arch: omniemv1
+encoder: emdinov1
+img_z: 1
+out_channels: 2
+kernel3d_z: null
+task_type: image2label
+resize4emdino: false
+mean: 0.5333333333333333
+std: 0.23137254901960785
+```
+
+Field guide:
+
+| Field | Meaning |
+|---|---|
+| `arch` | model architecture; see `omniem list-models` |
+| `encoder` | encoder architecture; see `omniem list-encoders` |
+| `img_z` | `1` for 2D heads; `>1` for 3D heads |
+| `out_channels` | model output channels |
+| `kernel3d_z` | z-kernel for 3D heads; usually `null` for 2D |
+| `task_type` | `image2label`, `image2image`, or `null` |
+| `resize4emdino` | whether the model uses resize-to-encoder-grid behavior |
+| `mean`, `std` | fixed training normalization for this head |
+
+`task_type` controls the canonical output transform:
+
+| `task_type` | Meaning | Output transform |
+|---|---|---|
+| `image2label` | segmentation / labels | `argmax` over channels |
+| `image2image` | restoration / denoise | `sigmoid`, clamp to `[0, 1]`, scale to uint |
+| omitted / `null` | model has no output opinion | raw float logits only |
+
+For a denoise/restoration head, `out_channels` is usually `1` and
+`task_type: image2image`. For segmentation, `out_channels` is the number of
+classes and `task_type: image2label`.
+
+## First Commands
+
+### Get the example inputs, configs, and weights
+
+The commands below read from three local folders. None of them ship inside the
+pip wheel, so gather them once before running anything:
+
+| Folder | What it holds | How to get it |
+|---|---|---|
+| `examples/` | small example EM images (`.tif`) | tracked in the repo (see below) |
+| `configs/` | model config YAMLs | Google Drive (see [Available Models](#available-models)) |
+| `weights/` | backbone + head weight files | Google Drive (see [Available Models](#available-models)) |
+
+**`examples/`** — if you installed by `git clone`, the example images are already
+in `examples/`. If you installed with `pip`, download them into a local
+`examples/` folder:
+
+```bash
+mkdir -p examples
+BASE=https://raw.githubusercontent.com/pku-maleilab/omniem-package/main/examples
+curl -L -o examples/2d_MitoEM_H_0_0_0.tif       "$BASE/2d_MitoEM_H_0_0_0.tif"
+curl -L -o examples/3d_AxonEM-H-0-0-0_0_0_0.tif "$BASE/3d_AxonEM-H-0-0-0_0_0_0.tif"
+curl -L -o "examples/gly-z=0.tif"               "$BASE/gly-z=0.tif"
+```
+
+**`configs/` and `weights/`** — these are distributed outside the wheel. Download
+the model config YAMLs and the backbone/head weight files from the Google Drive
+links in [Available Models](#available-models), then place them in local
+`configs/` and `weights/` folders so the paths below resolve:
+
+```text
+configs/   model_*.yaml         (config YAMLs)
+weights/   backbone_emdino_v1.pt, head_*.pt   (weight files)
+```
+
+Run the commands from the directory that contains these `examples/`, `configs/`,
+and `weights/` folders.
+
+### Run a model
+
+Run model inference from the CLI:
+
+```bash
+omniem infer \
+  -i examples/2d_MitoEM_H_0_0_0.tif \
+  -m configs/model_mito-seg-ViT-L-2D.yaml \
+  --backbone weights/backbone_emdino_v1.pt \
+  --head weights/head_mito-seg-ViT-L-2D.pt \
+  -o out/mito_labels.tif
+```
+
+Run the same model from Python:
+
+```python
+import numpy as np
+import tifffile
+import torch
+from omniem import OmniEM
+
+model = OmniEM.load(
+    "configs/model_mito-seg-ViT-L-2D.yaml",
+    backbone="weights/backbone_emdino_v1.pt",
+    head="weights/head_mito-seg-ViT-L-2D.pt",
+)
+
+img = tifffile.imread("examples/2d_MitoEM_H_0_0_0.tif")
+x = torch.from_numpy(img.astype(np.float32) / 255.0)
+logits = model.predict(x, axes="yx")
+labels = model.apply_output(logits, axes="yx", dtype="uint8")
+```
+
+### Output-size control (super-resolution)
+
+OmniEM models are shape-preserving (output XY == input XY). To get a larger
+output, for example super-resolution, resize the input up first with
+`--output-scale F`; the model then returns its output at the scaled size
+(`F > 1` upscales, `F < 1` is a quick-inference speed trade-off). It is XY-only
+(Z is never resized; 3D volumes warn) and orthogonal to `--conform`:
+
+```bash
+omniem infer \
+  -i examples/2d_MitoEM_H_0_0_0.tif \
+  -m configs/model_superreso-emdiffuse-l.yaml \
+  --backbone weights/backbone_emdino_v1.pt \
+  --head weights/head_superreso-emdiffuse-l.pt \
+  --output-scale 1.5 \
+  -o out/mito_1.5x.tif
+```
+
+### Split or merge weight files
+
+Convert between a merged whole-model `.pt` and a `backbone` + `head` pair. The
+boundary is the net's derived encoder prefix, so it is correct for any encoder.
+
+```bash
+# merged -> split pair
+omniem split -m configs/model_mito-seg-ViT-L-2D.yaml \
+  -i weights/merged_mito-seg.pt \
+  --backbone weights/backbone_emdino_v1.pt --head weights/head_mito-seg-ViT-L-2D.pt
+
+# split pair -> merged
+omniem merge -m configs/model_mito-seg-ViT-L-2D.yaml \
+  --backbone weights/backbone_emdino_v1.pt --head weights/head_mito-seg-ViT-L-2D.pt \
+  -o weights/merged_mito-seg.pt
+```
+
+## Full Guides
+
+- [CLI guide](docs/cli.md): all `omniem infer`, `omniem features`, `omniem split`,
+  and `omniem merge` options, with command examples.
+- [Python API guide](docs/api.md): `OmniEM`, `EMEncoder`, shared encoders,
+  lower-level calls, weight saving, errors, and API-doc generation.
+
+## Related Projects
+
+- [omniem-train](https://github.com/pku-maleilab/omniem-train): the recommended
+  training pipeline for OmniEM heads; it builds on this package's public API.
+- [napari-omniem](https://github.com/pku-maleilab/napari-omniem): a napari GUI
+  plugin for interactive OmniEM inference.
+
+## Future Features
+
+The current package focuses on the core model/encoder surface. These features are
+planned for later releases:
+
+- large-image tiling and blending (`Inferer`);
+- volume streaming and hdf5/zarr/n5 IO;
+- feature-export orchestration (`Exporter`);
+- install extras such as `[infer]`, `[volume]`, and `[full]`.
+
+## License
+
+[MIT](LICENSE).