如何在 TIAToolBox 中使用自定义 PyTorch HoVer-Net 模型进行核分割？

Question

我想使用 TIAToolBox 来分割组织图像中的细胞核。预训练的模型不符合我的要求。因此，我在 PanNuke 数据集上训练了一个自定义 HoVer-Net 模型，以检测/计数组织图像中的细胞核。

我得到的输出是一个 .pt 文件。

现在我正在寻找一种在 TIAToolBox 中使用自定义 PyTorch HoVer-Net 模型的方法。 我尝试将其放在预训练模型所在的同一文件夹中，但这不起作用。 我还尝试使用 torch 加载模型，并将其注入 NucleusInstanceSegmentor，但这也不起作用：

目前的做法：

n_classes_pannuke = 6

# load the model
hovernet = HoVerNet(n_classes=n_classes_pannuke)

# wrap model to use multi-GPU
hovernet = torch.nn.DataParallel(hovernet)

hovernet.load_state_dict(torch.load("hovernet_best_perf.pt"))
hovernet.eval();

# Tile prediction
inst_segmentor = NucleusInstanceSegmentor(
    model=hovernet,
    num_loader_workers=2,
    num_postproc_workers=2,
    batch_size=4,
)

# Analyse the images
tile_output = inst_segmentor.predict(
    imgs_to_predict,
    save_dir="sample_tile_results/",
    mode="tile",
    on_gpu=ON_GPU,
    crash_on_exception=True,
)

错误：

ValueError: Must provide either ioconfig or patch_input_shape and patch_output_shape

有人知道如何在 TIAToolBox 中使用自定义 PyTorch HoVer-Net 模型来检测原子核吗？

Answer 1

您必须先定义一个

ioconfig

对象并将其与

predict

函数一起使用，例如：

iostate = IOSegmentorConfig(
input_resolutions=[
    {"units": "mpp", "resolution": 1.0},
],
output_resolutions=[
    {"units": "mpp", "resolution": 1.0},
],
patch_input_shape=[512, 512],
patch_output_shape=[512, 512],
stride_shape=[512, 512],
save_resolution={"units": "mpp", "resolution": 1.0},)

您可以参考工具箱的示例来了解如何使用自定义模型这里。