cellmap/er-resunet-setup-67

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

pip install cellmap-models

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:

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
• CellMap Project

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.

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()