--- 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 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 “true” augmented features with an average cosine similarity of **about 81%** 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) --- ## Usage You can load the model from the Hub with Hugging Face’s `transformers`: ```python import torch from transformers import AutoModel # Load HistAug (H-optimus-1 latent augmentation model) model_id = "sofieneb/histaug-hoptimus1" model = AutoModel.from_pretrained(model_id, trust_remote_code=True) # Example: patch embeddings from H-optimus-1 num_patches = 50000 embedding_dim = 1536 patch_embeddings = torch.randn((num_patches, embedding_dim), device="cuda") # 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, "gamma": [-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`, `gamma`, `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 your cluster doesn’t allow internet access on compute nodes, **pre-download the model on the front-end/login node** (which has internet), so it’s cached locally, then run jobs offline: ```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)" # On your compute job (no internet): export HF_HUB_OFFLINE=1 export TRANSFORMERS_OFFLINE=1 ``` This prevents unnecessary network calls and ensures `transformers` loads HistAug from the local cache. --- ## Paper *Controllable Latent Space Augmentation for Digital Pathology* (Boutaj *et al.*, 2025).