| --- |
| library_name: cellmap-models |
| tags: |
| - pytorch |
| - onnx |
| - torchscript |
| - 3d |
| - segmentation |
| - electron-microscopy |
| - cellmap |
| - er |
| license: bsd-3-clause |
| --- |
| |
| # er_resunet_setup_67 |
| |
| Endoplasmic reticulum (er) from FIB-SEM at 8 nm, 3D residual U-Net (InstanceNorm, valid convs), exp_mito setup_67 at 20000 iterations. Valid convolutions (88 voxels of context per axis); outputs raw logits. ONNX web version: cellmap/onnx-web-models/setup_67. |
|
|
| ## Model Details |
|
|
| | | | |
| |---|---| |
| | **Architecture** | ResUNet3DValid | |
| | **Framework** | torch | |
| | **Spatial Dims** | 3 | |
| | **Input Channels** | 1 | |
| | **Output Channels** | 1 | |
| | **Channel Names** | er | |
| | **Iteration** | 20000 | |
| | **Input Voxel Size** | 8, 8, 8 nm | |
| | **Output Voxel Size** | 8, 8, 8 nm | |
| | **Inference Input Shape** | 300, 300, 300 | |
| | **Inference Output Shape** | 212, 212, 212 | |
| | **Input Normalization** | {"max": 1.0, "method": "min_max", "min": -1.0} | |
| | **Inference Output Activation** | {"apply_after_model": true, "name": "sigmoid"} | |
| |
| ## Available Formats |
| |
| | File | Format | Usage | |
| |---|---|---| |
| | `model.pt2` | torch.export program | `torch.export.load("model.pt2").module()` | |
| | `model_state_dict.pt` + `model_def.py` | PyTorch state_dict + module code | see "PyTorch original" below | |
| | `model.ts` | TorchScript | `torch.jit.load("model.ts")` | |
| | `model.onnx` (+ `model.onnx.data`) | ONNX | `onnxruntime.InferenceSession("model.onnx")` | |
| | `metadata.json` | JSON | Model metadata | |
| | `torch_source.json` | JSON | class + constructor arguments for `model_def.py` | |
| |
| ## Usage |
| |
| ```bash |
| pip install cellmap-models |
| ``` |
| |
| ```python |
| from cellmap_models.model_export.cellmap_model import CellmapModel |
|
|
| model = CellmapModel("path/to/model/folder") |
|
|
| # Inference |
| output = model.ts_model(input_tensor) |
|
|
| # Finetuning |
| trainable_model = model.train() |
| ``` |
| |
| Or download from this repo and load directly: |
| |
| ```python |
| from huggingface_hub import snapshot_download |
| from cellmap_models.model_export.cellmap_model import CellmapModel |
|
|
| path = snapshot_download(repo_id="cellmap/er-resunet-setup-67") |
| model = CellmapModel(path) |
| ``` |
| |
| ## Author |
| |
| Marwan Zouinkhi |
| |
| ## Links |
| |
| - [cellmap-models](https://github.com/janelia-cellmap/cellmap-models) |
| - [CellMap Project](https://www.janelia.org/project-team/cellmap) |
| |
| ## PyTorch original (for finetuning / re-export) |
| |
| `model_state_dict.pt` (plain tensors, `weights_only=True`) + the module code (`model_def.py`); `torch_source.json` |
| names the class and constructor arguments. |
| |
| ```python |
| import json, importlib.util, torch |
| src = json.load(open("torch_source.json")) |
| spec = importlib.util.spec_from_file_location("model_def", src["module"]); mod = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(mod) |
| model = getattr(mod, src["class"])(**src["kwargs"]) |
| model.load_state_dict(torch.load("model_state_dict.pt", weights_only=True)) |
| model.eval() |
| ``` |
| |
