ufmg-digital-pathology/breast-cancer-virtual-staining

Breast Cancer Virtual Staining (H&E to Ki-67 / pHH3)

CycleGAN generators that synthesise Ki-67 and pHH3 immunohistochemistry (IHC) appearance from H&E histopathology tiles of triple-negative breast cancer (TNBC).

The models were developed as part of a master's dissertation on predicting pathological complete response (pCR) in TNBC. Their purpose is to supply biomarker-like channels for a downstream spatial-attention pCR classifier in cases where real Ki-67 / pHH3 restains are unavailable (not to replace IHC).

Models

Each .pth is a full training checkpoint (~32 MB) and also contains the reverse generator (G_ki672he_state_dict), both discriminators (D_he_state_dict, D_ki67_state_dict) and the epoch counter.

Only the forward generator is needed for inference.

Note the pHH3 checkpoint stores its discriminator under D_ki67_state_dict. This is a variable-naming carryover from the shared training script, not a Ki-67 discriminator. It does not affect the forward generator.

Architecture

ResNet-style CycleGAN generator (~1.37M parameters):

• —ReflectionPad2d(3) → Conv2d(3, 32, 7) → InstanceNorm2d → ReLU
• —2 × downsampling Conv2d(stride=2) → 32 → 64 → 128 channels
• —4 × residual blocks at 128 channels
• —2 × upsampling ConvTranspose2d(stride=2) → 128 → 64 → 32 channels
• —ReflectionPad2d(3) → Conv2d(32, 3, 7) → Tanh

Input / output contract

This is the part most likely to trip you up, because training and the output activation disagree in range:

• —Input: RGB, 512×512, float32 scaled to `[0, 1]` (plain /255.0) — not ImageNet-normalised and not [-1, 1], despite the Tanh output layer.
• —Output: Tanh, so raw values land in [-1, 1]. Training, validation and test all apply torch.clamp(output, 0, 1) before use or saving. Reproduce that clamp, or your outputs will not match the dissertation's.
• —Tiles were resized with cv2.INTER_AREA. Channel order is RGB.

Usage

from huggingface_hub import hf_hub_download
import torch

from cyclegan import CycleGenerator, HES2KI67_STATE_DICT_KEY

path = hf_hub_download(
    repo_id="ufmg-digital-pathology/breast-cancer-virtual-staining",
    filename="models/ki67/best_cyclegan_ki67_model.pth",
    revision="weights-v1",
)

model = CycleGenerator("ki67")
model.load_state_dict(torch.load(path, map_location="cpu", weights_only=True)[HES2KI67_STATE_DICT_KEY])
model.eval()

with torch.no_grad():
    virtual_ki67 = torch.clamp(model(he_tile), 0, 1)

src/sample.py in this repo wraps the above for both markers.

Always pass an explicit revision. The weights-v1 tag pins the checkpoints as released; main will move as the code changes.

Training

Both markers were trained with the same procedure, differing only in the target stain.

Data. H&E tiles paired with registered restains of the same tissue, from TNBC cases labelled pCR / non-pCR. Pairs were filtered by registration quality: within each split, only pairs in the top 20% by SSIM (the 0.80 quantile) were kept. Splits are patient-level (test_size=0.1, val_size=0.15, random_state=42) with an assertion that no patient ID appears in two splits.

Because the pairs are spatially registered, this is a CycleGAN trained on paired data, which is what makes the added SSIM and identity terms meaningful.

Objective. Generator loss is a weighted sum:

The discriminator adds a gradient penalty (λ = 0.001) and gradient-norm clipping at 0.1.

Optimisation. Adam, lr = 4e-4, betas = (0.5, 0.999), for both G and D. ReduceLROnPlateau (factor=0.5, patience=5) on the respective validation losses. Batch size 2, mixed precision (torch.cuda.amp), max 150 epochs, early stopping after 15 epochs without improvement. Augmentation is random horizontal and vertical flips applied identically to tile and target.

Checkpoint selection. Best validation mIoU (thresholded at 0.5), with pixel accuracy (5% tolerance) as the tie-break. The released files are those best checkpoints.

Intended use

• —Generating virtual Ki-67 / pHH3 channels for TNBC H&E tiles, as input to biomarker-guided spatial attention in pCR prediction.
• —Methods research on virtual staining and stain-to-stain translation.

Out of scope

• —Any clinical or diagnostic use.
• —Quantitative biomarker readouts: Ki-67 proliferation index, mitotic counts.
• —Tissue other than breast, cancer subtypes other than TNBC, or non-TNBC cohorts.

Reproducibility

The released weights are pinned by the weights-v1 tag. Pass it as revision to any huggingface_hub download so results stay reproducible as main advances.

Developed and run on Python 3.9. The environment is pinned in requirements.txt (torch==2.8.0, torchvision==0.23.0).

Related work

Training code for these models, the cGAN and diffusion variants, and the downstream Kedro pCR-prediction pipeline live in the pcr-classification repository. Only the CycleGAN weights are released here.

License

MIT. Released for non-commercial academic research.

The training data is not released and is not covered by this license.