init commit

LICENSE

@@ -0,0 +1,58 @@
+EXAONEPath AI Model License Agreement 1.0 - NC
+
+This License Agreement (“Agreement”) is entered into between you (“Licensee”) and LG Management Development Institute Co., Ltd. (“Licensor”), governing the use of the EXAONEPath AI Model (“Model”). By downloading, installing, copying, or using the Model, you agree to comply with and be bound by the terms of this Agreement. If you do not agree to all the terms, you must not download, install, copy, or use the Model. This Agreement constitutes a binding legal agreement between the Licensee and Licensor.
+
+1. Definitions
+1.1 Model: The artificial intelligence model provided by Licensor, which includes any software, algorithms, machine learning models, or related components supplied by Licensor. This definition extends to encompass all updates, enhancements, improvements, bug fixes, patches, or other modifications that may be provided by Licensor from time to time, whether automatically or manually implemented.
+1.2 Derivatives: Any modifications, alterations, enhancements, improvements, adaptations, or derivative works of the Model created by Licensee or any third party. This includes changes made to the Model's architecture, parameters, data processing methods, or any other aspect of the Model that results in a modification of its functionality or output.
+1.3 Output: Any data, results, content, predictions, analyses, insights, or other materials generated by the Model or Derivatives, regardless of whether they are in their original form or have been further processed or modified by the Licensee. This includes, but is not limited to, textual or numerical produced directly or indirectly through the use of the Model.
+1.4 Licensor: LG Management Development Institute Co., Ltd., the owner, developer, and provider of the EXAONEPath AI Model. The Licensor holds all rights, title, and interest in the Model and is responsible for granting licenses to use the Model under the terms specified in this Agreement.
+1.5 Licensee: The individual, organization, corporation, academic institution, government agency, or other entity using or intending to use the Model under the terms and conditions of this Agreement. The Licensee is responsible for ensuring compliance with the Agreement by all authorized users who access or utilize the Model on behalf of the Licensee.
+
+2. License Grant
+2.1 Grant of License: Subject to the terms and conditions outlined in this Agreement, the Licensor hereby grants the Licensee a limited, non-exclusive, non-transferable, worldwide, and revocable license to:
+a. Access, download, install, and use the Model solely for research purposes. This includes evaluation, testing, academic research and experimentation.
+b. Publicly disclose research results and findings derived from the use of the Model or Derivatives, including publishing papers or presentations.
+c. Modify the Model and create Derivatives based on the Model, provided that such modifications and Derivatives are used exclusively for research purposes. The Licensee may conduct experiments, perform analyses, and apply custom modifications to the Model to explore its capabilities and performance under various scenarios. If the Model is modified, the modified Model must include "EXAONEPath" at the beginning of its name. 
+d. Distribute the Model and Derivatives in each case with a copy of this Agreement.
+2.2 Scope of License: The license granted herein does not authorize the Licensee to use the Model for any purpose not explicitly permitted under this Agreement. Any use beyond the scope of this license, including any commercial application or external distribution, is strictly prohibited unless explicitly agreed upon in writing by the Licensor.
+
+3. Restrictions
+3.1 Commercial Use: The Licensee is expressly prohibited from using the Model, Derivatives, or Output for any commercial purposes, including but not limited to, developing or deploying products, services, or applications that generate revenue, whether directly or indirectly. Any commercial exploitation of the Model or its derivatives requires a separate commercial license agreement with the Licensor. Furthermore, the Licensee shall not use the Model, Derivatives or Output to develop or improve other models, except for research purposes, which is explicitly permitted.
+3.2 Reverse Engineering: The Licensee shall not decompile, disassemble, reverse engineer, or attempt to derive the source code, underlying ideas, algorithms, or structure of the Model, except to the extent that such activities are expressly permitted by applicable law. Any attempt to bypass or circumvent technological protection measures applied to the Model is strictly prohibited.
+3.3 Unlawful Use: The Licensee shall not use the Model and Derivatives for any illegal, fraudulent, or unauthorized activities, nor for any purpose that violates applicable laws or regulations. This includes but is not limited to the creation, distribution, or dissemination of malicious, deceptive, or unlawful content.
+3.4 Ethical Use: The Licensee shall ensure that the Model or Derivatives is used in an ethical and responsible manner, adhering to the following guidelines:
+a. The Model and Derivatives shall not be used to generate, propagate, or amplify false, misleading, or harmful information, including fake news, misinformation, or disinformation.
+b. The Model and Derivatives shall not be employed to create, distribute, or promote content that is discriminatory, harassing, defamatory, abusive, or otherwise offensive to individuals or groups based on race, gender, sexual orientation, religion, nationality, or other protected characteristics.
+c. The Model and Derivatives shall not infringe on the rights of others, including intellectual property rights, privacy rights, or any other rights recognized by law. The Licensee shall obtain all necessary permissions and consents before using the Model and Derivatives in a manner that may impact the rights of third parties.
+d. The Model and Derivatives shall not be used in a way that causes harm, whether physical, mental, emotional, or financial, to individuals, organizations, or communities. The Licensee shall take all reasonable measures to prevent misuse or abuse of the Model and Derivatives that could result in harm or injury.
+
+4. Ownership
+4.1 Intellectual Property: All rights, title, and interest in and to the Model, including any modifications, Derivatives, and associated documentation, are and shall remain the exclusive property of the Licensor. The Licensee acknowledges that this Agreement does not transfer any ownership rights to the Licensee. All trademarks, service marks, and logos associated with the Model are the property of the Licensor.
+4.2 Output: All output generated by the Model from Licensee Data ("Output") shall be the sole property of the Licensee. Licensor hereby waives any claim of ownership or intellectual property rights to the Output. Licensee is solely responsible for the legality, accuracy, quality, integrity, and use of the Output.
+4.3 Attribution: In any publication or presentation of results obtained using the Model, the Licensee shall provide appropriate attribution to the Licensor, citing the Model's name and version, along with any relevant documentation or references specified by the Licensor.
+
+5. No Warranty
+5.1 “As-Is” Basis: The Model, Derivatives, and Output are provided on an “as-is” and “as-available” basis, without any warranties or representations of any kind, whether express, implied, or statutory. The Licensor disclaims all warranties, including but not limited to, implied warranties of merchantability, fitness for a particular purpose, accuracy, reliability, non-infringement, or any warranty arising from the course of dealing or usage of trade.
+5.2 Performance and Reliability: The Licensor does not warrant or guarantee that the Model, Derivatives or Output will meet the Licensee’s requirements, that the operation of the Model, Derivatives or Output will be uninterrupted or error-free, or that defects in the Model will be corrected. The Licensee acknowledges that the use of the Model, Derivatives or Output is at its own risk and that the Model, Derivatives or Output may contain bugs, errors, or other limitations.
+5.3 No Endorsement: The Licensor does not endorse, approve, or certify any results, conclusions, or recommendations derived from the use of the Model. The Licensee is solely responsible for evaluating the accuracy, reliability, and suitability of the Model for its intended purposes.
+
+6. Limitation of Liability
+6.1 No Liability for Damages: To the fullest extent permitted by applicable law, in no event shall the Licensor be liable for any special, incidental, indirect, consequential, exemplary, or punitive damages, including but not limited to, damages for loss of business profits, business interruption, loss of business information, loss of data, or any other pecuniary or non-pecuniary loss arising out of or in connection with the use or inability to use the Model, Derivatives or any Output, even if the Licensor has been advised of the possibility of such damages.
+6.2 Indemnification: The Licensee agrees to indemnify, defend, and hold harmless the Licensor, its affiliates, officers, directors, employees, and agents from and against any claims, liabilities, damages, losses, costs, or expenses (including reasonable attorneys' fees) arising out of or related to the Licensee's use of the Model, any Derivatives, or any Output, including any violation of this Agreement or applicable laws. This includes, but is not limited to, ensuring compliance with copyright laws, privacy regulations, defamation laws, and any other applicable legal or regulatory requirements.
+
+7. Termination
+7.1 Termination by Licensor: The Licensor reserves the right to terminate this Agreement and revoke the Licensee’s rights to use the Model at any time, with or without cause, and without prior notice if the Licensee breaches any of the terms or conditions of this Agreement. Termination shall be effective immediately upon notice.
+7.2 Effect of Termination: Upon termination of this Agreement, the Licensee must immediately cease all use of the Model, Derivatives, and Output and destroy all copies of the Model, Derivatives, and Output in its possession or control, including any backup or archival copies. The Licensee shall certify in writing to the Licensor that such destruction has been completed.
+7.3 Survival: The provisions of this Agreement that by their nature should survive termination, including but not limited to, Sections 4 (Ownership), 5 (No Warranty), 6 (Limitation of Liability), and this Section 7 (Termination), shall continue to apply after termination.
+
+8. Governing Law
+8.1 Governing Law: This Agreement shall be governed by and construed in accordance with the laws of the Republic of Korea, without regard to its conflict of laws principles.
+8.2 Arbitration: Any disputes, controversies, or claims arising out of or relating to this Agreement, including its existence, validity, interpretation, performance, breach, or termination, shall be referred to and finally resolved by arbitration administered by the Korean Commercial Arbitration Board (KCAB) in accordance with the International Arbitration Rules of the Korean Commercial Arbitration Board in force at the time of the commencement of the arbitration. The seat of arbitration shall be Seoul, Republic of Korea. The tribunal shall consist of one arbitrator. The language of the arbitration shall be English.
+
+9. Alterations
+9.1 Modifications: The Licensor reserves the right to modify or amend this Agreement at any time, in its sole discretion. Any modifications will be effective upon posting the updated Agreement on the Licensor’s website or through other means of communication. The Licensee is responsible for reviewing the Agreement periodically for changes. Continued use of the Model after any modifications have been made constitutes acceptance of the revised Agreement.
+9.2 Entire Agreement: This Agreement constitutes the entire agreement between the Licensee and Licensor concerning the subject matter hereof and supersedes all prior or contemporaneous oral or written agreements, representations, or understandings. Any terms or conditions of any purchase order or other document submitted by the Licensee in connection with the Model that are in addition to, different from, or inconsistent with the terms and conditions of this Agreement are not binding on the Licensor and are void.
+
+By downloading, installing, or using the EXAONEPath AI Model, the Licensee acknowledges that it has read, understood, and agrees to be bound by the terms and conditions of this Agreement.
+

README.md

@@ -0,0 +1,94 @@
+---
+license: other
+license_name: exaonepath
+license_link: LICENSE
+tags:
+- lg-ai
+- EXAONEPath-1.5
+- pathology
+---
+# EXAONE Path 1.5
+## Introduction
+EXAONE Path 1.5 is a whole slide image level(WSI-level) classification framework designed for downstream tasks in pathology, such as cancer subtyping, molecular subtyping and mutation prediction. It builds upon our previous work, EXAONE Path v1.0, which focused on patch-wise feature extraction by dividing a WSI into patches and embedding each patch into a feature vector.
+In EXAONE Path 1.5, we extend this pipeline to take an entire WSI as input. Each patch is first processed using the pretrained EXAONE Path 1.0 encoder to extract patch-level features. These features are then aggregated using a ViT-based (Vision Transformer) aggregator module to produce a slide-level representation.
+This aggregated representation is subsequently passed through a linear classifier to perform downstream tasks such as molecular subtyping, tumor subtyping, and mutation prediction.
+To effectively train the aggregator, we adopt a two-stage learning process:
+Pretraining: We employ multimodal learning by aligning slide images with various mRNA gene expression profiles to learn semantically meaningful slide-level representations.
+Fine-tuning: The pretrained model is then adapted to specific downstream classification tasks.
+In this repository, we release the model trained for EGFR mutation prediction in lung adenocarcinoma (LUAD), enabling researchers to leverage our pipeline for similar molecular pathology applications.
+## Quickstart
+### 1. Hardware Requirements ###
+- NVIDIA GPU is required
+- Minimum 40GB GPU memory recommended
+- Tested on Ubuntu 22.04 with NVIDIA driver version 550.144.03
+Note: This implementation requires NVIDIA GPU and drivers. The provided environment setup specifically uses CUDA-enabled PyTorch, making NVIDIA GPU mandatory for running the model.
+### 2. Environment Setup
+```
+pip install -r requirements.txt
+```
+### 3-a. Load the model & Inference
+#### Load model with HuggingFace
+ 
+```python
+from models.exaonepath import EXAONEPathV1p5Downstream
+hf_token = "YOUR_HUGGING_FACE_ACCESS_TOKEN"
+model = EXAONEPathV1p5Downstream.from_pretrained("LGAI-EXAONE/EXAONE-Path-1.5", use_auth_token=hf_token)
+slide_path = './samples/wsis/1/1.svs'
+probs = model(slide_path)
+```
+#### Fast CLI Inference
+Before running the command below, make sure you update your Hugging Face token.
+Open `tokens.py` and replace the placeholder with your actual token:
+
+```python
+HF_TOKEN = "YOUR_HUGGING_FACE_ACCESS_TOKEN"
+```
+
+Then, run inference with:
+```bash
+python inference.py --svs_path ./samples/wsis/1/1.svs
+```
+### 3-b. Fine-tuning with Pretrained Weights
+We provide example scripts and files to help you fine-tune the model on your own dataset.
+The provided script fine-tunes the model using pretrained weights stored in `./pretrained_weight.pth`.
+
+#### Extract Features from WSI Images
+To train the model using WSI images and their corresponding labels,  
+you must first extract patch-level features from each WSI using our provided feature extractor.
+
+```bash
+python feature_extract.py --input_dir ./samples/wsis/ --output_dir ./samples/feats/
+```
+This will generate .pt feature files in the output_dir.
+
+#### Fine-tuning
+```bash
+bash tuning_script.sh
+```
+Inside tuning_script.sh, you can modify the following variables to match your dataset:
+```bash
+FEAT_PATH=./samples/feats
+LABEL_PATH=./samples/label/label.csv
+LABEL_DICT="{'n':0, 'y':1}"
+SPLIT_PATH=./samples/splits
+```
+Change these paths to point to your own feature, label, and split files to start training.
+
+## Model Performance Comparison
+| Metric: AUC                      | Titan(Conch v1.5+iBot, image+text) | PRISM (virchow+pe receiver, Image+text) | CHIEF (CTransPath + CLAM, Image+text, clam+wsi contrastive) | Prov-GigaPath (GigaPath+LongNet, Image-only, mask precision manner) | UNI2-h + CLAM (Image-only) | EXAONE Path 1.5(image+gene expression) |
+|--------------------------|----------------------------------|-----------------------------------------|--------------------------------------------------------------|------------------------------------------------------------------------|-----------------------------|------------------|
+| **TMB (cutoff 10)**                | 0.74                             | 0.73                                    | 0.70                                                         | 0.69                                                                   | 0.71                        | 0.71             |
+| **LUAD-EGFR-mut**        | 0.76                             | 0.80                                    | 0.73                                                         | 0.73                                                                   | 0.79                        | 0.81             |
+| **LUAD-KRAS-mut**        | 0.61                             | 0.65                                    | 0.61                                                         | 0.66                                                                   | 0.60                        | 0.63             |
+| **LUAD-Gene-overexp[1]**    | 0.75                             | 0.68                                    | 0.71                                                         | 0.71                                                                   | 0.74                        | 0.72             |
+| **CRC-MSS/MSI**              | 0.89                             | 0.88                                    | 0.86                                                         | 0.90                                                                   | 0.90                        | 0.89             |
+| **BRCA-ER_PR_HER2**      | 0.82                             | 0.79                                    | 0.76                                                         | 0.79                                                                   | 0.81                        | 0.77             |
+| **Pan-cancer-Gene-mut[2]**  | 0.79                             | 0.77                                    | 0.73                                                         | 0.74                                                                   | 0.77                        | 0.76             |
+| **Avg. AUC**             | 0.77                             | 0.76                                    | 0.73                                                         | 0.74                                                                   | 0.77                        | 0.76             |
+ 
+[1]: **lung-gene-overexp**: total 11 genes were evaluated: LAG3, CLDN6, CD274, EGFR, ERBB2, ERBB3, CD276, VTCN1, TACSTD2, FOLR1, MET.
+
+[2]: **Pan-cancer-Gene-mut**: total 7 genes were evaluated: TP53, KRAS, ALK, PIK3CA, MET, EGFR, PTEN
+
+## License
+The model is licensed under [EXAONEPath AI Model License Agreement 1.0 - NC](./LICENSE)

config.json

@@ -0,0 +1,6 @@
+{
+    "step_size": 256,
+    "patch_size": 256,
+    "num_sampled_patch": 16384,
+    "architecture": "EXAONEPathV1p5Downstream"
+}

datasets/dataset_WSI.py

@@ -0,0 +1,52 @@
+import os
+import torch
+from openslide import OpenSlide
+from utils.preprocessor import MacenkoNormalizer, preprocessor
+from torch.utils.data import Dataset
+
+
+class WSIPatchDataset(Dataset):
+    def __init__(
+        self,
+        coords,
+        wsi_path,
+        pretrained=False,
+        patch_size=256,
+        patch_level=0,
+        macenko=True,
+        return_coord=False,
+    ):
+        self.pretrained = pretrained
+        self.wsi = OpenSlide(wsi_path)
+        self.patch_size = patch_size
+        self.patch_level = patch_level
+        self.return_coord = return_coord
+
+        if macenko:
+            normalizer = MacenkoNormalizer(
+                target_path=os.path.join(
+                    os.path.dirname(os.path.dirname(os.path.join(__file__))),
+                    "macenko_target",
+                    "macenko_param.pt",
+                )
+            )
+        else:
+            normalizer = None
+
+        self.roi_transforms = preprocessor(pretrained=pretrained, normalizer=normalizer)
+        self.coords = coords
+        self.length = len(self.coords)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx):
+        coord = self.coords[idx]
+        img = self.wsi.read_region(
+            coord, self.patch_level, (self.patch_size, self.patch_size)
+        ).convert("RGB")
+        img = self.roi_transforms(img)
+        if self.return_coord:
+            return img, torch.tensor(coord)
+        else:
+            return img

inference.py

@@ -0,0 +1,52 @@
+import argparse
+import torch
+import os
+
+from models.exaonepath import EXAONEPathV1p5Downstream
+from utils.constants import CLASS_NAMES
+from tokens import HF_TOKEN
+
+
+def infer(model, input_file):
+    print("Processing", input_file, "...")
+    probs = model(input_file)
+    result_str = "Result -- " + " / ".join(
+        [f"{name}: {probs[i].item():.4f}" for i, name in enumerate(CLASS_NAMES)]
+    )
+    print(result_str + "\n")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Inference")
+    parser.add_argument('--svs_path', type=str, default='./samples/wsis/1/1.svs', help="Path to the .svs file")
+    parser.add_argument('--svs_dir', type=str, default='./samples_CRC', help="")
+
+    args = parser.parse_args()
+
+    hf_token = HF_TOKEN
+    # model = EXAONEPathV1p5Downstream.from_pretrained("LGAI-EXAONE/EXAONE-Path-1.5", use_auth_token=hf_token)
+    model = EXAONEPathV1p5Downstream(num_sampled_patch=16384)
+
+
+    # qwe = torch.load('./pytorch_model_ori.bin')
+    # aaa = model.load_state_dict(qwe, strict=False)
+    # hw_w = torch.load('/mnt/shared/shared_medical/shared/hi.choi/MOM/logs_eval_25/closebench2/ours/BIOMARKER_SMC_SMC/CRCSensor/CRCSensor_exaone_mom3_MOM_batch_8_lr0.00003_wd0.1_do0.1/_s100/s_0_checkpoint.pt', map_location='cpu')
+
+    # new_state_dict = {}
+    # for k, v in hw_w.items():
+    #     if k.startswith("_orig_mod."):
+    #         new_k = k.replace("_orig_mod.", "agg_model.", 1)
+    #     else:
+    #         new_k = k
+    #     new_state_dict[new_k] = v
+
+    # load_result = model.load_state_dict(new_state_dict, strict=False)
+    model.load_state_dict(torch.load('./pytorch_model.bin'))
+    model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    model.eval()
+    model.feature_extractor = torch.compile(model.feature_extractor)
+    model.agg_model = torch.compile(model.agg_model)
+    
+    for svs_name in os.listdir(args.svs_dir):
+        infer(model, os.path.join(args.svs_dir, svs_name))
+

models/aggregator.py

@@ -0,0 +1,222 @@
+from typing import Any, Dict, Optional, Tuple, Union
+import torch
+from torch import nn
+from dataclasses import dataclass
+from functools import partial
+
+from models.transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer
+
+
+# @dataclass
+# class VisionCfg:
+#     layers: Union[Tuple[int, int, int, int], int] = 6
+#     width: int = 512
+#     head_width: int = 64
+#     mlp_ratio: float = 4.0
+
+#     ls_init_value: Optional[float] = None  # layer scale initial value
+#     patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+#     no_ln_pre: bool = False  # disable pre transformer LayerNorm
+#     pool_type: str = 'none'
+#     final_ln_after_pool: bool = True  # apply final LayerNorm after pooling
+#     output_tokens: bool = False
+#     act_kwargs: Optional[dict] = None
+#     norm_kwargs: Optional[dict] = None
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 6
+    width: int = 512
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer (overrides pool_type)
+    attn_pooler_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    no_ln_pre: bool = False  # disable pre transformer LayerNorm
+    pos_embed_type: str = 'none'
+    final_ln_after_pool: bool = True  # apply final LayerNorm after pooling
+    pool_type: str = 'none'
+    output_tokens: bool = False
+    act_kwargs: Optional[dict] = None
+    norm_kwargs: Optional[dict] = None
+
+    timm_model_name: Optional[str] = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+    img_embed: bool = False
+    cls_embed: bool = False
+    projection = False
+    use_flex = True
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+        dropout: float = 0.1,
+        num_registers: int = 0,
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    vision_heads = vision_cfg.width // vision_cfg.head_width
+    norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+    if vision_cfg.norm_kwargs:
+        norm_layer = partial(norm_layer, **vision_cfg.norm_kwargs)
+    if vision_cfg.act_kwargs is not None:
+        act_layer = partial(act_layer, **vision_cfg.act_kwargs)
+
+    visual = VisionTransformer(
+        width=vision_cfg.width,
+        layers=vision_cfg.layers,
+        heads=vision_heads,
+        mlp_ratio=vision_cfg.mlp_ratio,
+        ls_init_value=vision_cfg.ls_init_value,
+        output_dim=embed_dim,
+        patch_dropout=vision_cfg.patch_dropout,
+        no_ln_pre=vision_cfg.no_ln_pre,
+        pool_type=vision_cfg.pool_type,
+        final_ln_after_pool=vision_cfg.final_ln_after_pool,
+        act_layer=act_layer,
+        norm_layer=norm_layer,
+        output_tokens=vision_cfg.output_tokens,
+        img_embed = vision_cfg.img_embed,
+        use_flex = True,
+        dropout = dropout,
+        num_registers = num_registers,
+        use_rel_bias =True,
+    )
+
+    return visual
+
+
+class MixedOmicsModel(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            drop_rate: float = 0.25,
+            num_registers: int = 0,
+            *args,
+            **kwargs,
+            ):
+        super().__init__()
+
+        self.drop_prob = drop_rate
+        self.num_registers = num_registers
+
+        vision_cfg.cls_embed = False
+        
+        self.visual = _build_vision_tower(embed_dim,
+                                        vision_cfg, 
+                                        quick_gelu, 
+                                        cast_dtype, 
+                                        dropout=drop_rate, 
+                                        num_registers=0,
+                                        )
+        
+        self.image_proj = nn.Linear(embed_dim, embed_dim)
+        self.image_proj.apply(self.init_weights)
+
+        self.ln_post = LayerNorm(embed_dim)
+            
+
+
+    def init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+    def _check_tensor(self, tensor, name):
+        print(name, " : ", tensor.shape)
+        if torch.isnan(tensor).any():
+            print(tensor.shape)
+            print(f"Tensor {name} contains NaN values.")
+        if torch.isinf(tensor).any():
+            print(tensor.shape)
+            print(f"Tensor {name} contains Inf values.")
+    
+    def forward(
+        self,
+        image,
+        coords=None,
+        im_mask=None,
+        *args,
+        **kwargs,
+    ):
+          
+        ## image embedding
+        image_embeds = self.visual(image.contiguous(), coords=coords.contiguous(), key_padding_mask=None if im_mask is None else (~im_mask.bool()).contiguous())
+        image_embeds = self.ln_post(image_embeds)
+        
+        if im_mask is not None:
+            mask = im_mask.unsqueeze(-1).contiguous()
+            masked_embeds = image_embeds * mask
+            sum_embeds = masked_embeds.sum(dim=1)
+            valid_counts = mask.sum(dim=1).clamp(min=1)   # [N, 1]
+            mean_embeds = sum_embeds / valid_counts       # [N, dim]
+
+        else:
+            mean_embeds = image_embeds.mean(-2)
+
+        image_embeds_final = self.image_proj(mean_embeds)
+        
+        return image_embeds_final, image_embeds, mean_embeds
+
+
+
+def make_model(
+    embed_dim=768,
+    droprate=0.1,
+    num_registers=0,
+    depth=4,
+):  
+    vCfg = CLIPVisionCfg
+    vCfg.width = embed_dim
+    vCfg.layers = depth
+
+    model = MixedOmicsModel(
+        embed_dim=embed_dim,
+        vision_cfg=vCfg,
+        drop_rate=droprate,
+        num_registers=num_registers,
+        )
+    
+    return model

models/cls_modules.py

@@ -0,0 +1,168 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class GlobalClassificationHead(nn.Module):
+    def __init__(self, input_dim, num_classes=32, dropout_rate=0.1):
+        super().__init__()
+        self.norm = nn.LayerNorm(input_dim)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.fc = nn.Linear(input_dim, num_classes)
+    
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.dropout(x)
+        logits = self.fc(x)  # (batch_size, num_classes)
+        return logits
+    
+    
+class CLSHead(nn.Module):
+    def __init__(
+        self, 
+        embed_dim, 
+        num_classes, 
+        dropout=0.1, 
+        use_norm=True, 
+        hidden_dim=None,
+        activation="silu",
+        pooling_type="cls",
+    ):
+        super().__init__()
+        
+        if hidden_dim is None:
+            hidden_dim = embed_dim
+        
+        if activation == "gelu":
+            self.activation = nn.GELU()
+        elif activation == "relu":
+            self.activation = nn.ReLU(inplace=True)
+        elif activation == "silu":
+            self.activation = nn.SiLU(inplace=True)
+        else:
+            raise ValueError(f"Value Error: {activation}")
+        
+        self.pooling_type = pooling_type
+        
+        if pooling_type == "attention":
+            self.attention_pool = nn.Sequential(
+                nn.LayerNorm(embed_dim),
+                nn.Linear(embed_dim, 1),
+                # nn.Softmax(dim=1)
+            )
+        
+        if use_norm:
+            self.norm = nn.LayerNorm(embed_dim)
+        else:
+            self.norm = nn.Identity()
+            
+        self.fc1 = nn.Linear(embed_dim, hidden_dim)
+        self.dropout1 = nn.Dropout(dropout)
+        self.fc2 = nn.Linear(hidden_dim, num_classes)
+        
+        self._init_weights()
+    
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.fc1.weight, std=0.02)
+        nn.init.zeros_(self.fc1.bias)
+        nn.init.trunc_normal_(self.fc2.weight, std=0.02)
+        nn.init.zeros_(self.fc2.bias)
+        
+        if self.pooling_type == "attention":
+            nn.init.trunc_normal_(self.attention_pool[0].weight, std=0.02)
+            nn.init.zeros_(self.attention_pool[0].bias)
+    
+    def forward(self, x, x2=None, x3=None, attn_mask=None):
+        if self.pooling_type == "mlp":
+            pooled = x + x3
+        elif self.pooling_type == "attention":
+            x = torch.cat([x.unsqueeze(1), x2], dim=1)
+            weights = self.attention_pool(x)  # [batch_size, num_tokens, 1]
+            
+            if attn_mask is not None:
+                attn_mask = torch.cat([torch.zeros(attn_mask.shape[0], 1).to(attn_mask.dtype).to(attn_mask.device), attn_mask], dim=1)
+                new_attn_mask = torch.zeros_like(attn_mask, dtype=weights.dtype)
+                new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+                attn_mask = new_attn_mask
+                
+            if len(attn_mask.shape) ==2:
+                attn_mask = attn_mask[..., None] # [batch_size, num_tokens, 1]
+                
+            weights = weights + attn_mask
+            weights = F.softmax(weights, dim=1)
+            
+            pooled = torch.sum(x * weights, dim=1)
+        else:
+            raise ValueError(f"지원하지 않는 풀링 타입: {self.pooling_type}")
+
+        pooled = self.norm(pooled)
+        
+        x = self.fc1(pooled)
+        x = self.activation(x)
+        x = self.dropout1(x)
+        x = self.fc2(x)
+        
+        return x
+
+
+class BaseClassifier(nn.Module):
+    def __init__(
+        self,
+        base_model,
+        num_classes,
+        cls_head_kwargs=None
+    ):
+        super().__init__()
+        
+        self.backbone = base_model
+        embed_dim = self.backbone.embed_dim if hasattr(self.backbone, 'embed_dim') else 768
+        
+        cls_head_config = {
+            'embed_dim': embed_dim,
+            'num_classes': num_classes,
+        }
+
+        if cls_head_kwargs:
+            cls_head_config.update(cls_head_kwargs)
+        
+        self.cls_head = CLSHead(**cls_head_config)
+    
+    def forward(self, image, coords=None, im_mask=None):
+        feat_final, feats, m_feats = self.backbone(image=image, coords=coords, im_mask=im_mask)
+        logits = self.cls_head(feat_final, feats, m_feats, attn_mask=~im_mask.bool())
+        
+        Y_hat = torch.topk(logits, 1, dim=1)[1]
+        Y_prob = F.softmax(logits, dim=-1)
+        
+        return logits, Y_prob, Y_hat
+
+
+class LinearClassifier(nn.Module):
+    def __init__(
+        self, 
+        base_model,
+        num_classes=2,
+        pool="final"  
+    ):
+        super().__init__()
+        self.backbone = base_model
+        self.pool = pool
+        print("Linear classifier pool type : ", self.pool)
+        embed_dim = self.backbone.embed_dim if hasattr(self.backbone, 'embed_dim') else 768
+        
+        self.cls_head = nn.Sequential(
+                    nn.LayerNorm(embed_dim),
+                    nn.Linear(embed_dim, num_classes)
+        )
+    
+    def forward(self, image, coords=None, im_mask=None):
+        feat_final, _, feat_mean = self.backbone(image=image, coords=coords, im_mask=im_mask)
+        if self.pool == "mean":
+            logits = self.cls_head(feat_mean)
+        elif self.pool == "final":
+            logits = self.cls_head(feat_final)
+        
+        Y_hat = torch.topk(logits, 1, dim=1)[1]
+        Y_prob = F.softmax(logits, dim=-1)
+        
+        return logits, Y_prob, Y_hat

models/exaonepath.py

@@ -0,0 +1,158 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from models.aggregator import make_model
+from models.cls_modules import LinearClassifier
+from datasets.dataset_WSI import WSIPatchDataset
+from models.feature_extractor import vit_base
+from utils.wsi_utils import extract_tissue_patch_coords
+from torch.utils.data import DataLoader
+
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class EXAONEPathV1p5Downstream(nn.Module, PyTorchModelHubMixin):
+    def __init__(
+        self, step_size=256, patch_size=256, num_sampled_patch=999999, macenko=True
+    ):
+        super(EXAONEPathV1p5Downstream, self).__init__()
+        self.step_size = step_size
+        self.patch_size = patch_size
+        self.macenko = macenko
+        self.num_sampled_patch = num_sampled_patch
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.config = {
+            "step_size": step_size,
+            "patch_size": patch_size,
+            "macenko": macenko,
+            "num_sampled_patch": num_sampled_patch,
+        }
+
+        self.feature_extractor = vit_base()
+        self.feature_extractor = self.feature_extractor
+        # self.feature_extractor = self.feature_extractor.to(self.device)
+        # self.feature_extractor.eval()
+
+        self.agg_model = make_model(
+            embed_dim=768,
+            droprate=0.0,
+            num_registers=0,
+            depth=4,
+            )
+        
+        self.agg_model = LinearClassifier(self.agg_model, pool='mean')
+        
+        # self.agg_model.to(self.device)
+        # self.agg_model.eval()
+
+    @torch.no_grad()
+    def forward(self, svs_path: str, feature_extractor_batch_size: int = 8):
+        # Extract patches
+        coords = extract_tissue_patch_coords(
+            svs_path, patch_size=self.patch_size, step_size=self.step_size
+        )
+
+        # Extract patch-level features
+        self.feature_extractor.eval()
+        patch_dataset = WSIPatchDataset(
+            coords=coords,
+            wsi_path=svs_path,
+            pretrained=True,
+            macenko=self.macenko,
+            patch_size=self.patch_size,
+            return_coord=True,
+        )
+        patch_loader = DataLoader(
+            dataset=patch_dataset,
+            batch_size=feature_extractor_batch_size,
+            num_workers=(
+                feature_extractor_batch_size * 2 if self.device.type == "cuda" else 0
+            ),
+            pin_memory=self.device.type == "cuda",
+
+        )
+        features_list = []
+        coords_list = []
+        for count, items in enumerate(patch_loader):
+            patches, coords = items
+            print(
+                f"batch {count+1}/{len(patch_loader)}, {count * feature_extractor_batch_size} patches processed",
+                end="\r",
+            )
+            patches = patches.to(self.device, non_blocking=True)
+
+            feature = self.feature_extractor(patches)  # [B, 1024]
+            feature /= feature.norm(dim=-1, keepdim=True)  # use normalized featuren
+            feature = feature.to("cpu", non_blocking=True)
+            features_list.append(feature)
+
+            coords = coords.to(self.device, non_blocking=True)
+            coords_list.append(coords)
+
+        print("")
+        print("Feature extraction finished")
+
+        features = torch.cat(features_list)
+        coords = torch.cat(coords_list)
+        total_samples = features.shape[0]
+
+        num_samples = min(self.num_sampled_patch, total_samples)
+        indices = torch.randperm(total_samples)[:num_samples]
+        sampled_features = features[indices]
+        sampled_coords = coords[indices]
+        
+        # Aggregate features
+        self.agg_model.eval()
+
+        # sampled_features = torch.randn([8192, 768])
+        # sampled_coords = torch.randn([8192, 2])
+
+        logits, Y_prob, Y_hat = self.agg_model(sampled_features[None].to(self.device), sampled_coords[None].to(self.device))
+        probs = Y_prob[0].cpu()
+
+        return probs
+
+    @torch.no_grad()
+    def forward_feature_extractor(self, svs_path: str, feature_extractor_batch_size: int = 8):
+        # Extract patches
+        coords = extract_tissue_patch_coords(
+            svs_path, patch_size=self.patch_size, step_size=self.step_size
+        )
+
+        # Extract patch-level features
+        self.feature_extractor.eval()
+        patch_dataset = WSIPatchDataset(
+            coords=coords,
+            wsi_path=svs_path,
+            pretrained=True,
+            macenko=self.macenko,
+            patch_size=self.patch_size,
+            return_coord=False
+        )
+        patch_loader = DataLoader(
+            dataset=patch_dataset,
+            batch_size=feature_extractor_batch_size,
+            num_workers=(
+                feature_extractor_batch_size * 2 if self.device.type == "cuda" else 0
+            ),
+            pin_memory=self.device.type == "cuda",
+        )
+        features_list = []
+        for count, patches in enumerate(patch_loader):
+            print(
+                f"batch {count+1}/{len(patch_loader)}, {count * feature_extractor_batch_size} patches processed",
+                end="\r",
+            )
+            patches = patches.to(self.device, non_blocking=True)
+
+            feature = self.feature_extractor(patches)  # [B, 1024]
+            feature /= feature.norm(dim=-1, keepdim=True)  # use normalized featuren
+            feature = feature.to("cpu", non_blocking=True)
+            features_list.append(feature)
+        print("")
+        print("Feature extraction finished")
+
+        features = torch.cat(features_list)
+
+        return features

models/feature_extractor.py

@@ -0,0 +1,357 @@
+import math
+import warnings
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return _drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=[224],
+        patch_size=16,
+        in_chans=3,
+        num_classes=0,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        **kwargs
+    ):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0],
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        self.head = (
+            nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
+
+        _trunc_normal_(self.pos_embed, std=0.02)
+        _trunc_normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            _trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(
+                1, int(math.sqrt(N)), int(math.sqrt(N)), dim
+            ).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            # size=(int(w0), int(h0)),
+            mode="bicubic",
+        )
+        assert (
+            int(w0) == patch_pos_embed.shape[-2]
+            and int(h0) == patch_pos_embed.shape[-1]
+        )
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        # print(x.type())
+        return x[:, 0]
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def _trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    # type: (torch.Tensor, float, float, float, float) -> torch.Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def _drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output

models/flex_attn.py

@@ -0,0 +1,124 @@
+import math
+from typing import Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import (
+    _DEFAULT_SPARSE_BLOCK_SIZE,
+    create_block_mask,
+    create_mask,
+    flex_attention,
+)
+from torch.nn.attention.flex_attention import flex_attention, _vmap_for_bhqkv
+
+try:
+    from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
+except ImportError:
+    from torch._higher_order_ops.flex_attention import TransformGetItemToIndex
+    
+from torch._dynamo import disable
+
+
+def generate_alibi_bias(H=12):
+    alibi_bias = []
+    for h in range(H):
+        alibi_bias.append(-((h + 1) / H))
+    alibi_bias = torch.tensor(alibi_bias)
+    alibi_bias = torch.exp2(alibi_bias)
+    return alibi_bias
+
+
+def get_rel_bias_func(scale, coords=None, qk_scale=1.0):
+    def patch_coords_rel_bias(score, b, h, q_idx, kv_idx):
+        if coords is None:
+            return score
+        with torch.no_grad():
+            dx = coords[b, q_idx][0] - coords[b, kv_idx][0]
+            dy = coords[b, q_idx][1] - coords[b, kv_idx][1]
+            dist = torch.sqrt(dx * dx + dy * dy)
+            dist = dist.clamp(max=1000)  # max distance
+            dist = torch.log1p(dist)
+            bias = dist * scale[h] * qk_scale
+        return score - bias  # closer → larger score
+    return patch_coords_rel_bias
+
+
+def key_padding_mask(mask):
+    def padding_mask(b, h, q_idx, kv_idx):
+        return ~mask[b, kv_idx]
+    return padding_mask
+
+
+class FlexCore(nn.Module):
+    """
+        For using "forward hook"
+    """
+    def forward(self, q, k, v, score_mod=None, block_mask=None, return_lse=False):
+        """
+            "return_lse=True" should be used with ATTN_MAP_VIS wrapper.
+            Though return_lse is "True", _flex_attention(...) only have an output (attention output, not attention scores).
+        """
+        return flex_attention(q, k, v, score_mod=score_mod, block_mask=block_mask, return_lse=return_lse)
+
+
+class Flex_Attention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 12,
+            qkv_bias: bool = True,
+            proj_drop: float = 0.,
+            use_rel_bias: bool = True,
+    ):
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.scale = self.head_dim ** -0.5
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+            
+        self.f_attn = FlexCore()
+
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+        
+        self.max_distance=16
+    
+    def build_rel_bias(self, coords):
+        return torch.log1p(torch.cdist(coords, coords, p=2))
+
+    def forward(self, x, coords=None, attn_mask: Optional[torch.Tensor] = None, return_attn_score=False):
+        N, L, C = x.shape
+        # ensure contiguous projection before chunking
+        x_proj = F.linear(x, self.in_proj_weight, self.in_proj_bias).contiguous()
+        q, k, v = [t.contiguous() for t in x_proj.chunk(3, -1)]
+
+        q = q.reshape(N, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
+        k = k.reshape(N, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
+        v = v.reshape(N, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
+        
+        if attn_mask is not None:
+            maks_func = create_block_mask(
+                key_padding_mask(attn_mask), N, self.num_heads, L, L
+            )
+        
+        qk_scale = q.size(-1) ** -0.5
+        x = self.f_attn(
+                q, k, v,
+                score_mod = get_rel_bias_func(generate_alibi_bias(self.num_heads).to(coords.device), coords, qk_scale) if coords is not None else None,
+                block_mask = maks_func if attn_mask is not None else None,
+                return_lse=return_attn_score,
+        )
+        
+        x = x.permute(0, 2, 1, 3).contiguous()
+        x = x.reshape(N, L, C).contiguous()
+
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x

models/transformer.py

@@ -0,0 +1,424 @@
+from collections import OrderedDict
+import math
+from typing import Callable, List, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from einops import pack, repeat
+
+from .flex_attn import Flex_Attention
+
+
+
+class LayerNormFp32(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class QuickGELU(nn.Module):
+    # NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = True,
+            scaled_cosine: bool = True,
+            scale_heads: bool = False,
+            logit_scale_max: float = math.log(1. / 0.01),
+            batch_first: bool = True,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.
+    ):
+        super().__init__()
+        self.scaled_cosine = scaled_cosine
+        self.scale_heads = scale_heads
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.logit_scale_max = logit_scale_max
+        self.batch_first = batch_first
+        self.use_fsdpa = hasattr(nn.functional, 'scaled_dot_product_attention')
+
+        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+
+        if self.scaled_cosine:
+            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
+        else:
+            self.logit_scale = None
+        self.attn_drop = nn.Dropout(attn_drop)
+        if self.scale_heads:
+            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
+        else:
+            self.head_scale = None
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+        
+
+    def forward(self, x, coords, attn_mask: Optional[torch.Tensor] = None):
+        if self.batch_first:
+            x = x.transpose(0, 1)
+
+        L, N, C = x.shape
+        q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
+        q = q.reshape(L, N * self.num_heads, -1).transpose(0, 1)
+        k = k.reshape(L, N * self.num_heads, -1).transpose(0, 1)
+        v = v.reshape(L, N * self.num_heads, -1).transpose(0, 1)
+
+        if attn_mask is not None and attn_mask.dtype == torch.bool:
+            new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+            new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+            attn_mask = new_attn_mask
+
+        # if self.logit_scale is not None:
+        attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
+        logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
+        attn = attn.view(N, self.num_heads, L, L) * logit_scale
+
+        if attn_mask is not None:
+            attn = attn + attn_mask[:, None, None, :]
+        attn = attn.view(-1, L, L)
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        
+        x = torch.bmm(attn, v)
+        
+        if self.head_scale is not None:
+            x = x.view(N, self.num_heads, L, C) * self.head_scale
+            x = x.view(-1, L, C)
+
+        x = x.transpose(0, 1).reshape(L, N, C)
+
+        if self.batch_first:
+            x = x.transpose(0, 1)
+
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x
+
+
+class AttentionalPooler(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            context_dim: int,
+            n_head: int = 8,
+            n_queries: int = 256,
+            norm_layer: Callable = LayerNorm,
+    ):
+        super().__init__()
+        self.query = nn.Parameter(torch.randn(n_queries, d_model))
+        self.attn = nn.MultiheadAttention(d_model, n_head, kdim=context_dim, vdim=context_dim, batch_first=True)
+        self.ln_q = norm_layer(d_model)
+        self.ln_k = norm_layer(context_dim)
+
+    def forward(self, x: torch.Tensor):
+        N = x.shape[0]
+        x = self.ln_k(x)
+        q = self.ln_q(self.query)
+        out = self.attn(q.unsqueeze(0).expand(N, -1, -1), x, x, need_weights=False)[0]
+        return out
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            is_cross_attention: bool = False,
+            batch_first: bool = True,
+            use_flex:bool = False,
+            dropout:float = 0.2,
+            use_rel_bias:bool = True,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        
+        if use_flex:
+            print("Flex_Attention!")
+            self.attn = Flex_Attention(dim = d_model, num_heads=n_head, proj_drop=dropout, use_rel_bias=use_rel_bias)
+        else:
+            self.attn = Attention(dim = d_model, num_heads=n_head, batch_first=batch_first, proj_drop=dropout, attn_drop=dropout)
+            
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+        if is_cross_attention:
+            self.ln_1_kv = norm_layer(d_model)
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ("gelu", act_layer()),
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def attention(
+            self,
+            q_x: torch.Tensor,
+            k_x: Optional[torch.Tensor] = None,
+            v_x: Optional[torch.Tensor] = None,
+            coords = None,
+            attn_mask: Optional[torch.Tensor] = None,
+            key_padding_mask=None,
+    ):
+        k_x = k_x if k_x is not None else q_x
+        v_x = v_x if v_x is not None else q_x
+
+        attn_mask = attn_mask.to(q_x.dtype) if attn_mask is not None else None
+
+        return self.attn(
+            q_x, coords=coords, attn_mask=key_padding_mask
+        )
+
+    def forward(
+            self,
+            q_x: torch.Tensor,
+            k_x: Optional[torch.Tensor] = None,
+            v_x: Optional[torch.Tensor] = None,
+            coords = None,
+            attn_mask: Optional[torch.Tensor] = None,
+            key_padding_mask = None,
+    ):
+        k_x = self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None
+        v_x = self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None
+        x = q_x + self.ls_1(self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, coords=coords, attn_mask=attn_mask, key_padding_mask=key_padding_mask))
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+
+def _expand_token(token, batch_size: int):
+    return token.view(1, 1, -1).expand(batch_size, -1, -1)
+
+
+class Transformer(nn.Module):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            batch_first: bool = True,
+            use_flex: bool = False,
+            dropout: float = False,
+            use_rel_bias: bool = True,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.batch_first = batch_first
+        self.grad_checkpointing = False
+
+        self.resblocks = nn.ModuleList([
+            ResidualAttentionBlock(
+                width,
+                heads,
+                mlp_ratio,
+                ls_init_value=ls_init_value,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                batch_first=batch_first,
+                use_flex=use_flex,
+                dropout=dropout,
+                use_rel_bias=use_rel_bias
+            )
+            for _ in range(layers)
+        ])
+
+    def get_cast_dtype(self) -> torch.dtype:
+        if hasattr(self.resblocks[0].mlp.c_fc, 'int8_original_dtype'):
+            return self.resblocks[0].mlp.c_fc.int8_original_dtype
+        return self.resblocks[0].mlp.c_fc.weight.dtype
+
+    def forward(self, x: torch.Tensor, coords = None, attn_mask: Optional[torch.Tensor] = None, key_padding_mask=None):
+        if not self.batch_first:
+            x = x.transpose(0, 1).contiguous()    # NLD -> LND
+        for r in self.resblocks:
+            x = r(x, attn_mask=attn_mask, key_padding_mask=key_padding_mask, coords=coords)
+        if not self.batch_first:
+            x = x.transpose(0, 1).contiguous()    # LND -> NLD
+        return x
+
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float,
+            ls_init_value: float = None,
+            output_dim: int = 512,
+            patch_dropout: float = 0.,
+            no_ln_pre: bool = False,
+            pool_type: str = 'tok',
+            final_ln_after_pool: bool = False,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            output_tokens: bool = False,
+            img_embed: bool = False,
+            use_flex:bool = False,
+            dropout:float = 0.1,
+            num_registers: int = 0,
+            use_rel_bias: bool = True,
+    ):
+        super().__init__()
+        assert pool_type in ('tok', 'avg', 'none')
+        self.output_tokens = output_tokens
+        
+        self.final_ln_after_pool = final_ln_after_pool  # currently ignored w/ attn pool enabled
+        self.output_dim = output_dim
+        self.img_embed = img_embed
+        self.num_registers = num_registers
+        self.positional_embedding = None
+        self.pre_linear = nn.Linear(768, width)
+
+        
+        if num_registers>0:
+            self.register_token = nn.Parameter(torch.empty(num_registers, width))
+            nn.init.normal_(self.register_token, std=0.02)
+            
+            
+        self.positional_embedding = None
+        
+
+        self.positional_embedding = None
+        
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.ln_pre = nn.Identity() if no_ln_pre else norm_layer(width)
+        self.transformer = Transformer(
+            width,
+            layers,
+            heads,
+            mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            use_flex=use_flex,
+            dropout=dropout,
+            use_rel_bias=use_rel_bias,
+        )
+
+        pool_dim = width
+        self.pool_type = pool_type
+
+        self.ln_post = norm_layer(pool_dim)
+
+    def _global_pool(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.pool_type == 'avg':
+            pooled, tokens = x[:, 1:].mean(dim=1), x[:, 1:]
+        elif self.pool_type == 'tok':
+            pooled, tokens = x[:, 0], x[:, 1:]
+        else:
+            pooled = tokens = x
+
+        return pooled, tokens
+
+    def forward(self, x: torch.Tensor, coords=None, mask=None, key_padding_mask=None):
+        x = self.pre_linear(x)
+        
+        if self.num_registers > 0:
+            r = repeat(self.register_token, 'n d -> b n d', b=x.size(0))
+            x, ps = pack([x, r], 'b * d')
+        
+        x = self.patch_dropout(x)
+        x = self.ln_pre(x)
+        x = self.transformer(x, coords, mask, key_padding_mask=key_padding_mask)
+
+        if self.final_ln_after_pool:
+            pooled, tokens = self._global_pool(x)
+            pooled = self.ln_post(pooled)
+        else:
+            x = self.ln_post(x)
+            pooled, tokens = self._global_pool(x)
+
+        return pooled
+
+    

requirements.txt

@@ -0,0 +1,16 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+
+torch==2.6.0+cu124
+torchvision==0.21.0+cu124
+numpy==2.1.2
+opencv-python-headless==4.11.0.86
+openslide-bin==4.0.0.6
+openslide-python==1.4.1
+pandas==2.2.3
+torchstain==1.4.1
+fire==0.7.0
+einops==0.8.1
+huggingface_hub==0.32.2
+h5py==3.13.0
+scikit-learn==1.6.1
+tensorboard==2.19.0

utils/constants.py

@@ -0,0 +1 @@
+CLASS_NAMES = ["LABEL_LOW", "LABEL_HIGH"]

utils/file_utils.py

@@ -0,0 +1,35 @@
+import pickle
+import h5py
+
+def save_pkl(filename, save_object):
+	writer = open(filename,'wb')
+	pickle.dump(save_object, writer)
+	writer.close()
+
+def load_pkl(filename):
+	loader = open(filename,'rb')
+	file = pickle.load(loader)
+	loader.close()
+	return file
+
+
+def save_hdf5(output_path, asset_dict, attr_dict= None, mode='a'):
+    file = h5py.File(output_path, mode)
+    for key, val in asset_dict.items():
+        data_shape = val.shape
+        if key not in file:
+            data_type = val.dtype
+            chunk_shape = (1, ) + data_shape[1:]
+            maxshape = (None, ) + data_shape[1:]
+            dset = file.create_dataset(key, shape=data_shape, maxshape=maxshape, chunks=chunk_shape, dtype=data_type)
+            dset[:] = val
+            if attr_dict is not None:
+                if key in attr_dict.keys():
+                    for attr_key, attr_val in attr_dict[key].items():
+                        dset.attrs[attr_key] = attr_val
+        else:
+            dset = file[key]
+            dset.resize(len(dset) + data_shape[0], axis=0)
+            dset[-data_shape[0]:] = val
+    file.close()
+    return output_path

utils/preprocessor.py

@@ -0,0 +1,218 @@
+import os
+import random
+
+import torch
+from PIL import Image
+from torchstain.base.normalizers.he_normalizer import HENormalizer
+from torchstain.torch.utils import cov, percentile
+from torchvision import transforms
+from torchvision.transforms.functional import to_pil_image
+
+
+def preprocessor(pretrained=False, normalizer=None):
+    if pretrained:
+        mean = (0.485, 0.456, 0.406)
+        std = (0.229, 0.224, 0.225)
+    else:
+        mean = (0.5, 0.5, 0.5)
+        std = (0.5, 0.5, 0.5)
+
+    preprocess = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.Lambda(lambda x: x) if normalizer == None else normalizer,
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean, std=std),
+        ]
+    )
+
+    return preprocess
+
+
+"""
+Source code ported from: https://github.com/schaugf/HEnorm_python
+Original implementation: https://github.com/mitkovetta/staining-normalization
+"""
+
+
+class TorchMacenkoNormalizer(HENormalizer):
+    def __init__(self):
+        super().__init__()
+
+        self.HERef = torch.tensor(
+            [[0.5626, 0.2159], [0.7201, 0.8012], [0.4062, 0.5581]]
+        )
+        self.maxCRef = torch.tensor([1.9705, 1.0308])
+
+        # Avoid using deprecated torch.lstsq (since 1.9.0)
+        self.updated_lstsq = hasattr(torch.linalg, "lstsq")
+
+    def __convert_rgb2od(self, I, Io, beta):
+        I = I.permute(1, 2, 0)
+
+        # calculate optical density
+        OD = -torch.log((I.reshape((-1, I.shape[-1])).float() + 1) / Io)
+
+        # remove transparent pixels
+        ODhat = OD[~torch.any(OD < beta, dim=1)]
+
+        return OD, ODhat
+
+    def __find_HE(self, ODhat, eigvecs, alpha):
+        # project on the plane spanned by the eigenvectors corresponding to the two
+        # largest eigenvalues
+        That = torch.matmul(ODhat, eigvecs)
+        phi = torch.atan2(That[:, 1], That[:, 0])
+        # print(phi.size())
+
+        minPhi = percentile(phi, alpha)
+        maxPhi = percentile(phi, 100 - alpha)
+
+        vMin = torch.matmul(
+            eigvecs, torch.stack((torch.cos(minPhi), torch.sin(minPhi)))
+        ).unsqueeze(1)
+        vMax = torch.matmul(
+            eigvecs, torch.stack((torch.cos(maxPhi), torch.sin(maxPhi)))
+        ).unsqueeze(1)
+
+        # a heuristic to make the vector corresponding to hematoxylin first and the
+        # one corresponding to eosin second
+        HE = torch.where(
+            vMin[0] > vMax[0],
+            torch.cat((vMin, vMax), dim=1),
+            torch.cat((vMax, vMin), dim=1),
+        )
+
+        return HE
+
+    def __find_concentration(self, OD, HE):
+        # rows correspond to channels (RGB), columns to OD values
+        Y = OD.T
+
+        # determine concentrations of the individual stains
+        if not self.updated_lstsq:
+            return torch.lstsq(Y, HE)[0][:2]
+
+        return torch.linalg.lstsq(HE, Y)[0]
+
+    def __compute_matrices(self, I, Io, alpha, beta):
+        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)
+
+        # compute eigenvectors
+        _, eigvecs = torch.linalg.eigh(cov(ODhat.T))
+        eigvecs = eigvecs[:, [1, 2]]
+
+        HE = self.__find_HE(ODhat, eigvecs, alpha)
+
+        C = self.__find_concentration(OD, HE)
+        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])
+
+        return HE, C, maxC
+
+    def fit(self, I, Io=240, alpha=1, beta=0.15):
+        HE, _, maxC = self.__compute_matrices(I, Io, alpha, beta)
+
+        self.HERef = HE
+        self.maxCRef = maxC
+
+    def normalize(
+        self, I, Io=240, alpha=1, beta=0.15, stains=True, form="chw", dtype="int"
+    ):
+        """Normalize staining appearence of H&E stained images
+
+        Example use:
+            see test.py
+
+        Input:
+            I: RGB input image: tensor of shape [C, H, W] and type uint8
+            Io: (optional) transmitted light intensity
+            alpha: percentile
+            beta: transparency threshold
+            stains: if true, return also H & E components
+
+        Output:
+            Inorm: normalized image
+            H: hematoxylin image
+            E: eosin image
+
+        Reference:
+            A method for normalizing histology slides for quantitative analysis. M.
+            Macenko et al., ISBI 2009
+        """
+
+        c, h, w = I.shape
+
+        HE, C, maxC = self.__compute_matrices(I, Io, alpha, beta)
+
+        # normalize stain concentrations
+        C *= (self.maxCRef / maxC).unsqueeze(-1)
+
+        # recreate the image using reference mixing matrix
+        Inorm = Io * torch.exp(-torch.matmul(self.HERef, C))
+        Inorm = torch.clip(Inorm, 0, 255)
+
+        Inorm = Inorm.reshape(c, h, w).float() / 255.0
+        Inorm = torch.clip(Inorm, 0.0, 1.0)
+
+        H, E = None, None
+
+        if stains:
+            H = torch.mul(
+                Io,
+                torch.exp(
+                    torch.matmul(-self.HERef[:, 0].unsqueeze(-1), C[0, :].unsqueeze(0))
+                ),
+            )
+            H[H > 255] = 255
+            H = H.T.reshape(h, w, c).int()
+
+            E = torch.mul(
+                Io,
+                torch.exp(
+                    torch.matmul(-self.HERef[:, 1].unsqueeze(-1), C[1, :].unsqueeze(0))
+                ),
+            )
+            E[E > 255] = 255
+            E = E.T.reshape(h, w, c).int()
+
+        return Inorm, H, E
+
+
+class MacenkoNormalizer:
+    def __init__(self, target_path=None, prob=1):
+        self.transform_before_macenko = transforms.Compose(
+            [transforms.ToTensor(), transforms.Lambda(lambda x: x * 255)]
+        )
+        self.normalizer = TorchMacenkoNormalizer()
+
+        ext = os.path.splitext(target_path)[1].lower()
+        if ext in [".jpg", ".jpeg", ".png"]:
+            target = Image.open(target_path)
+            self.normalizer.fit(self.transform_before_macenko(target))
+        elif ext in [".pt"]:
+            target = torch.load(target_path)
+            self.normalizer.HERef = target["HERef"]
+            self.normalizer.maxCRef = target["maxCRef"]
+
+        else:
+            raise ValueError(f"Invalid extension: {ext}")
+        self.prob = prob
+
+    def __call__(self, image):
+        t_to_transform = self.transform_before_macenko(image)
+        try:
+            image_macenko, _, _ = self.normalizer.normalize(
+                I=t_to_transform, stains=False, form="chw", dtype="float"
+            )
+            if torch.any(torch.isnan(image_macenko)):
+                return image
+            else:
+                image_macenko = to_pil_image(image_macenko)
+                return image_macenko
+        except Exception as e:
+            if "kthvalue()" in str(e) or "linalg.eigh" in str(e):
+                pass
+            else:
+                print(str(e))
+            return image

utils/wsi_utils.py

@@ -0,0 +1,219 @@
+import cv2
+import numpy as np
+import torch
+from openslide import OpenSlide
+
+
+def extract_tissue_patch_coords(
+    wsi_path: str,
+    patch_size: int = 256,
+    step_size: int = 256,
+    downsample_threshold: float = 64,
+    threshold: int = 8,
+    max_val: int = 255,
+    median_kernel: int = 7,
+    close_size: int = 4,
+    min_effective_area_factor: float = 100,  # multiplied by (ref_area)^2
+    ref_area: int = 512,
+    min_hole_area_factor: float = 16,  # multiplied by (ref_area)^2
+    max_n_holes: int = 8,
+)-> list:
+    """
+    Extract patches from the full-resolution image whose centers fall within tissue regions.
+
+    Process:
+      1. Open the WSI.
+      2. Select a segmentation level and compute a binary mask.
+      3. Find contours and holes from the mask and filter them using effective area criteria.
+      4. Scale the external contours and holes to full resolution.
+      5. Slide a window over the full-resolution image and extract patches if the center is in tissue.
+
+    Returns:
+        A torch tensor of shape (N, 3, patch_size, patch_size) containing the patches.
+    """
+    slide = OpenSlide(wsi_path)
+    full_width, full_height = slide.level_dimensions[0]
+
+    seg_level, scale = select_segmentation_level(slide, downsample_threshold)
+    binary_mask = compute_segmentation_mask(
+        slide, seg_level, threshold, max_val, median_kernel, close_size
+    )
+
+    # Compute thresholds for effective area and hole area
+    effective_area_thresh = min_effective_area_factor * (
+        ref_area**2 / (scale[0] * scale[1])
+    )
+    hole_area_thresh = min_hole_area_factor * (ref_area**2 / (scale[0] * scale[1]))
+
+    ext_contours, holes_list = filter_contours_and_holes(
+        binary_mask, effective_area_thresh, hole_area_thresh, max_n_holes
+    )
+    if not ext_contours:
+        raise ValueError("No valid tissue contours found.")
+
+    tissue_contours = scale_contours(ext_contours, scale)
+    scaled_holes = [scale_contours(holes, scale) for holes in holes_list]
+
+    coords = []
+    for y in range(0, full_height - patch_size + 1, step_size):
+        for x in range(0, full_width - patch_size + 1, step_size):
+            center_x = x + patch_size // 2
+            center_y = y + patch_size // 2
+            if not point_in_tissue(center_x, center_y, tissue_contours, scaled_holes):
+                continue
+            coords.append((x, y))
+
+    if not coords:
+        raise ValueError("No available patches")
+    return coords
+
+
+def select_segmentation_level(slide: OpenSlide, downsample_threshold: float = 64):
+    """
+    Select a segmentation level whose downsample factor is at least the specified threshold.
+
+    Returns:
+        level (int): Chosen level index.
+        scale (tuple): Downsample factors (sx, sy) for that level.
+    """
+    level = slide.get_best_level_for_downsample(downsample_threshold)
+    ds = slide.level_downsamples[level]
+    if not isinstance(ds, (tuple, list)):
+        ds = (ds, ds)
+    return level, ds
+
+
+def compute_segmentation_mask(
+    slide: OpenSlide,
+    level: int,
+    threshold: int = 20,
+    max_val: int = 255,
+    median_kernel: int = 7,
+    close_size: int = 4,
+):
+    """
+    Compute a binary mask for tissue segmentation at the specified level.
+
+    Process:
+      - Read the image at the given level and convert to RGB.
+      - Convert the image to HSV and extract the saturation channel.
+      - Apply median blur.
+      - Apply binary thresholding (either fixed or Otsu).
+      - Apply morphological closing.
+
+    Returns:
+        binary (ndarray): Binary mask image.
+    """
+    img = np.array(
+        slide.read_region((0, 0), level, slide.level_dimensions[level]).convert("RGB")
+    )
+    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    sat = hsv[:, :, 1]
+    blurred = cv2.medianBlur(sat, median_kernel)
+    _, binary = cv2.threshold(blurred, threshold, max_val, cv2.THRESH_BINARY)
+    if close_size > 0:
+        kernel = np.ones((close_size, close_size), np.uint8)
+        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
+    return binary
+
+
+def filter_contours_and_holes(
+    binary_mask: np.ndarray,
+    min_effective_area: float,
+    min_hole_area: float,
+    max_n_holes: int,
+):
+    """
+    Find contours from the binary mask and filter them based on effective area.
+
+    For each external contour (one with no parent), identify child contours (holes),
+    sort them by area (largest first), and keep up to max_n_holes that exceed min_hole_area.
+    The effective area is computed as the area of the external contour minus the sum of areas
+    of the selected holes. Only contours with effective area above min_effective_area are retained.
+
+    Returns:
+        filtered_contours (list): List of external contours (numpy arrays).
+        holes_list (list): Corresponding list of lists of hole contours.
+    """
+    contours, hierarchy = cv2.findContours(
+        binary_mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE
+    )
+    if hierarchy is None:
+        return [], []
+    hierarchy = hierarchy[0]  # shape: (N, 4)
+    filtered_contours = []
+    holes_list = []
+    for idx, h in enumerate(hierarchy):
+        if h[3] != -1:
+            continue  # Only external contours
+        ext_cont = contours[idx]
+        ext_area = cv2.contourArea(ext_cont)
+        # Find child contours (holes)
+        hole_idxs = [i for i, hr in enumerate(hierarchy) if hr[3] == idx]
+        # Sort holes by area descending and keep up to max_n_holes
+        sorted_holes = sorted(
+            [contours[i] for i in hole_idxs], key=cv2.contourArea, reverse=True
+        )
+        selected_holes = [
+            hole
+            for hole in sorted_holes[:max_n_holes]
+            if cv2.contourArea(hole) > min_hole_area
+        ]
+        total_hole_area = sum(cv2.contourArea(hole) for hole in selected_holes)
+        effective_area = ext_area - total_hole_area
+        if effective_area > min_effective_area:
+            filtered_contours.append(ext_cont)
+            holes_list.append(selected_holes)
+    return filtered_contours, holes_list
+
+
+def scale_contours(contours: list, scale: tuple) -> list:
+    """
+    Scale contour coordinates by the provided scale factors.
+
+    Args:
+        contours: List of contours (each a numpy array of points).
+        scale: Tuple (sx, sy) for scaling.
+
+    Returns:
+        List of scaled contours.
+    """
+    scaled = []
+    for cont in contours:
+        scaled.append((cont * np.array(scale, dtype=np.float32)).astype(np.int32))
+    return scaled
+
+
+def point_in_tissue(x: int, y: int, ext_contours: list, holes_list: list) -> bool:
+    """
+    Check if point (x, y) lies within any external contour and not inside its corresponding holes.
+
+    For each external contour in ext_contours (paired with holes_list),
+    if the point is inside the contour and not inside any of its holes, return True.
+    """
+    for cont, holes in zip(ext_contours, holes_list):
+        if cv2.pointPolygonTest(cont, (x, y), False) >= 0:
+            inside_hole = False
+            for hole in holes:
+                if cv2.pointPolygonTest(hole, (x, y), False) >= 0:
+                    inside_hole = True
+                    break
+            if not inside_hole:
+                return True
+    return False
+
+
+def tile(x: torch.Tensor, size: int):
+    C, H, W = x.shape[-3:]
+
+    pad_h = (size - H % size) % size
+    pad_w = (size - W % size) % size
+    if pad_h > 0 or pad_w > 0:
+        x = torch.nn.functional.pad(x, (0, pad_w, 0, pad_h))
+
+    nh, nw = x.size(2) // size, x.size(3) // size
+    return (
+        x.view(-1, C, nh, size, nw, size)
+        .permute(0, 2, 4, 1, 3, 5)
+        .reshape(-1, C, size, size)
+    )