cellmap/mito-aff-unetiso-setup-53

mito_aff_unetiso_setup_53

Mitochondria affinities (mito_aff_z, mito_aff_y, mito_aff_x) from FIB-SEM at 16 nm, isotropic 3D U-Net (BatchNorm, MaxPool, ConvTranspose), exp_mito setup_53 at 340000 iterations. Valid convolutions (88 voxels of context per axis); outputs raw logits. ONNX web version: cellmap/onnx-web-models/setup_53.

Model Details

Architecture	Unet3dIso
Framework	torch
Spatial Dims	3
Input Channels	1
Output Channels	3
Channel Names	mito_aff_z, mito_aff_y, mito_aff_x
Iteration	340000
Input Voxel Size	16, 16, 16 nm
Output Voxel Size	16, 16, 16 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.pt	PyTorch pickle	torch.load("model.pt")
model.ts	TorchScript	torch.jit.load("model.ts")
model.onnx	ONNX	onnxruntime.InferenceSession("model.onnx")
metadata.json	JSON	Model metadata

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="mito_aff_unetiso_setup_53")
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()