Update README.md

README.md

@@ -70,43 +70,69 @@ A Vision Foundation Model for Affective Computing
 
 ---
 
-## Pre-trained Weights
+## Pre-trained Checkpoint
 
-Get the weights from the **[GitHub Releases](https://github.com/GkikasStefanos/PainFormer/releases)**.
+The checkpoint is stored under `checkpoint/` in this repository.
 
-| File             | Size    |
-| ---------------- | ------- |
-| `painformer.pth` | **TBA** |
+| File                        | Size    |
+| --------------------------- | ------- |
+| `checkpoint/painformer.pth` | **75 MB** |
+
+### Download options
 
 ```bash
-# download the latest checkpoint
-auto=https://github.com/GkikasStefanos/PainFormer/releases/latest/download/painformer.pth
-curl -L -o painformer.pth "$auto"
+# direct file download (PainFormer)
+mkdir -p checkpoint
+wget https://huggingface.co/stefanosgikas/PainFormer/resolve/main/checkpoint/painformer.pth
+```
+
+```python
+from huggingface_hub import hf_hub_download
+ckpt_path = hf_hub_download(
+    repo_id="stefanosgikas/PainFormer",
+    filename="checkpoint/painformer.pth"
+)
+print(ckpt_path)
+```
+
+Optional integrity check:
 
-# optional: verify
-sha256sum painformer.pth
+```bash
+sha256sum checkpoint/painformer.pth
 ```
 
-The checkpoint contains **one key**:
+The checkpoint contains:
 
-```text
+```
 model_state_dict    # PainFormer backbone weights
 ```
+
 ---
 
 ## Quick start
 
-> Assumes **PyTorch ≥ 2.0** and **timm ≥ 0.9** are installed.
+Assumes **PyTorch ≥ 2.0** and **timm ≥ 0.9** are installed.
+
+Repository layout expected:
+
+```
+.
+├── docs/                    # images for the model card
+├── architecture/            # Python modules (e.g., painformer.py)
+└── checkpoint/              # painformer.pth
+```
 
 ### Extract embeddings
 
 ```python
 import torch
 from timm.models import create_model
-from architecture import painformer
 from PIL import Image
 from torchvision import transforms
 
+# model code lives in the local "architecture" folder
+from architecture import painformer  # ensures registry / model class is imported
+
 # ---------------------------------------------------------------
 # Setup ---------------------------------------------------------
 # ---------------------------------------------------------------
@@ -126,8 +152,8 @@ to_tensor = transforms.Compose([
 # ---------------------------------------------------------------
 # Load PainFormer -----------------------------------------------
 # ---------------------------------------------------------------
-model = create_model('painformer').to(device)
-state = torch.load('./checkpoints/painformer.pth', map_location=device)
+model = create_model('painformer').to(device)   # class registered by architecture/painformer.py
+state = torch.load('checkpoint/painformer.pth', map_location=device)
 model.load_state_dict(state['model_state_dict'], strict=False)
 
 # expose embeddings (remove classification head)
@@ -158,15 +184,12 @@ import torch, torch.nn as nn
 from timm.models import create_model
 from architecture import painformer
 
-# ---------------------------------------------------------------
-# Setup ----------------------------------------------------------
-# ---------------------------------------------------------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 num_classes = 3  # set to your task
 
 # Backbone → 160-D embeddings
 model = create_model('painformer').to(device)
-state = torch.load('painformer.pth', map_location=device)
+state = torch.load('checkpoint/painformer.pth', map_location=device)
 model.load_state_dict(state['model_state_dict'], strict=False)
 
 # freeze if you only need fixed embeddings
@@ -182,16 +205,7 @@ head = nn.Sequential(
 optimizer = torch.optim.Adam(head.parameters(), lr=1e-3)
 criterion = nn.CrossEntropyLoss()
 
-# one step (sketch)
-def step(x, y):
-    model.eval()
-    with torch.no_grad():
-        z = model(x)      # [B, 160]
-    logits = head(z)      # [B, C]
-    loss = criterion(logits, y)
-    return loss, logits
-
-# --- optional: end-to-end fine-tune ---
+# optional: end-to-end fine-tune
 for p in model.parameters():
     p.requires_grad = True
 optimizer = torch.optim.AdamW(
@@ -221,7 +235,7 @@ optimizer = torch.optim.AdamW(
 
 ## Licence & acknowledgements
 
-* Code & weights: **MIT Licence** – see [`LICENSE`](./LICENSE)
+* Code & weights: **MIT Licence** – see `LICENSE`.
 
 ---