--- 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 [H-optimus-1 foundation model](https://huggingface.co/bioptimus/H-optimus-1). 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 H-optimus-1). 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 H-optimus-1 variant of HistAug can reconstruct the corresponding ground-truth augmented embeddings with an average cosine similarity of **about 82%** 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 (H-optimus-1 latent augmentation model) model_id = "sofieneb/histaug-hoptimus1" model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device) # Example: patch embeddings from H-optimus-1 num_patches = 50000 embedding_dim = 1536 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 (H-optimus-1 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-hoptimus1', trust_remote_code=True)" ``` Then in your offline job/script: ```python from transformers import AutoModel model = AutoModel.from_pretrained( "sofieneb/histaug-hoptimus1", 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-hoptimus1 --local-dir ./histaug-hoptimus1 ``` Then in your offline job/script: ```python from transformers import AutoModel model = AutoModel.from_pretrained( "./histaug-hoptimus1", # 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}, } ```