cellmap/mito-aff-unetiso-setup-53-2d

mito_aff_unetiso_setup_53_2d

Mitochondria affinities (mito_aff_z, mito_aff_y, mito_aff_x) from FIB-SEM at 16 nm, one z-slice at a time: 2D student of setup_53 (isotropic 3D U-Net (BatchNorm, MaxPool, ConvTranspose), every 3D layer made 2D), distilled for 60000 steps (train_2d_distill.py). Input [N,1,1,H,W], 88 pixels of context per axis; outputs raw logits. Teacher: cellmap/mito-aff-unetiso-setup-53. ONNX web version: cellmap/onnx-web-models/setup_53_2d.

Model Details

Architecture	Unet3dIso2d
Framework	torch
Spatial Dims	3
Input Channels	1
Output Channels	3
Channel Names	mito_aff_z, mito_aff_y, mito_aff_x
Iteration	60000
Input Voxel Size	16, 16, 16 nm
Output Voxel Size	16, 16, 16 nm
Inference Input Shape	1, 556, 556
Inference Output Shape	1, 468, 468
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_2d")
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 + student_2d.py); torch_source.json names the class and constructor arguments.

import json, importlib.util, torch
from student_2d import build_model   # 3D class from model_def.py with its 3D layers made 2D
model = build_model()
model.load_state_dict(torch.load("model_state_dict.pt", weights_only=True))
model.eval()