UCell

A small cell segmentation model with not so small generalizability.

UCell is a ~15M-parameter recursive transformer for instance segmentation of cells and nuclei in microscopy images. Rather than stacking layers, it applies the same two-block transformer repeatedly to refine an internal representation, so depth comes from recursion instead of from parameters.

• Code: https://github.com/jiyuuchc/ucell
• Paper: https://doi.org/10.48550/arXiv.2604.00243
• License: MIT

Checkpoints

file	parameters	size	hidden size
ucell-768.pt	14.7M	56 MB	768
ucell-1024.pt	26.4M	101 MB	1024

Installation

git clone https://github.com/jiyuuchc/ucell.git
cd ucell
python -m venv .venv && source .venv/bin/activate
pip install -e .
pip install huggingface_hub

Usage

import numpy as np
import tifffile
import torch
from huggingface_hub import hf_hub_download

from ucell.dynamics import compute_masks
from ucell.frm import FRMWrapper
from ucell.utils import pad_channel, patcherize

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

weights = hf_hub_download("jiyuuchc/ucell", "ucell-768.pt")
model = FRMWrapper.from_checkpoint(weights).eval().to(DEVICE)

# Any 2D or channel-first multichannel array, scaled to [0, 1] and padded to
# the three channels the patch embedding expects.
img = tifffile.imread("sample.tif").astype("float32")
img = pad_channel(img / (img.max() + 1e-5))

# patcherize tiles into 256x256 patches and stitches the result back, so any
# image size works.  The trailing 0 is the task id.
with torch.device(DEVICE):
    out = patcherize(model.inner.predict, GS=model.config.image_size)(img, 0)

flow, cell_prob = np.moveaxis(out[..., :2], -1, 0), out[..., 2]

mask = compute_masks(
    flow * 4.0,
    cell_prob,
    cellprob_threshold=-0.5,
    min_size=5,
    device=torch.device(DEVICE),
)

tifffile.imwrite("sample_mask.tif", mask.astype("uint16"))
print(f"{mask.max()} instances")

mask is a uint16 array of instance labels: 0 is background, each cell gets its own integer.

Inputs

Images are expected as float arrays scaled to [0, 1] and padded to three channels by pad_channel. Grayscale and channel-first multichannel inputs are both fine. There is no fixed input size — patcherize tiles the image.

Parameters

• cellprob_threshold (default -0.5) — lower detects more cells.
• min_size (default 5) — drops small fragments.
• recursion depth, via FRMWrapper.from_checkpoint(weights, overrides={"model.L_cycles": 7}). Checkpoints ship at 21; 7 runs three times shallower and faster, at some cost in accuracy. overrides takes any dotted config path.

Hardware

A GPU is optional. Measured on one 383x512 two-channel image:

checkpoint	peak GPU memory	GPU (L40S)	CPU
ucell-768.pt	376 MiB	0.2 s	12 s
ucell-1024.pt	525 MiB	0.2 s	—

Inference is tiled, so peak memory is set by the patch count of a single image rather than by its total size. Any CUDA GPU with ~2 GB free is ample.

Training data

Trained on a mixture of public cell and nucleus segmentation datasets: CellPose (and cyto2), LIVECell, the NeurIPS 2022 Cell Segmentation Challenge, TissueNet and MoNuSAC.

Fine-tuning

LoRA fine-tuning on a small labelled set is supported from the repository:

python train.py \
  --init ucell-768.pt \
  --config config.py:train \
  --config.data_dir=${DATADIR} \
  --config.n_iters=1 \
  --config.epochs_per_iter=1024 \
  --config.ema_decay=0.95 \
  --config.lora.rank=16

Save training data (*.tif, *_label.tif) under ${DATADIR}/train.

Limitations

• Trained on 2D images; there is no 3D model.
• Performance degrades on modalities far from the training mixture. Very sparse, very large cells are a known weak point.

Citation

@article{ucell,
  title         = {UCell: rethinking generalizability and scaling of bio-medical vision models},
  author        = {Kuang, Nicholas and Scalon, Vanessa and Yu, Ji},
  year          = {2026},
  eprint        = {2604.00243},
  archivePrefix = {arXiv},
  doi           = {10.48550/arXiv.2604.00243},
}