Update README with model loading code

README.md

@@ -14,32 +14,59 @@ Part of the [LAPVQA collection](https://huggingface.co/collections/dmusingu/lapv
 
 ## Description
 
-DETR-style detection heads for **chest abnormality detection** on VinDr-CXR (14 classes),
-trained on top of six frozen vision encoders.
-Each file contains the detection head weights; the encoder backbone is not included.
+DETR-style detection heads for 14-class chest abnormality detection on VinDr-CXR,
+trained on top of six **frozen** vision encoders.
+Each checkpoint is a dict: `{state_dict, vis_dim, d_model, num_queries, num_enc, num_dec, encoder, epoch, val_map40, val_map50}`.
 
-## Results (VinDr-CXR test set, mAP@IoU=0.4)
+## Architecture — `DetectionHead`
 
-| Encoder | mAP@0.4 (val) | mAP@0.4 (test) |
+```
+vis_proj        : Linear(vis_dim → 256)
+encoder         : 2 × TransformerEncoderLayer (self-attn, pre-norm)
+object_queries  : Parameter [1, 20, 256]
+decoder         : 3 × TransformerDecoderLayer (cross-attn to encoder output)
+class_head      : Linear(256 → 15)   # 14 classes + background
+box_head        : MLP(256 → 256 → 4) # (cx,cy,w,h) ∈ [0,1]
+```
+
+## Results (VinDr-CXR test, mAP@IoU=0.4)
+
+| Encoder | mAP@0.4 (test) |
+|---|---|
+| OWLv2 | 0.048 |
+| SigLIP | ~0.045 |
+| CLIP ViT-L/14 | ~0.040 |
+
+| File | Encoder | vis_dim |
 |---|---|---|
-| OWLv2 | ~0.14 | 0.048 |
-| SigLIP | ~0.14 | ~0.045 |
-| CLIP ViT-L/14 | ~0.14 | ~0.040 |
-| Florence-2 | ~0.14 | ~0.038 |
-| CoCa | ~0.14 | ~0.035 |
-| MAE-ViT-L/16 | ~0.14 | ~0.035 |
+| `clip-vit-l14.pt` | CLIP ViT-L/14 | 1024 |
+| `siglip.pt` | SigLIP | 1152 |
+| `florence2.pt` | Florence-2 | 1024 |
+| `coca.pt` | CoCa | 768 |
+| `owlv2.pt` | OWLv2 | 1024 |
+| `mae-vit-l16.pt` | MAE ViT-L/16 | 1024 |
 
-OWLv2 leads on detection — its pretraining objective (open-vocabulary detection) directly
-matches this task. The val→test drop (~3×) reflects distribution shift between
-multi-radiologist NMS-merged train annotations and single-consensus test labels.
+## Loading
 
-## Files
+```python
+import torch
+from lapvqa.ad.heads import DetectionHead
+from lapvqa.ad.heads import predict
 
-| 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 |
+ckpt = torch.load("owlv2.pt", map_location="cpu")
+head = DetectionHead(
+    vis_dim     = ckpt["vis_dim"],
+    d_model     = ckpt["d_model"],
+    num_queries = ckpt["num_queries"],
+    num_enc_layers = ckpt["num_enc"],
+    num_dec_layers = ckpt["num_dec"],
+)
+head.load_state_dict(ckpt["state_dict"])
+head.eval()
+
+with torch.no_grad():
+    # vis_tokens: [B, HW, vis_dim] — spatial patch tokens from the frozen encoder
+    outputs    = head(vis_tokens)
+    detections = predict(outputs, score_threshold=0.1, nms_iou=0.5)
+    # detections[i]: {'boxes': [K,4] xyxy, 'labels': [K], 'scores': [K]}
+```