Hugging Face's logo

xtxx
/
CellFM

Image Feature Extraction
PyTorch
English
histopathology
cell-embedding
vision-transformer
hoverfast
feature-extraction
Model card Files
xet
 Community
CellFM / README.md
xtxx's picture
xtxx
Use English README and remove model-training source descriptions
611e564 verified
preview code
|
Raw History Blame Contribute Delete
13.5 kB
---
pipeline_tag: image-feature-extraction
library_name: pytorch
tags:
- histopathology
- cell-embedding
- vision-transformer
- hoverfast
- feature-extraction
language:
- en
extra_gated_prompt: "Model downloads require manual approval by the repository owner."
---
# CellFM: Cell-Level Embedding Inference
This repository provides the **150,000-step EMA checkpoint of `ds1-vit-base-p4-loop-marb-150k`** and standalone PyTorch code for extracting a **768-dimensional embedding for each cell** from H&E images and HoverFast nuclear contours.
The workflow is **H&E ROI or WSI → HoverFast instance segmentation → CellFM patch tokens → pooling over each nuclear mask → one embedding per nucleus**. Nuclear regions serve as the spatial support for cell representations. HoverFast supplies the segmentation; CellFM extracts features. Segmentation masks are used only to pool output tokens and are not inputs to CellFM attention.
The model page is public. Downloads require a Hugging Face account, an access request, and **manual approval by the repository owner**. After approval, download with your own account token. HoverFast weights are obtained separately from its official project.
## 1. Installation and Download
Use Python 3.10 or later. Install [PyTorch](https://pytorch.org/get-started/locally/) for your hardware, then install the repository dependencies. Use your existing Hugging Face login or set your own token through the `HF_TOKEN` environment variable.
```bash
pip install "huggingface_hub>=0.27"
```
```python
from huggingface_hub import snapshot_download
snapshot_download(
repo_id="xtxx/CellFM",
local_dir="CellFM",
token=True, # Use the cached token of an approved account.
)
```
```bash
cd CellFM
pip install -r requirements.txt
```
Run inference through `cellfm.py` and the native PyTorch loader. `config.json` records the architecture and preprocessing settings; this package does not implement the Transformers `AutoModel` interface.
## 2. Generate Nuclear Contours with HoverFast
Install [HoverFast](https://github.com/choosehappy/HoverFast) and obtain its checkpoint according to the official instructions. Segmentation and CellFM can run in separate environments, exchanging contours through GeoJSON.
ROI segmentation:
```bash
HoverFast infer_roi 'inputs/*.png' \
-m /path/to/hoverfast_pretrained_pannuke.pth \
-o hoverfast_output
```
Whole-slide segmentation:
```bash
HoverFast infer_wsi 'slides/*.svs' \
-m /path/to/hoverfast_crosstissue_best_model.pth \
-n 20 -o hoverfast_output
```
Pass the resulting contour file as `--polygons` in the examples below. `cellfm.py` accepts plain GeoJSON and gzip-compressed `.json.gz` or `.geojson.gz` files. Supported structures include a `FeatureCollection`, an array of Features, or a single Feature. Both `Polygon` and `MultiPolygon` geometries are supported, including holes and overlapping instances.
Each Feature must represent one nucleus. Preserve individual instances rather than merging all nuclei into one foreground polygon. IDs are read from `properties.cell_id`, then the Feature's `id`, then `properties.id`; otherwise, the zero-based Feature index is used. Duplicate IDs raise an error. A minimal input is:
```json
{
"type": "FeatureCollection",
"features": [{
"type": "Feature",
"properties": {"cell_id": "nucleus-0"},
"geometry": {
"type": "Polygon",
"coordinates": [[[40,30],[51,30],[51,41],[40,41],[40,30]]]
}
}]
}
```
Coordinates must align with the image, use **(x, y)** order, and be expressed in pixels. By default, annotations are interpreted as **global level-0 pixel coordinates**, not micrometers or thumbnail coordinates.
For HoverFast WSI output, check the coordinate origin. In the [official implementation inspected for this release](https://github.com/choosehappy/HoverFast/blob/main/hoverfast/utils_wsi.py), image reading adds `openslide.bounds-x/y`, while exported contours use region coordinates relative to those bounds. For this form of raw HoverFast WSI output, add `--hoverfast-bounds-offset` to restore the slide's bounds origin. This has no effect when both offsets are zero. Omit the flag for annotations that already use global coordinates, including upstream versions or HEST files with that convention. Check alignment against your actual slide and HoverFast version.
For HEST-style `*_hoverfast_seg.parquet` files containing WKB `geometry` and instance `cell_id` columns:
```bash
pip install "pyarrow>=12" "shapely>=2"
python examples/parquet_to_geojson.py sample_hoverfast_seg.parquet nuclei.json.gz
```
The converter preserves IDs and coordinates without inferring units, scale, or image origin.
## 3. `token_mean`: Cell-Centered Context
```bash
python cellfm.py \
--image inputs/roi.png \
--polygons hoverfast_output/roi.json.gz \
--method token_mean \
--device cuda --batch-size 1 \
--output outputs/roi_token_mean.npz
```
This method follows the project's `token_mean` preprocessing definition:
1. Center a **224×224** window on the nucleus bounding-box center, keeping the native image resolution. Pad context outside the image with RGB black.
2. If the nucleus bounding box exceeds 224 pixels, enlarge the square window to a multiple of four that covers it, then resize the window to 224×224. Use bicubic interpolation for RGB and nearest-neighbor interpolation for the mask.
3. Divide RGB values by 255 and normalize with ImageNet mean `[0.485, 0.456, 0.406]` and standard deviation `[0.229, 0.224, 0.225]`.
4. Read `x_norm_patchtokens` after the final LayerNorm: shape `(B, 3136, 768)`, corresponding to a **56×56** patch grid.
5. Select a token if any pixel in its 4×4 patch belongs to the nuclear mask. Compute the **equally weighted mean of selected tokens**. CLS and register tokens are excluded; tokens are not weighted by foreground pixel area.
This gives each cell its own centered context. Nearby cells may require repeated encoding of overlapping image regions.
## 4. `slide`: Shared Image Windows
For an ROI image:
```bash
python cellfm.py \
--image inputs/roi.png \
--polygons hoverfast_output/roi.json.gz \
--method slide \
--device cuda --batch-size 1 \
--output outputs/roi_slide.npz
```
An image with both dimensions at most 224 is placed at the top-left of one window, with black padding on the right and bottom. Larger images are partitioned from `(0,0)` into **non-overlapping 224×224 windows with stride 224**. Only windows containing nuclear-mask pixels are encoded. Each window is encoded once, and its tokens are pooled separately for every nucleus it intersects.
For a nucleus crossing window boundaries, accumulate its **selected-token vectors and token counts** across all intersecting windows, then divide the total vector sum by the total token count. Do not average the window means equally. Each nucleus produces one output row. Overlapping instances may select the same token because each instance is pooled independently.
`slide` describes the windowing strategy; the output remains a **per-cell embedding**. Its context and patch alignment differ from `token_mean`, so the two methods need not produce identical vectors. Image padding is excluded from nuclear masks.
### Stream a Whole-Slide Image
```bash
pip install -r requirements-wsi.txt
python cellfm.py \
--image slides/sample.svs --wsi --level 0 \
--polygons hoverfast_output/sample.json.gz \
--hoverfast-bounds-offset \
--method slide --device cuda --batch-size 1 \
--output outputs/sample_slide.npz
```
OpenSlide reads RGB windows on demand without constructing a full-slide RGB array or instance-mask image. Polygons, the occupied-window index, and output embeddings remain in memory. One million float32 cell embeddings alone require approximately 3.07 GB; larger inputs can be divided into tissue regions with matching annotations.
`--level N` uses the slide's actual pyramid downsample factor to transform level-0 coordinates, after applying any annotation offset, into the selected level. The default is **level 0 at native resolution**. Magnification, MPP, and stain normalization are not chosen automatically. Keep the image scale fixed and record it when comparing embeddings.
If an ROI was cropped from level-0 origin `(10000,20000)` but its polygons retain global coordinates:
```bash
python cellfm.py --image roi.png --polygons roi_global.json.gz \
--origin 10000 20000 --method slide --output outputs/roi.npz
```
If that ROI was also downsampled by a factor of two, add `--coordinate-downsample 2`. To add a known translation to annotation coordinates, use `--polygon-offset X Y`. The complete transform is `(polygon + polygon_offset - origin) / downsample`.
Filter annotations to the ROI before inference. Instances entirely outside the image, or with no mask pixels after clipping and resizing, raise an error. In WSI mode, use `--level` rather than the ROI-specific `--origin` and `--coordinate-downsample` options.
## 5. Output Format and Python API
```python
import json
import numpy as np
data = np.load("outputs/roi_slide.npz", allow_pickle=False)
embedding = data["embeddings"] # (N, 768), float32
cell_ids = data["cell_ids"] # (N,), Unicode strings in input Feature order
boxes = data["bboxes_xyxy"] # (N, 4), x0,y0,x1,y1 in inference-image pixels
counts = data["token_counts"] # (N,), number of selected patch tokens per cell
metadata = json.loads(data["metadata_json"].item())
print(embedding.shape, cell_ids[:3], metadata["method"])
```
The default output contains pooled embeddings without L2 normalization. Add `--l2-normalize` for normalized vectors, for example for cosine retrieval. Keep this choice consistent when fitting and evaluating a probe. An empty annotation input produces an embedding array of shape `(0,768)`.
From Python, run the following in the repository directory:
```python
from cellfm import ImageSource, load_cells, load_model, extract_slide, extract_token_mean
model = load_model(device="cpu")
source = ImageSource("inputs/roi.png")
try:
cells = load_cells("hoverfast_output/roi.json.gz")
embedding, token_counts = extract_slide(model, source, cells, device="cpu", batch_size=1)
# extract_token_mean(...) accepts the same arguments.
finally:
source.close()
```
The default tensor-only weights, `cellfm_150k_inference.pt`, are loaded strictly with `torch.load(weights_only=True)`. CUDA inference uses bf16 autocast and requires a bf16-capable device; CPU inference uses float32. The default batch size is one. Release validation used CPU; CUDA output and performance have not been validated.
## 6. Minimal Installation Check
```bash
python examples/make_demo.py --output demo
python cellfm.py --image demo/roi.png --polygons demo/nuclei.json.gz \
--method token_mean --device cpu --output outputs/demo_token_mean.npz
python cellfm.py --image demo/roi.png --polygons demo/nuclei.json.gz \
--method slide --device cpu --output outputs/demo_slide.npz
python -m unittest test_cellfm -v
```
This creates a synthetic RGB image and two nuclear contours, including one crossing the window boundary at x=224. It checks installation and data flow without pathology data or running HoverFast segmentation.
`verification.json` records strict CPU loading of the actual weights, agreement with the original loader, and engineering checks of both 224-pixel inference paths. Numerical contract tests cover cross-window weighting, padding exclusion, overlapping instances, holes and MultiPolygons, gzip input, coordinate transforms, and empty input. HoverFast segmentation, OpenSlide WSI reading, and Parquet conversion were not executed in the release validation and require validation in their respective environments with actual data. No new downstream benchmark results are claimed here.
## Model Architecture and Files
The architecture is ViT-Base with 4×4 patches: 12 distinct blocks, hidden size 768, 12 attention heads, one CLS token, and eight register tokens. Blocks `0:4` run once, blocks `4:8` run four times with shared weights, and blocks `8:12` run once, for 24 block calls per forward pass. MARB combines attention with a local depthwise 3×3 morphology branch through a learned gate. The repeated blocks use dilation 1/2/3/4; prefix tokens are not convolved. The inference model has **99,478,284 parameters**.
- `epoch_0029_ema_student.pt`: original EMA checkpoint, 398,041,244 bytes, epoch 29, global step 150000.
- `cellfm_150k_inference.pt`: tensor-only state dict produced by renaming the original 335 tensor keys, preserving tensor values, for strict inference loading.
- `load_epoch_0029.py`: original self-contained model definition, key conversion, and loader. Its historical class name `ViTP8` does not denote patch8; this model uses patch4.
- `cellfm.py`: model loading, image reading, polygon rasterization, both pooling methods, CLI, and `.npz` output.
- `config.json`: architecture and preprocessing metadata.
- `SHA256SUMS`: release file checksums.
Original checkpoint SHA256:
```text
f7c2a769012377f560688c05e239a9a6c0591c39d45a0cf182217086908566cd
```
Verify the original file with:
```bash
sha256sum epoch_0029_ema_student.pt
```
No license is specified for this repository. HoverFast code, checkpoints, and their licenses are provided independently by the official project. When using HoverFast, follow its citation instructions and cite [Liakopoulos et al., JOSS 2024](https://doi.org/10.21105/joss.07022).