tyfei216/TomoGNN

Cryo-ET Particle and Organelle Detection - Release Models & Example Data

This package contains pre-trained models and example tomography data for detecting particles (ribosome, HSP60) and organelles (mitochondria, nucleus) in cryo-electron tomography (cryo-ET) volumes.

📦 Contents

Models

1. Ribosome Detection (ribosome/)

• Model Files:
  • last.ckpt (262 MB): Full transformer-based detector checkpoint
  • 3DCNN.ckpt (65 MB): 3D CNN particle classifier for scoring refinement
• Configuration: config.json - Training hyperparameters and architecture specs
• Trained on: 13 tomography volumes with manual ribosome annotations
• Task: Slice-wise detection + 3D CNN scoring for ribosome localization
• Output: Ribosomal particle coordinates (z, y, x) with confidence scores

2. HSP60 Detection (hsp60/)

• Model Files:
  • last.ckpt (262 MB): Transformer-based detector checkpoint
  • 3DCNN.ckpt (29 MB): 3D CNN particle classifier
• Configuration: config.json - Training hyperparameters
• Trained on: 5 tomography volumes with HSP60 chaperonin annotations
• Task: Slice-wise detection + 3D CNN scoring for HSP60 particles
• Output: HSP60 particle coordinates with confidence scores

3. Organelle Detection (mitochondria&nucleus/)

• Model Files:
  • last.ckpt (262 MB): DETR-based detector for organelles
• Task: Detection and segmentation of mitochondria and nucleus structures
• Output: Organelle localization in 3D tomography

4. Pretrained Models (pretrained_models/)

• conditionaldetr.ckpt: Base conditional DETR model used as initialization

Example Data

All example data is located in data_example/ directory:

Tomography Volumes (MRC Format)

• ribosome.mrc (2.0 GB): Sample cryo-ET tomogram containing ribosomes
• hsp60.mrc (2.0 GB): Sample cryo-ET tomogram containing HSP60 particles
• mitochondria_nucleus.mrc (2.0 GB): Sample tomogram with organelle structures
• ribosome_label.mrc (1.0 GB): Segmentation mask for ribosome volume

Ground-Truth Annotations (TXT Format)

• ribosome_label.txt: Tab-separated coordinates [z, y, x] of ribosomal particles
• hsp60_label.txt: Tab-separated coordinates [z, y, x] of HSP60 particles
• shrec_labels.txt: Coordinates of SHREC benchmark particles

Pickle Annotations (Python Objects)

• ribosome_label.pkl (52 MB): Serialized annotation dictionary
• hsp60_label.pkl (200 MB): Serialized annotation dictionary
• hsp60_label_corrected.pkl (200 MB): Corrected annotation version

🚀 Quick Start

Prerequisites

# Create conda environment
conda create -n cryo-detection python=3.12 -y
conda activate cryo-detection

# Install dependencies
cd /path/to/cryoem
pip install -r requirements.txt
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

Usage with Notebooks

1. Load and Test Ribosome Model:

   import sys
   sys.path.append("../src")
   import utils, data
   
   model = utils.loadModel("/path/to/release_models/ribosome", "last.ckpt")
   model = model.eval().cuda(0)
   
   dataset = data.TestDatasetMrc(
       "/path/to/release_models/data_example/ribosome.mrc",
       norm="hist",
       reshape=800,
       length_for_average=3,
       gap=1
   )
   

2. Run Inference:

   • See notebooks/scan_particles.ipynb for slice-wise detection
   • See notebooks/scan_particle_with3DCNN_pipeline.ipynb for full pipeline with 3D CNN scoring
   • See revise_notebooks/SHREC.ipynb for benchmark evaluation

3. Evaluate with Ground Truth:

   • Use provided label files for validation
   • Compute metrics (AUPR, AUROC) against ground-truth coordinates

📋 Model Architecture

DETR-Based Detector

• Backbone: Transformer feature extraction on 2D slices
• Head: Transformer decoder with learnable queries (300 per image)
• Output: Bounding boxes + objectness scores per slice
• Inference: Slide through tomography volume with configurable gap (typically gap=1 for dense detection)

3D CNN Classifier

• Input: 41×41×41 (or 65×65×65) voxel crop around candidate
• Layers: Multiple 3D convolutional blocks
• Output: Binary or multi-class particle classification
• Purpose: Filter low-confidence detections and improve precision

📊 Recommended Thresholds

Model	Objectness Threshold	Sweep Distance	Matching Distance
Ribosome	0.20	15	15
HSP60	0.20	10	15
Organelle	0.30	20	20

📝 Input/Output Formats

Input

• MRC File: 3D electron microscopy volume (shape: [D, H, W])
• Normalization: Histogram equalization recommended

Output

• Predictions DataFrame: columns = [z, y, x, prediction_score, revised_z, revised_y, revised_x]
• Filtering: Apply threshold on prediction_score then spatial clustering
• Refinement: Optional 3D CNN re-scoring for final candidate ranking

🔍 Key Parameters

• reshape: Target size for memory efficiency (typically 800 for 1024×1024 slices)
• length_for_average: Number of slices to average for feature computation (typically 3)
• gap: Stride for slice sampling (1 = every slice, 2 = every other slice)
• crop_size: 3D crop size around predicted center for CNN classifier (41 or 65)
• norm: Normalization type ("hist" for histogram equalization, "std" for standardization)

📚 Referenced Notebooks

Located in /home/feity/cryoem/notebooks/:

1. scan_particles.ipynb - Basic particle detection pipeline
2. scan_particle_with3DCNN_pipeline.ipynb - Full detection + 3D CNN scoring + evaluation
3. trainModel.ipynb - Model training from scratch
4. buildDataset.ipynb - Dataset creation from raw tomograms

Located in /home/feity/cryoem/revise_notebooks/:

1. SHREC.ipynb - Benchmark evaluation on SHREC dataset
2. test_GNN.ipynb - Graph neural network post-processing tests

📖 Data Format Details

MRC Format

• Binary format for electron microscopy data
• Read via: mrcfile library (included in requirements)
• Numpy array access: mrcfile.open(path).data

Label Text Format

z  y  x
125 256 512
130 260 520
...

Pickle Format

Dictionary containing:

{
    "mapclass": {"ribosome": 0},
    "annotations": {0: {slice_idx: [instance_ids]}},
    "masks": {0: {slice_idx: scipy.sparse.csr_matrix}},
    "bboxes": {0: {slice_idx: {instance_id: [x_min, y_min, w, h]}}},
    "mrc_path": "/path/to/volume.mrc",
    "mrc_shape": (500, 1024, 1024)
}

🔗 Model Information

• Framework: PyTorch + PyTorch Lightning
• Checkpoint Format: .ckpt (Lightning checkpoint)
• Quantization: None (full precision FP32)
• Memory Requirements: ~8 GB GPU memory recommended for inference

📄 License & Citation

[Add appropriate license information]

If you use these models and data, please cite:

@article{your_paper_title,
  author={Your Authors},
  journal={Journal Name},
  year={2024}
}

🤝 Support & Contact

For issues or questions:

• Check notebook examples for usage patterns
• Review model config.json files for architecture details
• See utils.py and postprocess.py for utility functions

Troubleshooting

Model loading error: Ensure PyTorch Lightning version matches checkpoint format

pip install pytorch-lightning==2.0.0  # adjust version as needed

Out of memory: Reduce reshape parameter or use smaller batches Poor predictions: Verify input normalization matches training setup

Directory Structure

release_models/
├── README.md (this file)
├── ribosome/
│   ├── last.ckpt
│   ├── 3DCNN.ckpt
│   └── config.json
├── hsp60/
│   ├── last.ckpt
│   ├── 3DCNN.ckpt
│   └── config.json
├── mitochondria&nucleus/
│   └── last.ckpt
├── pretrained_models/
│   └── conditionaldetr.ckpt
└── data_example/
    ├── ribosome.mrc
    ├── ribosome_label.txt
    ├── ribosome_label.mrc
    ├── ribosome_label.pkl
    ├── hsp60.mrc
    ├── hsp60_label.txt
    ├── hsp60_label.pkl
    ├── hsp60_label_corrected.pkl
    ├── mitochondria_nucleus.mrc
    └── shrec_labels.txt