Patho3dMatrix-Liver

Using Patho3dMatrix_Liver to extract features from pathology image

import torch
import timm
from PIL import Image
from torchvision import transforms
from safetensors.torch import load_file

MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]

if __name__ == '__main__':
    # Init Patho3DMatrix_Liver Foundation Model
    patho3dmatrix_Liver = timm.create_model(
        "vit_large_patch16_224",
        pretrained=False,
        init_values=1e-5,
        dynamic_img_size=True,
        num_classes=0,
    )

    # Load safetensors weights
    patho3dmatrix_Liver_weights_path = 'pytorch_model.safetensors'
    state_dict = load_file(patho3dmatrix_Liver_weights_path, device='cpu')
    msg = patho3dmatrix_Liver.load_state_dict(state_dict, strict=True)
    print(msg)
    print('weights loaded successfully')

    # Set device
    device = torch.device('cuda:0')
    patho3dmatrix_Liver = patho3dmatrix_Liver.to(device)
    patho3dmatrix_Liver.eval()

    # Image preprocess: Resize(224) -> ToTensor -> Normalize(ImageNet mean/std)
    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=MEAN, std=STD),
    ])

    # Encode one image
    img_path = 'test.png'
    img = Image.open(img_path).convert('RGB')
    img_tensor = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        feat = patho3dmatrix_Liver(img_tensor)
        if feat.dim() == 3:
            feat = feat[:, 0, :]  # CLS token

    print('feature shape:', feat.shape)  # [1, 1024]

For a whole-slide bag, encode every patch the same way, then concatenate to a bare tensor [N, 1024] (float32) and torch.save it. That is the input expected by the OS/PFS heads.

Downstream prognosis: features to OS / PFS risk score

downstream/ ships one AB-MIL head per endpoint, trained on TCGA-LIHC with Patho3dMatrix_Liver (tea) features. See downstream/data_description.md for the cohort, endpoints, and training protocol.

Input must be tea features [N, 1024] float32 (a saved .pt bag, or a stacked batch of encoder outputs). A dict with a features / feature key is also accepted. The head returns a scalar risk score (higher = higher risk). Risk is negatively associated with survival time, so C-Index should be computed as concordance_index(times, -risk, events).

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file


def initialize_weights(module):
    for m in module.modules():
        if isinstance(m, nn.Linear):
            nn.init.xavier_normal_(m.weight)
            if m.bias is not None:
                m.bias.data.zero_()


class AB_MIL(nn.Module):
    def __init__(self, L=512, D=128, num_classes=1, dropout=0.25, in_dim=1024):
        super(AB_MIL, self).__init__()
        self.L = L
        self.D = D
        self.feature = nn.Sequential(
            nn.Linear(in_dim, self.L),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.attention = nn.Sequential(
            nn.Linear(self.L, self.D),
            nn.Tanh(),
            nn.Linear(self.D, 1),
        )
        self.classifier = nn.Linear(self.L, num_classes)
        self.apply(initialize_weights)

    def forward(self, x):
        x = x.squeeze(0)          # [N, 1024]
        h = self.feature(x)
        A = self.attention(h)     # [N, 1]
        A = torch.transpose(A, 1, 0)
        A = F.softmax(A, dim=1)
        M = torch.mm(A, h)        # [1, 512]
        logits = self.classifier(M)
        return {'logits': logits, 'A': A}


def load_features(path):
    content = torch.load(path, map_location='cpu', weights_only=True)
    if isinstance(content, dict):
        feat = content.get('features', content.get('feature', None))
    else:
        feat = content
    if feat is None:
        raise ValueError(f'no features/feature key in {path}')
    if not isinstance(feat, torch.Tensor):
        feat = torch.tensor(feat)
    return feat.float()           # [N, 1024]


if __name__ == '__main__':
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # Switch to 'downstream/PFS.safetensors' for progression-free survival
    head_path = 'downstream/OS.safetensors'
    model = AB_MIL(in_dim=1024)
    model.load_state_dict(load_file(head_path))
    model.to(device)
    model.eval()

    features = load_features('features.pt')   # [N, 1024] float32
    with torch.no_grad():
        risk = model(features.unsqueeze(0).to(device))['logits'].cpu().item()

    print('risk score:', risk)    

Evaluation Pipeline

• WSI Classification: https://github.com/lingxitong/MIL_BASELINE
• ROI Classification: https://github.com/lingxitong/HistoROIBench
• ROI Segmentation: https://github.com/lingxitong/PFM_Segmentation