Update README with model loading code

README.md

@@ -13,37 +13,54 @@ Part of the [LAPVQA collection](https://huggingface.co/collections/dmusingu/lapv
 
 ## Description
 
-Phrase grounding heads trained on MIMIC-CXR, predicting the bounding box of a described
-abnormality given the image and a text phrase (e.g. "Pleural Effusion").
-Six frozen encoder backbones are covered; each file contains only the grounding head weights.
+TransVG-style phrase grounding heads trained on MIMIC-CXR, predicting the bounding
+box of a described abnormality given the chest X-ray and a text phrase.
+Each checkpoint is a dict: `{state_dict, vis_dim, txt_dim, d_model, num_layers, encoder, epoch, val_miou, val_acc50}`.
 
-## Results (MIMIC-CXR test set)
+## Architecture — `VisualGroundingHead`
+
+```
+vis_proj   : Linear(vis_dim → 256)
+txt_proj   : Linear(txt_dim → 256)
+reg_token  : Parameter [1, 1, 256]
+sequence   : [REG | vis_tokens | txt_token]
+transformer: 3 × TransformerEncoderLayer (self-attn, pre-norm)
+box_head   : MLP(256 → 256 → 4)   # sigmoid → (cx,cy,w,h) ∈ [0,1]
+```
 
-**Zero-shot (no fine-tuning):**
+## Results (MIMIC-CXR test set)
 
-| Encoder | mIoU | Acc@0.25 |
-|---|---|---|
-| SigLIP | 0.086 | 0.042 |
-| Florence-2 | 0.089 | 0.046 |
-| CLIP ViT-L/14 | 0.085 | 0.028 |
-| CoCa | 0.082 | 0.023 |
-| OWLv2 | 0.082 | 0.023 |
+**Zero-shot:** mIoU ≈ 0.082–0.089 across all encoders.
 
-**Fine-tuned (MAE-ViT-L/16):**
+**Fine-tuned (MAE-ViT-L/16):** mIoU 0.320, Acc@0.25 0.569, Pointing Acc 0.593.
 
-| mIoU | Acc@0.25 | Acc@0.50 | Pointing Acc |
+| File | Encoder | vis_dim | txt_dim |
 |---|---|---|---|
-| 0.320 | 0.569 | 0.273 | 0.593 |
+| `clip-vit-l14.pt` | CLIP ViT-L/14 | 1024 | 768 |
+| `siglip.pt` | SigLIP | 1152 | 1152 |
+| `florence2.pt` | Florence-2 | 1024 | 768 |
+| `coca.pt` | CoCa | 768 | 768 |
+| `owlv2.pt` | OWLv2 | 1024 | 768 |
+| `mae-vit-l16.pt` | MAE ViT-L/16 | 1024 | 768 |
+
+## Loading
 
-Fine-tuning provides a ~4× improvement over zero-shot across all encoders.
+```python
+import torch
+from lapvqa.pg.heads import VisualGroundingHead
 
-## Files
+ckpt = torch.load("mae-vit-l16.pt", map_location="cpu")
+head = VisualGroundingHead(
+    vis_dim    = ckpt["vis_dim"],
+    txt_dim    = ckpt["txt_dim"],
+    d_model    = ckpt["d_model"],
+    num_layers = ckpt["num_layers"],
+)
+head.load_state_dict(ckpt["state_dict"])
+head.eval()
 
-| File | Encoder backbone |
-|---|---|
-| `clip-vit-l14.pt` | CLIP ViT-L/14 |
-| `siglip.pt` | SigLIP |
-| `florence2.pt` | Florence-2 |
-| `coca.pt` | CoCa |
-| `owlv2.pt` | OWLv2 |
-| `mae-vit-l16.pt` | MAE ViT-L/16 (fine-tuned) |
+with torch.no_grad():
+    # vis_tokens: [B, HW, vis_dim] — spatial patch tokens from frozen encoder
+    # txt_vec:    [B, txt_dim]     — pooled text representation from frozen encoder
+    pred_boxes = head(vis_tokens, txt_vec)  # [B, 4] (cx,cy,w,h) in [0,1]
+```