elliotzhangQx/path_encoder

WSI 特征提取：五个病理基础模型的统一封装

用同一个接口调用 TITAN / UNI2-h / CONCH / MUSK / KEEP 提取 WSI patch 特征。

纯本地运行，不联网，不需要 HuggingFace token。 权重已下载完毕。

encoder	特征维度	模型输入	归一化
uni2h	1536	224	ImageNet
conch	512	448	CLIP
keep	768	224	ImageNet（输出已 L2 归一化）
musk	1024	384	(0.5, 0.5, 0.5)
titan	768（slide-level）	448	ImageNet

---

1. 怎么调用

用 example.py 里的 example() 函数。 输入一整张 slide 的 patch tensor 和 coordinates，内部自动分 batch，提完特征存成一个 h5：

from example import example

example(imgs,                      # (N, 3, H, W) 一整张 slide 的所有 patch
        coords,                    # (N, 2) level-0 坐标，与 imgs 逐行对应
        "uni2h",                   # encoder 名
        "out/slide_001.h5",        # 输出 h5 路径
        batch_size=256)            # 默认 256

TITAN 多传一个 patch_size_lv0（它需要坐标来做 slide 级聚合）：

example(imgs, coords, "titan", "out/slide_001.h5", patch_size_lv0=512)

example.py 可以直接运行，它会造一份 fake slide（600 个 patch）走完整流程， 拿它当模板改：

python example.py --encoder uni2h
python example.py --encoder musk

更底层的用法

如果不想用 example() 的存盘逻辑，直接用 OfflineEncoder——构造一次，每个 batch 调一次，自己决定怎么存：

from extractor import OfflineEncoder

enc = OfflineEncoder("uni2h", weights_dir="/your/weights")
feats = enc(imgs)          # (B,3,H,W) -> (B,1536)，等价于 enc.encode(imgs)

---

2. 传入数据的格式要求

example() 的参数

参数	类型	要求
imgs	(N, 3, H, W) tensor	一整张 slide 的所有 patch，RGB，未归一化
coords	(N, 2) 整数 tensor	level-0 像素坐标，(x, y) 列顺序，与 imgs 逐行对应
encoder_name	str	uni2h / conch / keep / musk / titan
out_path	str	输出 .h5 路径，父目录自动创建
batch_size	int	默认 256，只影响显存占用，不改变结果
patch_size_lv0	int	TITAN 必填，其余 encoder 忽略

可选：weights_dir、device、amp_dtype。

patch tensor

• 必须是未归一化的 RGB。五个模型的 mean/std 各不相同（见上表），由 encoder 内部各自归一化。如果你的读图流程已经做过 ImageNet 归一化，要先反归一化，否则 CONCH 和 MUSK 的特征是错的。
• 数值范围 [0,1] 的 float 即可。若是 0-255，example() 里构造 OfflineEncoder 时要加 input_range="255"。
• H/W 任意，encoder 内部会 resize 到上表的尺寸。直接给原始 patch 尺寸就行。
• N 是整张 slide 的 patch 总数。函数内部按 batch_size 切分（最后一批可以不 足），提完特征按原顺序拼回，所以 features[i] 一定对应 coords[i]。

coordinates

两个硬要求：

1. **列顺序是 (x, y)**，即 (列, 行)。
2. 必须是 level-0 像素坐标——每个 patch 左上角在全分辨率图上的位置。不是 patch 索引，也不是 level-1/2 的坐标。索引 i、步长 s 的 patch 要传 i * s。

patch_size_lv0 是相邻 coords 之间的步长。512px 的 patch 若在 20x 下读取而 level-0 是 40x，这里要传 1024 而不是 512。

只有 TITAN 会用 coords 参与计算，其余 encoder 只是把它原样存进 h5。 coords 或 patch_size_lv0 传错不会报错，只会得到一个错误的 slide embedding。

3. 输出格式

一张 slide 一个 h5：

dataset	形状	说明
features	(N, D) float32	patch 级特征，与 coords 逐行对应
coords	(N, 2) int32	原样写回
slide_embedding	(768,) float32	仅 TITAN，slide 级特征

attrs：encoder、hf_repo、feat_dim、input_size、n_patches （TITAN 另有 slide_dim、patch_size_lv0）。

import h5py
with h5py.File("out/slide_001.h5") as f:
    feats  = f["features"][:]       # (N, D)
    coords = f["coords"][:]         # (N, 2)
    print(dict(f.attrs))

TITAN 是 slide-level 模型：features 是 CONCH v1.5 的 patch 级中间结果， slide_embedding 才是它的最终输出，两级都会保存。

---

4. 环境安装

MUSK 必须单独一个环境：它依赖的 timm 版本与其余四个冲突（UNI2-h 需要 timm>=1.0.3 的 SwiGLUPacked，MUSK 的 model registry 要 timm 0.9.x）。

# 环境 A：UNI2-h / CONCH / TITAN / KEEP
conda create -y -n pathenc python=3.10
~/anaconda3/envs/pathenc/bin/pip install -r requirements.txt \
    --extra-index-url https://download.pytorch.org/whl/cu121
~/anaconda3/envs/pathenc/bin/pip install --no-deps \
    git+https://github.com/mahmoodlab/CONCH.git

# 环境 B：MUSK
conda create -y -n pathenc_musk python=3.10
~/anaconda3/envs/pathenc_musk/bin/pip install -r requirements_musk.txt \
    --extra-index-url https://download.pytorch.org/whl/cu121
~/anaconda3/envs/pathenc_musk/bin/pip install --no-deps \
    git+https://github.com/lilab-stanford/MUSK.git

两点说明：CONCH 和 MUSK 不在 PyPI，必须从 GitHub 装且带 --no-deps（否则会把 timm 降级）；**不要装 PyPI 的 torchscale**，MUSK 自带 vendored 版本，PyPI 那个硬 pin timm==0.6.13 会导致解析失败。

5. 放置权重

五个目录放在同一个根目录下，结构必须是 <weights_dir>/<encoder_name>/：

/your/weights/
├── uni2h/   2.6G
├── conch/   766M
├── keep/    3.1G
├── musk/    3.7G
└── titan/   1.9G

共约 12G。路径不必和原机器一致。优先级：构造参数 weights_dir > 环境变量 ENCODER_WEIGHTS_DIR > encoders.py 里的 DEFAULT_WEIGHTS_DIR。 固定用一个路径的话设环境变量最省事：

export ENCODER_WEIGHTS_DIR=/your/weights

运行时会在 <weights_dir>/hf_cache/ 下自动生成缓存目录（暂存模型的 remote code）， 不需要随权重一起交接，程序会自己重建。

6. 验证环境

五个 encoder 各跑一次 example.py，都打印 OK 就说明环境和权重没问题：

export ENCODER_WEIGHTS_DIR=/your/weights
for E in uni2h conch keep titan; do
  ~/anaconda3/envs/pathenc/bin/python example.py --encoder $E
done
~/anaconda3/envs/pathenc_musk/bin/python example.py --encoder musk

---

7. 构造参数

OfflineEncoder(name, weights_dir=None, device="cuda",
               amp_dtype=None, input_range="01", **extras)

参数	说明
name	"uni2h" / "conch" / "keep" / "musk" / "titan"
weights_dir	权重根目录，不传则用环境变量或默认值
device	默认 "cuda"，多卡可用 "cuda:1"
amp_dtype	"fp16" / "bf16"，提速明显，默认关（fp32）
input_range	"01"（默认）或 "255"
ms_aug	仅 MUSK，多尺度增强，维度 1024 → 2048
l2_normalize	仅 KEEP，冗余（输出本就是单位长度）

可读属性：enc.feat_dim、enc.input_size、enc.name、enc.weights_dir。

必须在 import transformers / huggingface_hub 之前构造——离线标志是这两个库 在 import 时读取的。顺序错了会给 RuntimeWarning 提示。

8. 文件说明

文件	用途
extractor.py	OfflineEncoder 类，主入口
example.py	可运行的调用示例，拿它当模板
encoders.py	五个 encoder 的底层封装；权重目录默认值在这里
requirements.txt	环境 A 依赖
requirements_musk.txt	环境 B 依赖

9. 常见问题

FileNotFoundError: weights directory not found: .../uni2h weights_dir 指错了，或子目录名不对。必须是 <weights_dir>/<encoder_name>/， 名字用 uni2h / conch / keep / musk / titan。

ModuleNotFoundError: No module named 'conch' / 'musk' 这两个包要从 GitHub 装，见第 4 节。另外确认用的是对应环境的 python：MUSK 在 pathenc_musk，其余四个在 pathenc。

ResolutionImpossible（装 MUSK 环境时） 装了 PyPI 的 torchscale。不要装。

ValueError: patch_size_lv0 is required for TITAN TITAN 必须传这个参数，见第 2 节。

显存不够：降 batch_size（不改变结果），或传 amp_dtype="fp16"。

特征数值不对：最可能是归一化重复了。确认传入的是未归一化 RGB，且 input_range 与实际数值范围一致。

10. 已验证的内容

五个 encoder 都通过了冷启动测试：全新 HOME、新的权重路径、无 token、无缓存、 HF 端点不可达。

分 batch 的正确性：同一份输入用 batch_size 256 / 16 / 7 得到的特征一致 （差异 ~1e-5，来自 cuDNN 对不同 batch 选择不同 kernel），且逐行比对单 patch 单独编码的结果也对得上，顺序没有错乱。

与官方 example 的数值比对：喂预先 resize 到目标尺寸的 tensor 时，五个模型的输出与 官方路径位级完全一致；端到端（含 resize）的 cosine similarity ≥ 0.9998，差异仅 来自 PIL 与 GPU 的 resize 实现不同。