cellmap/nuc-resunet-setup-61

nuc_resunet_setup_61

Nuclei (nuc) from FIB-SEM at 256 nm, 3D residual U-Net (InstanceNorm, valid convs), exp_mito setup_61 at 20000 iterations. Valid convolutions (88 voxels of context per axis); outputs raw logits. ONNX web version: cellmap/onnx-web-models/setup_61.

Model Details

Architecture	ResUNet3DValid
Framework	torch
Spatial Dims	3
Input Channels	1
Output Channels	1
Channel Names	nuc
Iteration	20000
Input Voxel Size	256, 256, 256 nm
Output Voxel Size	256, 256, 256 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/nuc-resunet-setup-61")
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()