Duplicate from Astaxanthin/KEEP

Co-authored-by: Xiao Zhou <Astaxanthin@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,130 @@
+---
+license: mit
+---
+
+# Model Card for KEEP
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+[Paper](https://www.cell.com/cancer-cell/fulltext/S1535-6108(26)00058-9) | [Github](https://github.com/MAGIC-AI4Med/KEEP) | [Webpage](https://loiesun.github.io/keep/) | [Cite](#citation)
+
+**KEEP** (**K**nowledg**E**-**E**nhanced **P**athology) is a foundation model designed for cancer diagnosis that integrates disease knowledge into vision-language pre-training. It utilizes a comprehensive disease knowledge graph (KG) containing 11,454 human diseases and 139,143 disease attributes, such as synonyms, definitions, and hierarchical relationships. KEEP reorganizes millions of publicly available noisy pathology image-text pairs into 143K well-structured semantic groups based on the hierarchical relations of the disease KG. By incorporating disease knowledge into the alignment process, KEEP achieves more nuanced image and text representations. The model is validated on 18 diverse benchmarks with over 14,000 whole-slide images (WSIs), demonstrating state-of-the-art performance in zero-shot cancer diagnosis, including an average sensitivity of 89.8% for cancer detection across 7 cancer types. KEEP also excels in subtyping rare cancers, achieving strong generalizability in diagnosing rare tumor subtypes.
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+- **Developed by:** MAGIC-AI4Med team from Shanghai Jiao Tong University and Shanghai AI Lab.
+- **Model type:** Vision-language models (vision encoder: ViT-L/16; text encoder: Bert)
+- **Pretrain datasets:** 143K pathology semantic groups, each with a single caption and multiple images.
+- **License:** MIT
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** https://github.com/MAGIC-AI4Med/KEEP
+- **Paper [optional]:** https://www.cell.com/cancer-cell/fulltext/S1535-6108(26)00058-9
+- **Demo [optional]:** [More Information Needed]
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+## Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+```python
+from transformers import AutoModel, AutoTokenizer
+from torchvision import transforms
+from PIL import Image
+
+model = AutoModel.from_pretrained("Astaxanthin/KEEP", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("Astaxanthin/KEEP", trust_remote_code=True)
+model.eval()
+transform = transforms.Compose([
+    transforms.Resize(size=224, interpolation=transforms.InterpolationMode.BICUBIC),
+    transforms.CenterCrop(size=(224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
+])
+
+example_image_path = './example.tif'
+example_text = ['an H&E image of breast invasive carcinoma.', 'an H&E image of normal tissue.', 'an H&E image of lung adenocarcinoma.']
+
+img_input =  transform(Image.open(example_image_path).convert('RGB')).unsqueeze(0)
+token_input = tokenizer(example_text,max_length=256,padding='max_length',truncation=True, return_tensors='pt')
+
+img_feature = model.encode_image(img_input)
+text_feature = model.encode_text(token_input)
+
+```
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+We present benchmark results for a range of representative tasks. A complete set of benchmarks can be found in the [paper](https://arxiv.org/abs/2412.18***). These results will be updated with each new iteration of KEEP. 
+
+
+### Results
+
+
+#### Zero-shot Cancer Region Segmentation (DICE) 
+| Models | PLIP[[1]](https://www.nature.com/articles/s41591-023-02504-3) | QuiltNet [[2]](https://proceedings.neurips.cc/paper_files/paper/2023/hash/775ec578876fa6812c062644964b9870-Abstract-Datasets_and_Benchmarks.html) |   MI-Zero (Pub) [[3]](https://openaccess.thecvf.com/content/CVPR2023/html/Lu_Visual_Language_Pretrained_Multiple_Instance_Zero-Shot_Transfer_for_Histopathology_Images_CVPR_2023_paper.html) | CONCH [[4]](https://www.nature.com/articles/s41591-024-02856-4) | **KEEP(Ours)**  |
+|:---------------|--------------:|---------------------:|-------------------------:|-----------------:|------------------:|
+| CAMELYON16 | 0.253 | 0.157 | 0.186 | 0.292 | **0.361** |
+| PANDA | 0.295 | 0.309 | 0.276 | 0.315 | **0.334** |
+| AGGC22 | 0.284 | 0.282 | 0.324 | 0.449 | **0.530** |
+
+#### Zero-shot Cancer Detection (AUROC)
+| Models | CHIEF[[1]](https://www.nature.com/articles/s41586-024-07894-z) |   PLIP [[2]](https://www.nature.com/articles/s41591-023-02504-3)   |   QuiltNet [[3]](https://proceedings.neurips.cc/paper_files/paper/2023/hash/775ec578876fa6812c062644964b9870-Abstract-Datasets_and_Benchmarks.html)     |   MI-Zero (Pub) [[4]](https://openaccess.thecvf.com/content/CVPR2023/html/Lu_Visual_Language_Pretrained_Multiple_Instance_Zero-Shot_Transfer_for_Histopathology_Images_CVPR_2023_paper.html) |   CONCH [[5]](https://www.nature.com/articles/s41591-024-02856-4) | KEEP(Ours) |
+|:---------------|--------------:|--------------------:|-----------------:|-----------------:|------------------:| -----------------:|
+| CPTAC-CM | 0.915 | 0.970 | 0.972 | 0.985 | **0.994** | **0.994** | 
+| CPTAC-CCRCC | 0.723 | 0.330 | 0.755 | 0.886 | 0.871 | **0.999** |
+| CPTAC-PDA | 0.825 | 0.391 | 0.464 | 0.796 | 0.920 | **0.929** |
+| CPTAC-UCEC | 0.955 | 0.945 | 0.973 | 0.979 | 0.996 | **0.998** | 
+| CPTAC-LSCC | 0.901 | 0.965 | 0.966 | 0.910 | **0.987** | 0.983 | 
+| CPTAC-HNSCC | 0.946 | 0.898 | 0.874 | 0.918 | **0.982** | 0.976 | 
+| CPTAC-LUAD | 0.891 | 0.988 | 0.991 | 0.981 | 0.999 | **1.000** |  
+
+#### Zero-shot Cancer Subtyping (BACC) 
+| Models     | PLIP [[1]](https://www.nature.com/articles/s41591-023-02504-3) |   QuiltNet [[2]](https://proceedings.neurips.cc/paper_files/paper/2023/hash/775ec578876fa6812c062644964b9870-Abstract-Datasets_and_Benchmarks.html)      |   MI-Zero (Pub) [[3]](https://openaccess.thecvf.com/content/CVPR2023/html/Lu_Visual_Language_Pretrained_Multiple_Instance_Zero-Shot_Transfer_for_Histopathology_Images_CVPR_2023_paper.html)     |   CONCH [[4]](https://www.nature.com/articles/s41591-024-02856-4) |   **KEEP(Ours)**  |
+|:---------------|--------------:|---------------------------:|-------------------------:|-----------------:|------------------:|
+| TCGA-BRCA | 0.519 | 0.500 | 0.633 | 0.727 | **0.774** |
+| TCGA-NSCLC | 0.699 | 0.667 | 0.753 | 0.901 | **0.902** |
+| TCGA-RCC | 0.735 | 0.755 | 0.908 | 0.921 | **0.926** |
+| TCGA-ESCA | 0.614 | 0.746 | 0.954 | 0.923 | **0.977** |
+| TCGA-BRAIN | 0.361 | 0.346 | 0.361 | 0.453 | **0.604** |
+| UBC-OCEAN | 0.343 | 0.469 | 0.652 | **0.674** | 0.661 |
+| CPTAC-NSCLC | 0.647 | 0.607 | 0.643 | 0.836 | **0.863** |
+| EBRAINS | 0.096 | 0.093 | 0.325 | 0.371 | **0.456** |
+
+### Summary
+
+Validated on 18 diverse benchmarks with more than 14,000 whole slide images (WSIs), KEEP achieves state-of-the-art performance in zero-shot cancer diagnostic tasks. Notably, for cancer detection, KEEP demonstrates an average sensitivity of 89.8% at a specificity of 95.0% across 7 cancer types, significantly outperforming vision-only foundation models and highlighting its promising potential for clinical application. For cancer subtyping, KEEP achieves a median balanced accuracy of 0.456 in subtyping 30 rare brain cancers, indicating strong generalizability for diagnosing rare tumors.
+
+
+## Citation
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+```
+@article{ZHOU2026,
+title = {Knowledge-enhanced pretraining for vision-language pathology foundation model on cancer diagnosis},
+journal = {Cancer Cell},
+year = {2026},
+issn = {1535-6108},
+doi = {https://doi.org/10.1016/j.ccell.2026.01.019},
+url = {https://www.sciencedirect.com/science/article/pii/S1535610826000589},
+author = {Xiao Zhou and Luoyi Sun and Dexuan He and Wenbin Guan and Ge Wang and Ruifen Wang and Lifeng Wang and Xiaojun Yuan and Xin Sun and Ya Zhang and Kun Sun and Yanfeng Wang and Weidi Xie},
+keywords = {pathology foundation model, cancer diagnosis, vision-language, knowledge enhancement},
+}
+```

config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "KEEPModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_keep.KEEPConfig",
+    "AutoModel": "modeling_keep.KEEPModel"
+  },
+  "model_type": "keep",
+  "projection_dim": 768,
+  "text_config": {
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 512,
+    "model_type": "bert",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "type_vocab_size": 2,
+    "vocab_size": 30522
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "vision_config": {
+    "attn_drop_rate": 0.0,
+    "drop_rate": 0.0,
+    "dynamic_img_size": true,
+    "hidden_size": 1024,
+    "img_size": 224,
+    "init_values": 1e-05,
+    "mlp_ratio": 4.0,
+    "model_type": "vit",
+    "num_classes": 0,
+    "num_heads": 16,
+    "num_layers": 24,
+    "patch_size": 16,
+    "qkv_bias": true
+  }
+}

model.safetensors

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

modeling_keep.py

@@ -0,0 +1,88 @@
+from transformers import PretrainedConfig, PreTrainedModel, BertModel, BertConfig
+import timm
+import torch.nn as nn
+import torch
+import numpy
+from torchvision import transforms
+from PIL import Image
+
+class RenameLayerScale(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.weight = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.weight) if self.inplace else x * self.weight
+
+timm.models.vision_transformer.LayerScale = RenameLayerScale
+
+class KEEPConfig(PretrainedConfig):
+    model_type = "keep"  # 
+
+    def __init__(
+        self,
+        vision_config=None,  # Vision Encoder
+        text_config=None,    # Text Encoder
+        projection_dim=768,  
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vision_config = vision_config
+        self.text_config = text_config
+        self.projection_dim = projection_dim
+
+class KEEPModel(PreTrainedModel):
+    config_class = KEEPConfig  # 
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        # Vision Encoder 
+        vision_config = config.vision_config
+        self.visual = timm.create_model(
+            "vit_large_patch16_224",
+            pretrained=False,
+            img_size=vision_config["img_size"],
+            patch_size=vision_config["patch_size"],
+            init_values=vision_config["init_values"],
+            num_classes=vision_config["num_classes"],
+        )
+
+        self.visual_head = nn.Sequential(
+                    nn.Linear(self.visual.num_features, config.projection_dim),
+                    nn.GELU(),
+                    nn.Linear(config.projection_dim, config.projection_dim)
+                )
+
+        # Text Encoder
+        text_config =  BertConfig(**config.text_config)
+        self.text = BertModel(text_config)
+        
+        self.logit_scale = nn.Parameter(torch.ones([]) * numpy.log(1 / 0.04))
+
+    def encode_image(self, image_inputs):
+        vision_features = self.visual(image_inputs)  # [batch_size, vision_dim]
+        vision_features =  torch.nn.functional.normalize(self.visual_head(vision_features), dim=-1)  # [batch_size, projection_dim]
+        
+        return vision_features
+    
+    def encode_text(self, text_inputs):
+        text_features = torch.nn.functional.normalize(self.text(**text_inputs).pooler_output, dim=-1)  # [batch_size, text_dim]
+        return text_features
+    
+    
+    def forward(self, image_inputs, text_inputs):
+        vision_features = self.encode_image(image_inputs)
+        
+        text_features = self.encode_text(text_inputs)
+
+        return {
+            "vision_features": vision_features,  
+            "text_features": text_features       
+        }

pytorch_model.bin

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

tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "do_lower_case": true,
+  "tokenizer_class": "BertTokenizer" 
+}