initial commit

.gitignore

@@ -0,0 +1,172 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc
+

LICENSE

@@ -0,0 +1,57 @@
+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

@@ -1,5 +1,35 @@
----
-license: other
-license_name: license
-license_link: LICENSE
----
+# Description
+MSI classification of CRC tumors using EXAONEPath 1.0.0 - a patch-level foundation model for pathology
+
+# Model Overview
+This model serves as a reference for predicting MSI status using CRC (colorectal cancer) tumor images as input. When the model receives an H&E-stained whole slide image as input, it removes artifacts observed in the image and extracts only tissue-related objects. These objects are then reconstructed into a set of tiles with a size of 256 by 256 pixels at an mpp (micron per pixel) of 0.5.
+
+The tiles pass through the EXAONEPath v1.0 patch-level foundation model (https://huggingface.co/LGAI-EXAONE/EXAONEPath), which converts them into a set of features. These features are then integrated into a slide-level feature representation through an aggregator. Finally, a linear classifier predicts the MSI status (MSS or MSI-H/L).
+
+The model achieves an average performance of AUROC 0.93 on TCGA-COAD + TCGA-READ data and 0.84 on in-house data.
+
+This open-source release aims to demonstrate that the combination of EXAONEPath and the aggregator can effectively perform pathological tasks. It is hoped that this source code will serve as an important reference not only for CRC MSI prediction but also as an image-based solution for various disease-related problems, including molecular subtyping, tumor subtyping, and mutation prediction.
+
+This open-source code is designed to be compatible with the MONAI framework (https://monai.io/), and users must check the accompanying license before use.
+
+![EXAONEPath](assets/exaonepath_v1.png)
+
+# Input and Output Formats
+The input is a `.svs` file containing whole slide image.
+The output is an array with probabilities for each of the two classes (MSS / MSI(H/L)).
+
+# Execute Inference
+The inference can be executed as follows
+1. Copy your `.svs` files into `samples` directory
+2. Run inference
+    ```
+    python -m monai.bundle run evaluating --meta_file configs/metadata.json --config_file configs/inference.json configs/logging.conf
+    ```
+
+## Contact
+LG AI Research Technical Support: <a href="mailto:[email protected]">[email protected]</a>
+
+# License
+Copyright (c) LG AI Research
+
+The model is licensed under [EXAONEPath AI Model License Agreement 1.0 - NC](LICENSE).

configs/inference.yaml

@@ -0,0 +1,26 @@
+imports:
+- $import os
+- $import glob
+- $import torch
+- $import scripts
+- $import scripts.inference
+
+device: '$torch.device("cuda" if torch.cuda.is_available() else "cpu")'
+
+model_config:
+  _target_: scripts.exaonepath.EXAONEPathV1Downstream
+  step_size: 256
+  patch_size: 256
+  macenko: true
+  device: '$@device'
+
+model: '$@model_config'
+
+input_dir: 'samples'
+input_files: '$sorted(glob.glob(@input_dir+''/*.svs''))'
+
+root_dir: '$os.path.dirname(os.path.dirname(scripts.__file__))'
+
+inference:
+- [email protected]_state_dict(torch.load(os.path.join(@root_dir, "models/exaonepath_v1.0.0_msi.pt")))
+- $scripts.inference.infer(@model, @input_files)

configs/metadata.json

@@ -0,0 +1,50 @@
+{
+    "version": "0.1.0",
+    "changelog": {
+        "0.1.0": "Initial release"
+    },
+    "monai_version": "1.5",
+    "pytorch_version": "2.6.0",
+    "numpy_version": "1.26.4",
+    "required_packages_version": {
+        "torchvision": "0.21.0",
+        "openslide-bin": "4.0.0.6",
+        "openslide-python": "1.4.1",
+        "pandas": "2.2.3",
+        "torchstain": "1.4.1",
+        "opencv-python-headless": "4.11.0.86",
+        "einops": "0.8.1"
+    },
+    "name": "MSI predictor from WSI using EXAONEPath v1.0",
+    "task": "MSI Classification of Colorectal Cancer (CRC) Tumors",
+    "description": "MSI classification of CRC tumors using EXAONEPath 1.0.0 - a patch-level foundation model for pathology",
+    "authors": "LG AI Research",
+    "copyright": "Copyright (c) LG AI Research",
+    "data_source": "Whole Slide Images",
+    "data_type": "float32",
+    "intended_use": "Research only",
+    "network_data_format": {
+        "inputs": {
+            "image": {
+                "type": "image",
+                "format": "svs",
+                "modality": "pathology"
+            }
+        },
+        "outputs": {
+            "pred": {
+                "type": "probabilities",
+                "format": "classes",
+                "num_channels": 2,
+                "spatial_shape": [2],
+                "dtype": "float32",
+                "value_range": [0, 1],
+                "is_patch_data": false,
+                "channel_def": {
+                    "0": "MSS",
+                    "1": "MSI"
+                }
+            }
+        }
+    }
+}

models/exaonepath_v1.0.0_msi.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c7c3199335922de3c156ede12e2d655b23ec988529be18611981832d6ab8095
+size 625193090

models/macenko_param.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74a6e178f1d3ea897152e544857eb880045a8ffb2be1e05b1639386083e68435
+size 999

requirements.txt

@@ -0,0 +1,13 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+
+torch==2.6.0+cu124
+torchvision==0.21.0+cu124
+numpy==1.26.4
+opencv-python-headless==4.11.0.86
+openslide-bin==4.0.0.6
+openslide-python==1.4.1
+monai-weekly[ignite,pyyaml]==1.5.dev2506
+pandas==2.2.3
+torchstain==1.4.1
+fire==0.7.0
+einops==0.8.1

scripts/aggregator.py

@@ -0,0 +1,334 @@
+from functools import wraps
+from math import log, pi
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange, Reduce
+from torch import einsum, nn
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def cache_fn(f):
+    cache = dict()
+
+    @wraps(f)
+    def cached_fn(*args, _cache=True, key=None, **kwargs):
+        if not _cache:
+            return f(*args, **kwargs)
+        nonlocal cache
+        if key in cache:
+            return cache[key]
+        result = f(*args, **kwargs)
+        cache[key] = result
+        return result
+
+    return cached_fn
+
+
+def fourier_encode(x, max_freq, num_bands=4):
+    x = x.unsqueeze(-1)
+    device, dtype, orig_x = x.device, x.dtype, x
+
+    scales = torch.linspace(1.0, max_freq / 2, num_bands, device=device, dtype=dtype)
+    scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]
+
+    x = x * scales * pi
+    x = torch.cat([x.sin(), x.cos()], dim=-1)
+    x = torch.cat((x, orig_x), dim=-1)
+    return x
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn, context_dim=None):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+        self.norm_context = nn.LayerNorm(context_dim) if exists(context_dim) else None
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+
+        if exists(self.norm_context):
+            context = kwargs["context"]
+            normed_context = self.norm_context(context)
+            kwargs.update(context=normed_context)
+
+        return self.fn(x, **kwargs)
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gates = x.chunk(2, dim=-1)
+        return x * F.gelu(gates)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4, dropout=0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult * 2),
+            GEGLU(),
+            nn.Linear(dim * mult, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.scale = dim_head**-0.5
+        self.heads = heads
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+        self.to_out = nn.Linear(inner_dim, query_dim)
+
+    def forward(self, x, context=None, mask=None):
+        h = self.heads
+
+        q = self.to_q(x)
+        context = default(context, x)
+        k, v = self.to_kv(context).chunk(2, dim=-1)
+
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
+
+        sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+
+        if exists(mask):
+            mask = rearrange(mask, "b ... -> b (...)")
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, "b j -> (b h) () j", h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        attn = sim.softmax(dim=-1)
+        attn = self.dropout(attn)
+
+        out = einsum("b i j, b j d -> b i d", attn, v)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
+        return self.to_out(out)
+
+
+class Perceiver(nn.Module):
+    def __init__(
+        self,
+        *,
+        num_freq_bands,
+        depth,
+        max_freq,
+        input_channels=3,
+        input_axis=2,
+        num_latents=512,
+        latent_dim=512,
+        cross_heads=1,
+        latent_heads=8,
+        cross_dim_head=64,
+        latent_dim_head=64,
+        num_classes=1000,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        weight_tie_layers=False,
+        fourier_encode_data=True,
+        self_per_cross_attn=1,
+        final_classifier_head=True,
+        pool="mean",
+        latent_init=None,
+    ):
+        """The shape of the final attention mechanism will be:
+        depth * (cross attention -> self_per_cross_attn * self attention)
+
+        Args:
+          num_freq_bands: Number of freq bands, with original value (2 * K + 1)
+          depth: Depth of net.
+          max_freq: Maximum frequency, hyperparameter depending on how
+              fine the data is.
+          freq_base: Base for the frequency
+          input_channels: Number of channels for each token of the input.
+          input_axis: Number of axes for input data (2 for images, 3 for video)
+          num_latents: Number of latents, or induced set points, or centroids.
+              Different papers giving it different names.
+          latent_dim: Latent dimension.
+          cross_heads: Number of heads for cross attention. Paper said 1.
+          latent_heads: Number of heads for latent self attention, 8.
+          cross_dim_head: Number of dimensions per cross attention head.
+          latent_dim_head: Number of dimensions per latent self attention head.
+          num_classes: Output number of classes.
+          attn_dropout: Attention dropout
+          ff_dropout: Feedforward dropout
+          weight_tie_layers: Whether to weight tie layers (optional).
+          fourier_encode_data: Whether to auto-fourier encode the data, using
+              the input_axis given. defaults to True, but can be turned off
+              if you are fourier encoding the data yourself.
+          self_per_cross_attn: Number of self attention blocks per cross attn.
+          final_classifier_head: mean pool and project embeddings to number of classes (num_classes) at the end
+        """
+        super().__init__()
+        self.input_axis = input_axis
+        self.max_freq = max_freq
+        self.num_freq_bands = num_freq_bands
+        self.self_per_cross_attn = self_per_cross_attn
+
+        self.fourier_encode_data = fourier_encode_data
+        fourier_channels = (
+            (input_axis * ((num_freq_bands * 2) + 1)) * 2 if fourier_encode_data else 0
+        )
+        input_dim = fourier_channels + input_channels
+
+        self.latents = nn.Parameter(torch.randn(num_latents, latent_dim))
+        if latent_init is not None:
+            latent_init_feat = torch.load(latent_init)
+            if type(latent_init_feat) != torch.Tensor:
+                latent_init_feat = torch.Tensor(latent_init_feat)
+            if len(latent_init_feat.shape) == 3:
+                latent_init_feat = latent_init_feat[0]
+                with torch.no_grad():
+                    self.latents.copy_(latent_init_feat)
+                print(f"load latent feature: , {latent_init}")
+
+        get_cross_attn = lambda: PreNorm(
+            latent_dim,
+            Attention(
+                latent_dim,
+                input_dim,
+                heads=cross_heads,
+                dim_head=cross_dim_head,
+                dropout=attn_dropout,
+            ),
+            context_dim=input_dim,
+        )
+        get_cross_ff = lambda: PreNorm(
+            latent_dim, FeedForward(latent_dim, dropout=ff_dropout)
+        )
+        get_latent_attn = lambda: PreNorm(
+            latent_dim,
+            Attention(
+                latent_dim,
+                heads=latent_heads,
+                dim_head=latent_dim_head,
+                dropout=attn_dropout,
+            ),
+        )
+        get_latent_ff = lambda: PreNorm(
+            latent_dim, FeedForward(latent_dim, dropout=ff_dropout)
+        )
+
+        get_cross_attn, get_cross_ff, get_latent_attn, get_latent_ff = map(
+            cache_fn, (get_cross_attn, get_cross_ff, get_latent_attn, get_latent_ff)
+        )
+
+        self.layers = nn.ModuleList([])
+        for i in range(depth):
+            should_cache = i > 0 and weight_tie_layers
+            cache_args = {"_cache": should_cache}
+
+            self_attns = nn.ModuleList([])
+
+            for block_ind in range(self_per_cross_attn):
+                self_attns.append(
+                    nn.ModuleList(
+                        [
+                            get_latent_attn(**cache_args, key=block_ind),
+                            get_latent_ff(**cache_args, key=block_ind),
+                        ]
+                    )
+                )
+            if self_per_cross_attn == 0:
+                self_attns.append(get_latent_ff(**cache_args, key=block_ind))
+
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        get_cross_attn(**cache_args),
+                        get_cross_ff(**cache_args),
+                        self_attns,
+                    ]
+                )
+            )
+
+        if final_classifier_head:
+            if pool == "cat":
+                self.to_logits = nn.Sequential(
+                    Rearrange("b n d -> b (n d)"),
+                    nn.LayerNorm(num_latents * latent_dim),
+                    nn.Linear(num_latents * latent_dim, num_classes),
+                )
+            elif pool == "mlp":
+                self.to_logits = nn.Sequential(
+                    Reduce("b n d -> b d", "mean"),
+                    nn.LayerNorm(latent_dim),
+                    nn.Linear(latent_dim, latent_dim),
+                    nn.ReLU(),
+                    nn.LayerNorm(latent_dim),
+                    nn.Linear(latent_dim, num_classes),
+                )
+            else:
+                self.to_logits = nn.Sequential(
+                    Reduce("b n d -> b d", pool),
+                    nn.LayerNorm(latent_dim),
+                    nn.Linear(latent_dim, num_classes),
+                )
+
+    def forward(self, h, label=None, mask=None, pretrain=False, coords=None):
+        b, *axis, _, device, dtype = *h.shape, h.device, h.dtype
+        assert (
+            len(axis) == self.input_axis
+        ), "input data must have the right number of axis"
+
+        if self.fourier_encode_data:
+            # calculate fourier encoded positions in the range of [-1, 1], for all axis
+
+            # axis_pos = list(map(lambda size: torch.linspace(-1., 1., steps=size, device=device, dtype=dtype), axis))
+            # pos = torch.stack(torch.meshgrid(*axis_pos, indexing = 'ij'), dim = -1)
+            # enc_pos = fourier_encode(pos, self.max_freq, self.num_freq_bands)
+            # enc_pos = rearrange(enc_pos, '... n d -> ... (n d)')
+            # enc_pos = repeat(enc_pos, '... -> b ...', b = b)
+
+            enc_pos = fourier_encode(coords, self.max_freq, self.num_freq_bands)
+            enc_pos = rearrange(enc_pos, "... n d -> ... (n d)")
+
+            h = torch.cat((h, enc_pos), dim=-1)
+
+        # concat to channels of data and flatten axis
+
+        h = rearrange(h, "b ... d -> b (...) d")
+
+        x = repeat(self.latents, "n d -> b n d", b=b)
+
+        # layers
+
+        for cross_attn, cross_ff, self_attns in self.layers:
+            x = cross_attn(x, context=h, mask=mask) + x
+            x = cross_ff(x) + x
+
+            if self.self_per_cross_attn > 0:
+                for self_attn, self_ff in self_attns:
+                    x = self_attn(x) + x
+                    x = self_ff(x) + x
+            else:
+                x = self_attns[0](x) + x
+        # allow for fetching embeddings
+
+        if pretrain:
+            return x.mean(dim=1)
+
+        # to logits
+        logits = self.to_logits(x)
+        Y_hat = torch.topk(logits, 1, dim=1)[1]
+        Y_prob = F.softmax(logits, dim=1)
+        return logits, Y_prob, Y_hat

scripts/constants.py

@@ -0,0 +1 @@
+CLASS_NAMES = ["MSS", "MSI"]

scripts/dataset.py

@@ -0,0 +1,47 @@
+import os
+
+from openslide import OpenSlide
+from scripts.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,
+    ):
+        self.pretrained = pretrained
+        self.wsi = OpenSlide(wsi_path)
+        self.patch_size = patch_size
+        self.patch_level = patch_level
+
+        if macenko:
+            normalizer = MacenkoNormalizer(
+                target_path=os.path.join(
+                    os.path.dirname(os.path.dirname(os.path.join(__file__))),
+                    "models",
+                    "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)
+        return img

scripts/exaonepath.py

@@ -0,0 +1,92 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from scripts.aggregator import Perceiver
+from scripts.dataset import WSIPatchDataset
+from scripts.feature_extractor import vit_base
+from scripts.wsi_utils import extract_tissue_patch_coords
+from torch.utils.data import DataLoader
+
+
+class EXAONEPathV1Downstream(nn.Module):
+    def __init__(
+        self, device: torch.device, step_size=256, patch_size=256, macenko=True
+    ):
+        super(EXAONEPathV1Downstream, self).__init__()
+        self.step_size = step_size
+        self.patch_size = patch_size
+        self.macenko = macenko
+        self.device = device
+
+        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 = Perceiver(
+            input_channels=768,
+            input_axis=1,
+            num_freq_bands=6,
+            max_freq=10.0,
+            depth=6,
+            num_latents=256,
+            latent_dim=512,
+            cross_heads=1,
+            latent_heads=8,
+            cross_dim_head=64,
+            latent_dim_head=64,
+            num_classes=2,
+            fourier_encode_data=False,
+            self_per_cross_attn=2,
+            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,
+        )
+        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 feature
+            feature = feature.to("cpu", non_blocking=True)
+            features_list.append(feature)
+        print("")
+        print("Feature extraction finished")
+
+        features = torch.cat(features_list)
+
+        # Aggregate features
+        self.agg_model.eval()
+        logits, Y_prob, Y_hat = self.agg_model(features[None].to(self.device))
+        probs = Y_prob[0].cpu()
+
+        return probs

scripts/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

scripts/inference.py

@@ -0,0 +1,11 @@
+from scripts.constants import CLASS_NAMES
+
+
+def infer(model, input_files):
+    for p in input_files:
+        print("Processing", p, "...")
+        probs = model(p)
+        result_str = "Result -- " + " / ".join(
+            [f"{name}: {probs[i].item():.4f}" for i, name in enumerate(CLASS_NAMES)]
+        )
+        print(result_str + "\n")

scripts/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

scripts/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)
+    )