Upload epoch=17-step=1836.ckpt

README.md

@@ -53,6 +53,36 @@ extra_gated_prompt: >-
 # ODELIA Classification Baseline Model
 For a comprehensive description of the model and its intended use, please refer to our paper: [Read the paper](https://arxiv.org/abs/2506.00474)
 
+## Setup
+
+To run the code, we recommend creating a Python virtual environment.
+
+### Using venv
+
+```bash
+# Create a virtual environment
+python -m venv venv
+
+# Activate the environment
+# On Linux/Mac:
+source venv/bin/activate
+# On Windows:
+# venv\Scripts\activate
+
+# Install dependencies
+pip install torch torchvision numpy huggingface_hub torchio matplotlib transformers einops x_transformers
+```
+
+### Using Conda
+
+```bash
+# Create a conda environment
+conda create -n odelia_hf python=3.10
+conda activate odelia_hf
+
+# Install dependencies
+pip install torch torchvision numpy huggingface_hub torchio matplotlib transformers einops x_transformers
+```
 
 ## Get Probabilities and Attention
 

model_config.json

@@ -1,6 +1,6 @@
 {
   "checkpoint_source": "epoch=17-step=1836.ckpt",
-  "created_at": "2025-10-26T16:51:08.236236Z",
+  "created_at": "2025-11-23T17:24:34.909430Z",
   "hparams": {
     "backbone_type": "dinov3",
     "in_ch": 1,

models.py

@@ -2,7 +2,7 @@ from einops import rearrange
 import torch.nn as nn
 import torch
 import math
-from transformers import AutoModel
+from transformers import AutoModel, Dinov2WithRegistersModel, Dinov2WithRegistersConfig, DINOv3ViTConfig, DINOv3ViTModel
 from x_transformers import Encoder
 
 
@@ -13,22 +13,44 @@ class _MST(nn.Module):
         backbone_type="dinov3",
         model_size = "s", # 34, 50, ... or 's', 'b', 'l'
         slice_fusion_type = "transformer", # transformer, linear, average, none 
+        weights=True,
     ):
         super().__init__()
         self.backbone_type = backbone_type
         self.slice_fusion_type = slice_fusion_type
 
         if backbone_type == "dinov2":
-            model_size = {'s':'small', 'b':'base', 'l':'large'}.get(model_size)
-            self.backbone = AutoModel.from_pretrained(f"facebook/dinov2-with-registers-{model_size}")
+            model_size_key = {'s':'small', 'b':'base', 'l':'large'}.get(model_size)
+            model_name = f"facebook/dinov2-with-registers-{model_size_key}"
+            if weights:
+                self.backbone = AutoModel.from_pretrained(model_name)
+            else:
+                configs = {
+                    'small': Dinov2WithRegistersConfig(hidden_size=384, num_hidden_layers=12, num_attention_heads=6),
+                    'base': Dinov2WithRegistersConfig(hidden_size=768, num_hidden_layers=12, num_attention_heads=12),
+                    'large': Dinov2WithRegistersConfig(hidden_size=1024, num_hidden_layers=24, num_attention_heads=16),
+                }
+                config = configs.get(model_size_key)
+                config.image_size = 518
+                config.patch_size = 14
+                self.backbone = Dinov2WithRegistersModel(config)
             emb_ch = self.backbone.config.hidden_size
         elif backbone_type == "dinov3":
-            self.backbone = AutoModel.from_pretrained(f"facebook/dinov3-vit{model_size}16-pretrain-lvd1689m")
-            emb_ch = self.backbone.config.hidden_size
+            model_name = f"facebook/dinov3-vit{model_size}16-pretrain-lvd1689m"
+            if weights:
+                self.backbone = AutoModel.from_pretrained(model_name)
+            else:
+                configs = {
+                    's': DINOv3ViTConfig(hidden_size=384, num_hidden_layers=12, num_attention_heads=6, intermediate_size=1536, patch_size=16, num_register_tokens=4),
+                    'b': DINOv3ViTConfig(hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, patch_size=16, num_register_tokens=4),
+                    'l': DINOv3ViTConfig(hidden_size=1024, num_hidden_layers=24, num_attention_heads=16, intermediate_size=4096, patch_size=16, num_register_tokens=4),
+                }
+                config = configs.get(model_size)
+                self.backbone = DINOv3ViTModel(config)
         else:
             raise ValueError("Unknown backbone_type")
 
-
+        emb_ch = self.backbone.config.hidden_size
         self.emb_ch = emb_ch 
         if slice_fusion_type == "transformer":
             self.slice_fusion = Encoder(
@@ -144,9 +166,9 @@ class _MST(nn.Module):
 
 
 class MSTRegression(nn.Module):
-    def __init__(self, in_ch=1, out_ch=1, spatial_dims=3, backbone_type="dinov3", model_size="s", slice_fusion_type="transformer", optimizer_kwargs={'lr':1e-5}, **kwargs):
+    def __init__(self, in_ch=1, out_ch=1, spatial_dims=3, backbone_type="dinov3", model_size="s", slice_fusion_type="transformer", weights=True, **kwargs):
         super().__init__()
-        self.mst = _MST(out_ch=out_ch, backbone_type=backbone_type, model_size=model_size, slice_fusion_type=slice_fusion_type)
+        self.mst = _MST(out_ch=out_ch, backbone_type=backbone_type, model_size=model_size, slice_fusion_type=slice_fusion_type, weights=weights)
     
     def forward(self, x):
         return self.mst(x)

predict_attention.py

@@ -124,7 +124,7 @@ def load_model(repo_id= "ODELIA-AI/MST") -> MSTRegression:
         config = json.load(fp)
 
     hparams = config.get("hparams", {})
-    model = MSTRegression(**hparams)
+    model = MSTRegression(weights=False, **hparams)
 
     state_dict_path = hf_hub_download(repo_id=repo_id, repo_type="model", filename="state_dict.pt")
     state_dict = torch.load(state_dict_path, map_location="cpu")