README.md

---
tags:
- model_hub_mixin
- pytorch_model_hub_mixin
language:
- en
---

## Model Summary

**HistAug** is a lightweight transformer-based generator for **controllable latent-space augmentations** in the feature space of the [CONCH foundation model](https://www.nature.com/articles/s41591-024-02856-4). Instead of applying costly image-space augmentations on millions of WSI patches, HistAug operates **directly on patch embeddings** extracted from a given foundation model(here CONCH). By conditioning on explicit transformation parameters (e.g., hue shift, erosion, HED color transform), HistAug generates realistic augmented embeddings while preserving semantic content. In practice, the CONCH variant of HistAug can reconstruct the corresponding ground-truth augmented embeddings with an average cosine similarity of **about 93%** at **10X, 20X, and 40X magnification**.

This enables training of Multiple Instance Learning (MIL) models with:  
- ⚡ **Fast augmentation** 
- 🧠 **Low memory usage** (up to 200k patches in parallel on a single V100 32GB GPU)  
- 🎛 **Controllable and WSI-consistent augmentations** (bag-wise or patch-wise)

Need HistAug for a different foundation model? Explore the full collection: [**HistAug models collection**](https://huggingface.co/collections/sofieneb/histaug-models-68a334437f71d35c7037a54e).


📄 **Paper**: [*Controllable Latent Space Augmentation for Digital Pathology* (Boutaj *et al.*, 2025)](https://arxiv.org/abs/2508.14588)




---

## Usage

You can load the model from the Hub with Hugging Face’s `transformers`:

```python
import torch
from transformers import AutoModel

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load HistAug (CONCH latent augmentation model)
model_id = "sofieneb/histaug-conch"
model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device)

# Example: patch embeddings from CONCH
num_patches = 50000
embedding_dim = 512
patch_embeddings = torch.randn((num_patches, embedding_dim), device=device)

# Sample augmentation parameters
# mode="wsi_wise" applies the same transformation across the whole slide
# mode="instance_wise" applies different transformations per patch
aug_params = model.sample_aug_params(
    batch_size=num_patches,
    device=patch_embeddings.device,
    mode="wsi_wise"
)

# Apply augmentation in latent space
augmented_embeddings = model(patch_embeddings, aug_params)

print(augmented_embeddings.shape)  # (num_patches, embedding_dim)
```

## Default Transform Configuration

The original transform configuration (shipped in the model config) is:

```json
{
  "transforms": {
    "parameters": {
      "brightness": [-0.5, 0.5],
      "contrast": [-0.5, 0.5],
      "crop": 0.75,
      "dilation": 0.75,
      "erosion": 0.75,
      "powerlaw": [-0.5, 0.5],
      "gaussian_blur": 0.75,
      "h_flip": 0.75,
      "hed": [-0.5, 0.5],
      "hue": [-0.5, 0.5],
      "rotation": 0.75,
      "saturation": [-0.5, 0.5],
      "v_flip": 0.75
    }
  }
}
```

* **Continuous transforms** (e.g., `brightness`, `hue`, `hed`, `powerlaw`, `saturation`) use an **interval** `[min, max]` from which parameters are sampled.
* **Discrete/binary transforms** (e.g., `h_flip`, `v_flip`, `dilation`, `erosion`, `rotation`, `gaussian_blur`, `crop`) use a **probability** (e.g., `0.75`) indicating how likely the transform is applied during sampling.

> You can access and modify this at runtime via:
>
> ```python
> print(model.histaug.transforms_parameters)
> ```

---

## Controlling Transformations

You can **inspect, modify, or delete** transformations at runtime via `model.histaug.transforms_parameters`.  
- To **remove** a transform, simply `pop` the key; during sampling it will appear with parameter **`0`** (effectively disabled).  
- You can also narrow a transform’s interval or change a transform’s probability, then re-sample to observe the effects.  
- Sampling mode: `mode="wsi_wise"` (same parameters for all patches) or `mode="instance_wise"` (per-patch parameters).

```python
## Controlling Transformations — pop vs. change params (continuous & discrete)

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
num_to_sample = 5

# start: sample once and inspect current config
sample_1 = model.sample_aug_params(batch_size=num_to_sample, device=device, mode="wsi_wise")
print("initial sample:\n", sample_1, "\n")

print("initial transforms_parameters:\n", model.histaug.transforms_parameters, "\n")

# pop examples
# pop a continuous transform: remove "hue" (interval transform)
model.histaug.transforms_parameters.pop("hue", None)

# pop a discrete transform: remove "rotation" (probability-based)
model.histaug.transforms_parameters.pop("rotation", None)

sample_2 = model.sample_aug_params(batch_size=num_to_sample, device=device, mode="wsi_wise")
print("after popping 'hue' (continuous) and 'rotation' (discrete):\n", sample_2, "\n")

# change param examples
# change a continuous transform interval: narrow 'brightness' from [-0.5, 0.5] to [-0.25, 0.25]
model.histaug.transforms_parameters["brightness"] = [-0.25, 0.25]

# change a discrete transform probability: lower 'h_flip' from 0.75 to 0.10
model.histaug.transforms_parameters["h_flip"] = 0.10

sample_3 = model.sample_aug_params(batch_size=num_to_sample, device=device, mode="wsi_wise")
print("after changing 'brightness' interval and 'h_flip' probability:\n", sample_3, "\n")

````
---

## During MIL

You can apply latent-space augmentation **during MIL training** with a probability (e.g., **60%**). We generally recommend applying augmentation with a non-trivial probability (e.g., 0.3–0.7) rather than always-on.

```python
import torch

# histaug: the loaded HistAug model (CONCH variant)
# mil_model: your MIL aggregator (e.g., ABMIL/CLAM/TransMIL head)
# criterion, optimizer, loader already defined

device = "cuda" if torch.cuda.is_available() else "cpu"
histaug = histaug.to(device).eval()   # histaug generator is frozen during MIL training
for p in histaug.parameters():
    p.requires_grad_(False)

def maybe_augment_bag(bag_features: torch.Tensor,
                      p_apply: float = 0.60,
                      mode: str = "wsi_wise") -> torch.Tensor:
    """
    bag_features: (num_patches, embed_dim) on device
    p_apply: probability to apply augmentation
    mode: "wsi_wise" (same params for all patches) or "instance_wise"
    """
    if torch.rand(()) >= p_apply:
        return bag_features
    with torch.no_grad():
        aug_params = histaug.sample_aug_params(
            batch_size=bag_features.size(0),
            device=bag_features.device,
            mode=mode  # "wsi_wise" or "instance_wise"
        )
        bag_features = histaug(bag_features, aug_params)
    return bag_features

# --- single-bag training example ---
for bag_features, label in loader:  # bag_features: (num_patches, embed_dim)
    bag_features = bag_features.to(device)

    # apply augmentation with 60% probability (WSI-wise by default)
    bag_features = maybe_augment_bag(bag_features, p_apply=0.60, mode="wsi_wise") # output : (num_patches, embed_dim)

    logits = mil_model(bag_features)              # forward through your MIL head
    loss = criterion(logits, label.to(device))
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

```

---

## Offline usage (HPC clusters without internet)

If compute nodes don’t have internet, **always** run jobs with the offline flags to **prevent unnecessary network calls** and force local loads:

```bash
# On your compute job (no internet):
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
```

Prepare the model **in advance** on a front-end/login node (with internet), then choose **either** approach below.

### Option — Warm the cache (simplest)

```bash
# On the front-end/login node (with internet):
python -c "from transformers import AutoModel; AutoModel.from_pretrained('sofieneb/histaug-conch', trust_remote_code=True)"
```

Then in your offline job/script:

```python
from transformers import AutoModel
model = AutoModel.from_pretrained(
    "sofieneb/histaug-conch",
    trust_remote_code=True,
    local_files_only=True,  # uses local cache only
)
```

### Option — Download to a local folder with `hf download`

```bash
# On the front-end/login node (with internet):
hf download sofieneb/histaug-conch --local-dir ./histaug-conch
```

Then in your offline job/script:

```python
from transformers import AutoModel
model = AutoModel.from_pretrained(
    "./histaug-conch",   # local path instead of hub ID
    trust_remote_code=True,
    local_files_only=True,  # uses local files only
)
```

---
## Citation
If our work contributes to your research, or if you incorporate part of this code, please consider citing our paper:

```bibtex
@misc{boutaj2025controllablelatentspaceaugmentation,
      title={Controllable Latent Space Augmentation for Digital Pathology}, 
      author={Sofiène Boutaj and Marin Scalbert and Pierre Marza and Florent Couzinie-Devy and Maria Vakalopoulou and Stergios Christodoulidis},
      year={2025},
      eprint={2508.14588},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2508.14588}, 
}
```