Initial release: ViT-B/16 and ResNet-50v2 for Gravity Spy glitch classification

README.md

@@ -1,3 +1,173 @@
----
-license: cc-by-4.0
----
+# Gravity Spy Glitch Classifier: ViT-B/16 and ResNet-50v2 BiT
+
+Two deep learning models for classifying LIGO gravitational-wave detector
+glitch morphologies from Q-transform spectrograms, trained on Gravity Spy O3
+data (Zevin et al. 2017, CQG 34 064003).
+
+## Model Overview
+
+| Property | ViT-B/16 | ResNet-50v2 BiT |
+|---|---|---|
+| Architecture | Vision Transformer (patch 16, 224px) | Big Transfer ResNet-50v2 |
+| Pretrained weights | AugReg ImageNet-21k + ImageNet-1k | ImageNet-21k + ImageNet-1k |
+| Parameters | ~86M | ~25M |
+| Framework | timm | timm |
+
+Both models classify 224x224 RGB Q-transform spectrograms into 23 Gravity Spy
+glitch classes (22 glitch morphologies + No_Glitch). The class taxonomy follows
+Zevin et al. (2017).
+
+## Training Data
+
+- **Source:** Gravity Spy O3 (H1 + L1), filtered to ml_confidence > 0.9
+- **Total samples:** 227,943 training / 48,844 validation / 48,845 test
+- **Split:** Temporal split (70/15/15%) with 60-second gap enforcement
+- **Classes:** 23 (see `src/class_labels.json`)
+- **Preprocessing:** Q-transform spectrograms resized to 224x224, normalized with ImageNet statistics
+- **Rare classes:** Chirp (11 train), Wandering_Line (30), Helix (33), Light_Modulation (142)
+
+## Performance (O3 Test Set)
+
+**Primary metric: macro-F1** (averaged equally across all 23 classes).
+
+| Metric | ViT-B/16 | ResNet-50v2 BiT |
+|---|---|---|
+| **Macro-F1 (PRIMARY)** | **0.7230** [0.7031, 0.7397] | **0.6786** [0.6598, 0.6944] |
+| Rare-class macro-F1 | 0.2412 [0.2019, 0.2957] | 0.3028 [0.2085, 0.3751] |
+| Overall accuracy (secondary) | 0.9343 | 0.9181 |
+
+*95% bootstrap confidence intervals (10,000 resamples). Macro-F1 difference is
+statistically significant (p = 0.0002).*
+
+> **Note:** Overall accuracy is reported as a secondary sanity check only.
+> It is not the primary metric because it masks rare-class performance
+> differences (see "Limitations" below).
+
+### Per-Class Highlights
+
+Architecture preference is **class-morphology-dependent**, not a uniform
+advantage for either model:
+
+| Class | ViT F1 | CNN F1 | Difference | Favors |
+|---|---|---|---|---|
+| Power_Line | 0.742 | 0.235 | **+0.507** | ViT |
+| Light_Modulation | 0.859 | 0.691 | +0.168 | ViT |
+| Scratchy | 0.875 | 0.503 | +0.372 | ViT |
+| Chirp | 0.000 | 0.471 | **+0.471** | CNN |
+| Scattered_Light | 0.719 | 0.811 | +0.092 | CNN |
+| Paired_Doves | 0.613 | 0.099 | +0.514 | ViT |
+| Violin_Mode | 0.544 | 0.683 | +0.139 | CNN |
+
+The ViT achieves higher overall macro-F1 but performs **worse on rare classes**
+(macro-F1 0.241 vs CNN 0.303). This is a forbidden-proxy scenario: overall
+accuracy improvement does not extend to rare classes.
+
+## O4 Generalization
+
+Both models were evaluated on 38,587 O4a Gravity Spy spectrograms:
+
+| Metric | ViT-B/16 | ResNet-50v2 BiT |
+|---|---|---|
+| O4 macro-F1 | 0.6695 [0.6555, 0.6816] | 0.6674 [0.6567, 0.6765] |
+| Relative degradation from O3 | -7.4% | -1.7% |
+
+Both models pass the <20% degradation threshold. The CNN shows better
+generalization stability across observing runs.
+
+## Limitations
+
+- **O3-trained only:** Models have not been fine-tuned on O4 data.
+  Performance may degrade further on later observing runs with new
+  instrumental configurations.
+- **Single-view:** Models use only the 1.0-second duration Q-transform view.
+  Gravity Spy's full pipeline uses four duration views (0.5s, 1.0s, 2.0s, 4.0s).
+- **23 classes:** New glitch morphologies appearing in O4+ will be
+  misclassified into existing categories.
+- **Rare-class regression for ViT:** The ViT's macro-F1 advantage does not
+  extend to rare classes (< 200 training samples). The CNN outperforms on
+  Chirp, Helix, and Violin_Mode.
+- **Label quality:** Training labels are from ml_confidence > 0.9 filtering
+  of Gravity Spy citizen science classifications, not expert-reviewed labels.
+
+## Usage
+
+```bash
+# Install dependencies
+pip install torch timm albumentations numpy Pillow
+
+# Classify a spectrogram with ViT
+python release/src/inference.py --model vit --image path/to/spectrogram.png
+
+# Classify with CNN, showing top-5 predictions
+python release/src/inference.py --model cnn --image path/to/spectrogram.png --top-k 5
+
+# JSON output
+python release/src/inference.py --model vit --image path/to/spectrogram.png --json
+```
+
+### Input Format
+
+- PNG or JPG image of a Q-transform spectrogram
+- Any resolution (automatically resized to 224x224)
+- RGB color (grayscale images are converted to 3-channel)
+
+### Output
+
+```
+Model: ViT-B/16
+Image: spectrogram.png
+
+Predictions:
+Rank   Class                     Probability
+-------------------------------------------
+1      Blip                      0.9823
+2      Blip_Low_Frequency        0.0091
+3      Koi_Fish                  0.0034
+```
+
+## File Structure
+
+```
+release/
+  README.md                              # This file (model card)
+  checkpoints/
+    vit_b16_gravityspy_o3.pt             # ViT-B/16 weights (~983 MB)
+    resnet50v2_gravityspy_o3.pt          # ResNet-50v2 BiT weights (~270 MB)
+    checksums.sha256                     # SHA-256 checksums for integrity
+  src/
+    inference.py                         # Standalone CLI inference script
+    preprocessing.py                     # Locked evaluation transforms
+    class_labels.json                    # 23-class index-to-label mapping
+    model_config.json                    # Architecture and training config
+  examples/
+    expected_output.json                 # Validated predictions for test images
+```
+
+## Verifying Checkpoint Integrity
+
+```bash
+cd release/checkpoints
+sha256sum -c checksums.sha256
+```
+
+## Citation
+
+```bibtex
+@misc{gravitspy_vit_cnn_2026,
+  title={Vision Transformer vs. CNN for LIGO Glitch Classification:
+         Class-Morphology-Dependent Architecture Preferences},
+  author={[Authors]},
+  year={2026},
+  note={Models trained on Gravity Spy O3 data (Zevin et al. 2017, CQG 34 064003)}
+}
+```
+
+## License
+
+MIT
+
+## References
+
+- Zevin, M. et al. (2017). "Gravity Spy: integrating advanced LIGO detector
+  characterization, machine learning, and citizen science." *Classical and
+  Quantum Gravity*, 34(6), 064003.

checkpoints/checksums.sha256

@@ -0,0 +1,2 @@
+ca5635865a9ad8886e7eed2022b4b2005e439bd9033ce177d656c5068867461d  vit_b16_gravityspy_o3.pt
+65bfe7db29dce148b45dc2eba1d8a728c947c6b2432bbfc741a407801c67e9b3  resnet50v2_gravityspy_o3.pt

checkpoints/resnet50v2_gravityspy_o3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65bfe7db29dce148b45dc2eba1d8a728c947c6b2432bbfc741a407801c67e9b3
+size 282766133

checkpoints/vit_b16_gravityspy_o3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca5635865a9ad8886e7eed2022b4b2005e439bd9033ce177d656c5068867461d
+size 1029991217

examples/expected_output.json

@@ -0,0 +1,217 @@
+[
+  {
+    "filename": "qnPsv3hL2x_1.0s.png",
+    "true_label": "Blip",
+    "image_path": "data/spectrograms/Blip/qnPsv3hL2x_1.0s.png",
+    "vit_prediction": "Blip",
+    "vit_confidence": 0.990561,
+    "vit_top3": [
+      {
+        "rank": 1,
+        "class": "Blip",
+        "probability": 0.990561
+      },
+      {
+        "rank": 2,
+        "class": "Blip_Low_Frequency",
+        "probability": 0.003823
+      },
+      {
+        "rank": 3,
+        "class": "Repeating_Blips",
+        "probability": 0.001406
+      }
+    ],
+    "cnn_prediction": "Blip",
+    "cnn_confidence": 0.655181,
+    "cnn_top3": [
+      {
+        "rank": 1,
+        "class": "Blip",
+        "probability": 0.655181
+      },
+      {
+        "rank": 2,
+        "class": "Repeating_Blips",
+        "probability": 0.323093
+      },
+      {
+        "rank": 3,
+        "class": "Extremely_Loud",
+        "probability": 0.002679
+      }
+    ]
+  },
+  {
+    "filename": "EMxVe5MHAz_1.0s.png",
+    "true_label": "Power_Line",
+    "image_path": "data/spectrograms/Power_Line/EMxVe5MHAz_1.0s.png",
+    "vit_prediction": "Power_Line",
+    "vit_confidence": 0.997734,
+    "vit_top3": [
+      {
+        "rank": 1,
+        "class": "Power_Line",
+        "probability": 0.997734
+      },
+      {
+        "rank": 2,
+        "class": "No_Glitch",
+        "probability": 0.000965
+      },
+      {
+        "rank": 3,
+        "class": "Blip_Low_Frequency",
+        "probability": 0.000165
+      }
+    ],
+    "cnn_prediction": "Power_Line",
+    "cnn_confidence": 0.912969,
+    "cnn_top3": [
+      {
+        "rank": 1,
+        "class": "Power_Line",
+        "probability": 0.912969
+      },
+      {
+        "rank": 2,
+        "class": "Air_Compressor",
+        "probability": 0.057046
+      },
+      {
+        "rank": 3,
+        "class": "No_Glitch",
+        "probability": 0.007583
+      }
+    ]
+  },
+  {
+    "filename": "kwYnhvabQ8_1.0s.png",
+    "true_label": "Chirp",
+    "image_path": "data/spectrograms/Chirp/kwYnhvabQ8_1.0s.png",
+    "vit_prediction": "Blip_Low_Frequency",
+    "vit_confidence": 0.944321,
+    "vit_top3": [
+      {
+        "rank": 1,
+        "class": "Blip_Low_Frequency",
+        "probability": 0.944321
+      },
+      {
+        "rank": 2,
+        "class": "Blip",
+        "probability": 0.033233
+      },
+      {
+        "rank": 3,
+        "class": "Repeating_Blips",
+        "probability": 0.005682
+      }
+    ],
+    "cnn_prediction": "Chirp",
+    "cnn_confidence": 0.514547,
+    "cnn_top3": [
+      {
+        "rank": 1,
+        "class": "Chirp",
+        "probability": 0.514547
+      },
+      {
+        "rank": 2,
+        "class": "Blip_Low_Frequency",
+        "probability": 0.274587
+      },
+      {
+        "rank": 3,
+        "class": "Blip",
+        "probability": 0.14339
+      }
+    ]
+  },
+  {
+    "filename": "mRGHzX9W9X_1.0s.png",
+    "true_label": "No_Glitch",
+    "image_path": "data/spectrograms/No_Glitch/mRGHzX9W9X_1.0s.png",
+    "vit_prediction": "No_Glitch",
+    "vit_confidence": 0.991169,
+    "vit_top3": [
+      {
+        "rank": 1,
+        "class": "No_Glitch",
+        "probability": 0.991169
+      },
+      {
+        "rank": 2,
+        "class": "1400Ripples",
+        "probability": 0.004831
+      },
+      {
+        "rank": 3,
+        "class": "Power_Line",
+        "probability": 0.000853
+      }
+    ],
+    "cnn_prediction": "No_Glitch",
+    "cnn_confidence": 0.989384,
+    "cnn_top3": [
+      {
+        "rank": 1,
+        "class": "No_Glitch",
+        "probability": 0.989384
+      },
+      {
+        "rank": 2,
+        "class": "1400Ripples",
+        "probability": 0.003378
+      },
+      {
+        "rank": 3,
+        "class": "Helix",
+        "probability": 0.002471
+      }
+    ]
+  },
+  {
+    "filename": "XvHH2XRheR_1.0s.png",
+    "true_label": "Scattered_Light",
+    "image_path": "data/spectrograms/Scattered_Light/XvHH2XRheR_1.0s.png",
+    "vit_prediction": "Scattered_Light",
+    "vit_confidence": 0.995357,
+    "vit_top3": [
+      {
+        "rank": 1,
+        "class": "Scattered_Light",
+        "probability": 0.995357
+      },
+      {
+        "rank": 2,
+        "class": "Extremely_Loud",
+        "probability": 0.001687
+      },
+      {
+        "rank": 3,
+        "class": "Fast_Scattering",
+        "probability": 0.000669
+      }
+    ],
+    "cnn_prediction": "Scattered_Light",
+    "cnn_confidence": 0.962074,
+    "cnn_top3": [
+      {
+        "rank": 1,
+        "class": "Scattered_Light",
+        "probability": 0.962074
+      },
+      {
+        "rank": 2,
+        "class": "Low_Frequency_Lines",
+        "probability": 0.017069
+      },
+      {
+        "rank": 3,
+        "class": "Extremely_Loud",
+        "probability": 0.004967
+      }
+    ]
+  }
+]

src/class_labels.json

@@ -0,0 +1 @@
+["1080Lines", "1400Ripples", "Air_Compressor", "Blip", "Blip_Low_Frequency", "Chirp", "Extremely_Loud", "Fast_Scattering", "Helix", "Koi_Fish", "Light_Modulation", "Low_Frequency_Burst", "Low_Frequency_Lines", "No_Glitch", "Paired_Doves", "Power_Line", "Repeating_Blips", "Scattered_Light", "Scratchy", "Tomte", "Violin_Mode", "Wandering_Line", "Whistle"]

src/inference.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Standalone inference script for Gravity Spy glitch classification.
+
+Classifies LIGO glitch spectrograms using either a ViT-B/16 or ResNet-50v2 BiT
+model trained on O3 Gravity Spy data (23 classes).
+
+Usage:
+    python inference.py --model vit --image path/to/spectrogram.png
+    python inference.py --model cnn --image path/to/spectrogram.png --top-k 5
+
+Requirements: torch, timm, albumentations, numpy, Pillow
+No training code imports required.
+"""
+# ASSERT_CONVENTION: primary_metric=macro_f1, input_format=224x224_RGB_PNG_0to1
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import timm
+
+from preprocessing import load_image
+
+
+def load_model(model_key, config_path=None, checkpoint_dir=None):
+    """Load a model from config and checkpoint.
+
+    Parameters
+    ----------
+    model_key : str
+        Either 'vit' or 'cnn'.
+    config_path : Path or None
+        Path to model_config.json. Defaults to same directory as this script.
+    checkpoint_dir : Path or None
+        Path to checkpoints directory. Defaults to ../checkpoints/ relative to this script.
+
+    Returns
+    -------
+    model : torch.nn.Module
+        Loaded model in mode for inference.
+    class_labels : list of str
+        Ordered class label names (index -> label).
+    """
+    script_dir = Path(__file__).resolve().parent
+    if config_path is None:
+        config_path = script_dir / "model_config.json"
+    if checkpoint_dir is None:
+        checkpoint_dir = script_dir.parent / "checkpoints"
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    if model_key not in config["models"]:
+        available = list(config["models"].keys())
+        raise ValueError(f"Unknown model key '{model_key}'. Choose from: {available}")
+
+    model_cfg = config["models"][model_key]
+
+    # Build architecture via timm (no pretrained weights -- we load our own)
+    model = timm.create_model(
+        model_cfg["model_name"],
+        pretrained=False,
+        num_classes=model_cfg["num_classes"],
+    )
+
+    # Load trained weights
+    checkpoint_path = checkpoint_dir / model_cfg["checkpoint_file"]
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {checkpoint_path}\n"
+            f"Expected file: {model_cfg['checkpoint_file']}"
+        )
+
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    # Handle both full training checkpoints (with 'model_state_dict' key)
+    # and bare state_dict files
+    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+        state_dict = checkpoint["model_state_dict"]
+    else:
+        state_dict = checkpoint
+    model.load_state_dict(state_dict)
+
+    # Set to inference mode
+    model.requires_grad_(False)
+    model.eval()
+
+    # Load class labels
+    labels_path = script_dir / "class_labels.json"
+    with open(labels_path) as f:
+        class_labels = json.load(f)
+
+    return model, class_labels
+
+
+def predict(model, image_tensor, class_labels, top_k=3):
+    """Run inference on a preprocessed image tensor.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        Model in inference mode.
+    image_tensor : torch.Tensor
+        Preprocessed image of shape (3, 224, 224).
+    class_labels : list of str
+        Ordered class label names.
+    top_k : int
+        Number of top predictions to return.
+
+    Returns
+    -------
+    predictions : list of dict
+        Top-k predictions, each with 'rank', 'class', 'probability'.
+    """
+    with torch.no_grad():
+        logits = model(image_tensor.unsqueeze(0))  # (1, num_classes)
+        probs = torch.softmax(logits, dim=1).squeeze(0)  # (num_classes,)
+
+    top_k = min(top_k, len(class_labels))
+    top_probs, top_indices = torch.topk(probs, top_k)
+
+    predictions = []
+    for rank, (prob, idx) in enumerate(zip(top_probs, top_indices), 1):
+        predictions.append({
+            "rank": rank,
+            "class": class_labels[idx.item()],
+            "probability": round(prob.item(), 6),
+        })
+
+    return predictions
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Gravity Spy glitch classification inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=(
+            "Examples:\n"
+            "  python inference.py --model vit --image spectrogram.png\n"
+            "  python inference.py --model cnn --image spectrogram.png --top-k 5\n"
+            "  python inference.py --model vit --image spectrogram.png --json\n"
+        ),
+    )
+    parser.add_argument("--model", required=True, choices=["vit", "cnn"],
+                        help="Model to use: vit (ViT-B/16) or cnn (ResNet-50v2 BiT)")
+    parser.add_argument("--image", required=True, type=str,
+                        help="Path to spectrogram image (PNG/JPG)")
+    parser.add_argument("--top-k", type=int, default=3,
+                        help="Number of top predictions to show (default: 3)")
+    parser.add_argument("--json", action="store_true",
+                        help="Output predictions as JSON")
+    args = parser.parse_args()
+
+    # Validate image path
+    image_path = Path(args.image)
+    if not image_path.exists():
+        print(f"Error: Image not found: {image_path}", file=sys.stderr)
+        sys.exit(1)
+
+    # Load model
+    model, class_labels = load_model(args.model)
+
+    # Preprocess image
+    image_tensor = load_image(str(image_path))
+
+    # Run inference
+    predictions = predict(model, image_tensor, class_labels, top_k=args.top_k)
+
+    # Output results
+    if args.json:
+        result = {
+            "model": args.model,
+            "image": str(image_path),
+            "predictions": predictions,
+        }
+        print(json.dumps(result, indent=2))
+    else:
+        model_name = "ViT-B/16" if args.model == "vit" else "ResNet-50v2 BiT"
+        print(f"\nModel: {model_name}")
+        print(f"Image: {image_path}")
+        print(f"\nPredictions:")
+        print(f"{'Rank':<6} {'Class':<25} {'Probability':<12}")
+        print("-" * 43)
+        for pred in predictions:
+            print(f"{pred['rank']:<6} {pred['class']:<25} {pred['probability']:.4f}")
+        print()
+
+
+if __name__ == "__main__":
+    main()

src/model_config.json

@@ -0,0 +1,35 @@
+{
+  "models": {
+    "vit": {
+      "model_name": "vit_base_patch16_224.augreg_in21k_ft_in1k",
+      "architecture": "ViT-B/16",
+      "num_classes": 23,
+      "pretrained_source": "AugReg ImageNet-21k fine-tuned on ImageNet-1k",
+      "checkpoint_file": "vit_b16_gravityspy_o3.pt",
+      "input_size": 224,
+      "normalization": {
+        "mean": [0.485, 0.456, 0.406],
+        "std": [0.229, 0.224, 0.225]
+      }
+    },
+    "cnn": {
+      "model_name": "resnetv2_50x1_bit.goog_in21k_ft_in1k",
+      "architecture": "ResNet-50v2 BiT",
+      "num_classes": 23,
+      "pretrained_source": "ImageNet-21k fine-tuned on ImageNet-1k",
+      "checkpoint_file": "resnet50v2_gravityspy_o3.pt",
+      "input_size": 224,
+      "normalization": {
+        "mean": [0.485, 0.456, 0.406],
+        "std": [0.229, 0.224, 0.225]
+      }
+    }
+  },
+  "training": {
+    "dataset": "Gravity Spy O3",
+    "num_classes": 23,
+    "split": "temporal 70/15/15% with 60s gap enforcement",
+    "train_samples": 227943,
+    "seed": 42
+  }
+}

src/preprocessing.py

@@ -0,0 +1,66 @@
+"""Standalone preprocessing for Gravity Spy spectrogram inference.
+
+Extracts eval_transforms() from the training pipeline with NO training
+dependencies (no wandb, no dataloader, no training-specific imports).
+
+Preprocessing is locked to match training exactly:
+  - Resize to 224x224
+  - Normalize with ImageNet statistics
+  - Convert to PyTorch tensor
+"""
+# ASSERT_CONVENTION: primary_metric=macro_f1, input_format=224x224_RGB_PNG_0to1
+
+import numpy as np
+from PIL import Image
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+# ImageNet statistics for pretrained model normalization
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
+def eval_transforms(image_size=224):
+    """Evaluation transform -- resize + normalize only, no augmentation.
+
+    This is identical to the eval_transforms used during training/validation.
+    Input images are expected to be RGB numpy arrays with pixel values in [0, 255].
+    Output tensors have pixel values normalized by ImageNet statistics.
+
+    Parameters
+    ----------
+    image_size : int
+        Target spatial dimension (default 224 for ViT-B/16 and ResNet-50v2).
+
+    Returns
+    -------
+    transform : albumentations.Compose
+        Evaluation transform pipeline.
+    """
+    return A.Compose([
+        A.Resize(image_size, image_size),
+        A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+        ToTensorV2(),
+    ])
+
+
+def load_image(image_path, image_size=224):
+    """Load an image file and apply evaluation transforms.
+
+    Parameters
+    ----------
+    image_path : str
+        Path to a PNG/JPG spectrogram image.
+    image_size : int
+        Target spatial dimension (default 224).
+
+    Returns
+    -------
+    tensor : torch.Tensor
+        Preprocessed image tensor of shape (3, image_size, image_size).
+    """
+    img = Image.open(image_path).convert("RGB")
+    img_np = np.array(img)  # shape (H, W, 3), dtype uint8, values [0, 255]
+    transform = eval_transforms(image_size)
+    transformed = transform(image=img_np)
+    return transformed["image"]  # torch.Tensor (3, 224, 224)