---
license: cc-by-nc-sa-4.0
datasets:
- mwalmsley/gz2
metrics:
- accuracy 
---
# Important Disclaimer
This model is a part of my **bachelor thesis** (VUT in Brno, FIT).  

# CosmoFormer Model

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
- **Model accuracy:** 75%  
- **Input Resolution:** 224×224 RGB  
- **Number of Classes:** 8 (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

```python
import torch
import torchvision.transforms.v2 as v2
from huggingface_hub import hf_hub_download
from PIL import Image

label_mapping = {
    0: 'barred_spiral',
    1: 'edge_on_disk',
    2: 'featured_without_bar_or_spiral',
    3: 'irregular',
    4: 'smooth_cigar',
    5: 'smooth_inbetween',
    6: 'smooth_round',
    7: 'unbarred_spiral'
}

# 1. Define the path to the hugging face repo
ts_path = hf_hub_download(
    repo_id="artursultanov/cosmoformer-model",
    filename="cosmoformer_traced_cpu.pt"
)

# 2. Load the model from the hugging face repo
model = torch.jit.load(ts_path, map_location="cpu")
model.eval()

# 3. Define image transform to match model's internal representation 
transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
])

# 4. Load the image
image_path = "test_image.jpg"
image = Image.open(image_path).convert("RGB")

tensor = transform(image)  # shape [3, 224, 224]
tensor = tensor.unsqueeze(0)  shape [1, 3, 224, 224]

# 5. Inference
with torch.no_grad():
    output = model(tensor)
    predicted_idx = torch.argmax(output, dim=1).item()

predicted_label = label_mapping[predicted_idx]
print("Predicted class:", predicted_label)
```

```
@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},
}
```