Update README.md

README.md

@@ -1,12 +1,61 @@
+---
+license: cc-by-nc-sa-4.0
+datasets:
+- mwalmsley/gz2
+metrics:
+- accuracy 
+---
+
 # CosmoFormer Model
 
-This is a TorchScript version of our CrossFormer-based model.
-You can load it in plain PyTorch without needing `vit_pytorch`.
+This is a **TorchScript** version of our CrossFormer-based image classification model.  
+It was trained on [Galaxy Zoo 2 (GZ2)](https://www.zooniverse.org/projects/zookeeper/galaxy-zoo/about/research) data to classify galaxy morphologies (spirals, ellipticals, and other morphological types).
+I also leveraged the [galaxy-datasets pip package](https://github.com/mwalmsley/galaxy-datasets) by [Michael Walmsley](https://github.com/mwalmsley) for data loading and handling.
+
+
+## Model Details
+
+- **Architecture:** [CrossFormer](https://github.com/lucidrains/vit-pytorch) variant  
+- **Input Resolution:** 224×224 RGB  
+- **Number of Classes:** Depends on your label encoder (e.g., galaxy morphology classes)  
+- **Checkpoint Format:** TorchScript (`.pt`) file  
+- **Frameworks:** Originally in PyTorch with `vit_pytorch`. Now self-contained in TorchScript.  
+
+## Usage
+
+You can load and run this model **directly in PyTorch** **without** installing `vit_pytorch`. Just make sure you have an environment with:
+
+- `torch` >= 1.13.0  
+- `torchvision` (optional, if you need standard transforms)  
+
+### Quick Start Example
 
-Usage Example:
 ```python
 import torch
-model = torch.jit.load('cosmoformer_traced.pt')
+
+# 1. Load the model
+model = torch.jit.load("cosmoformer_traced.pt")
 model.eval()
-...
-```
+
+# 2. Inference
+# Suppose you have a 3-channel image tensor (1, 3, 224, 224)
+dummy_input = torch.randn(1, 3, 224, 224)
+
+with torch.no_grad():
+    outputs = model(dummy_input)
+
+print(outputs.shape)  # e.g., [1, num_classes]
+
+
+@article{10.1093/mnras/stt1458,
+author = {Willett, Kyle W. and Lintott, Chris J. and Bamford, Steven P. and Masters, Karen L. and Simmons, Brooke D. and Casteels, Kevin R. V. and Edmondson, Edward M. and Fortson, Lucy F. and Kaviraj, Sugata and Keel, William C. and Melvin, Thomas and Nichol, Robert C. and Raddick, M. Jordan and Schawinski, Kevin and Simpson, Robert J. and Skibba, Ramin A. and Smith, Arfon M. and Thomas, Daniel},
+title = "{Galaxy Zoo 2: detailed morphological classifications for 304 122 galaxies from the Sloan Digital Sky Survey}",
+journal = {Monthly Notices of the Royal Astronomical Society},
+volume = {435},
+number = {4},
+pages = {2835-2860},
+year = {2013},
+month = {09},
+issn = {0035-8711},
+doi = {10.1093/mnras/stt1458},
+}