Update README.md

README.md

@@ -8,9 +8,10 @@ tags:
 - retinal-imaging
 ---
 
-This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
-- Library: ronnief1/RetFiner
-- Docs: [More Information Needed]
+## RetFiner-VisionFM
+This repo contains the weights of RetFiner-VisionFM from the paper [RetFiner: A Vision-Language Refinement Scheme for Retinal Foundation Models](https://arxiv.org/abs/2506.22149).
+
+Project page: [RetFiner](https://github.com/ronnief1/RetFiner)
 
 ### Required Libraries
 This model requires specific Python libraries:
@@ -20,13 +21,32 @@ timm==0.4.12
 torchvision==0.19.1+cu118
 ```
 
+To use the model, please download the ViT implementation from [visionfm_hf.py](https://github.com/ronnief1/RetFiner/blob/main/RetFiner/visionfm_hf.py)
+
 Note: if you are using this for downstream inference, adjust the num_classes based on the target dataset.
 
 ```python
 import torch
 import torch.nn as nn
 from huggingface_hub import PyTorchModelHubMixin
-from vit import VisionTransformer, vit_base_patch16, vit_large_patch16
+from visionfm_hf import VisionTransformer
+
+class RetFiner(nn.Module, PyTorchModelHubMixin):
+    """RetFiner: Fine-tuned ViT models for retinal image analysis"""
+    
+    def __init__(self, model_name: str = "RetFiner-VisionFM", num_classes: int = 2, **kwargs):
+        super().__init__()
+        
+        self.model = VisionTransformer(return_all_tokens=True, qkv_bias=True)
+        
+        self.config = {
+            **kwargs
+        }
+
+    def forward(self, x):
+        return self.model(x)
+
+model = RetFiner.from_pretrained('ronnief1/RetFiner-VisionFM')
 ```
 
 ## Citation