| --- |
| license: mit |
| library_name: pytorch |
| tags: |
| - biology |
| - protein-localization |
| - diffusion-model |
| - 3d |
| - cell-imaging |
| pipeline_tag: image-text-to-image |
| --- |
| |
| # DeepGPS-3D — baseline model weights |
|
|
| Pretrained weights for the **baseline DeepGPS-3D** model: a conditional 3D |
| denoising diffusion model that predicts a protein's 3D subcellular localization |
| volume from a matched nuclear 3D volume and the protein's ESM2 sequence |
| embedding. |
|
|
| - **Code / usage:** <https://github.com/royal-dargon/DeepGPS_3D> |
| - **File:** `best_model.pth` (~185 MB) |
|
|
| ## Architecture |
|
|
| - **Diffusion:** DDPM, 1000 linear noise steps (β 1e-4 → 0.02). |
| - **Backbone:** conditional 3D U-Net, channel widths 32/64/128/256, two residual |
| blocks per level, 8-head attention at the bottleneck. Input is two spatial |
| channels (noisy protein volume + nucleus condition); output is the predicted |
| noise for a single protein channel. |
| - **Sequence conditioning:** precomputed ESM2-650M |
| (`facebook/esm2_t33_650M_UR50D`) 1,280-d embedding injected at every U-Net |
| scale. |
| - **Volume size:** 64 × 144 × 144 voxels. |
| - **Sampling:** classifier-free guided reverse diffusion (default guidance 3.0), |
| retaining nucleus and sequence conditions throughout. |
|
|
| ## Training data |
|
|
| OpenCell confocal volumes, protein-disjoint split: 114,855 train images |
| (1,048 proteins) / 14,280 val (131) / 14,282 test (131). Trained with AdamW |
| (lr 1e-4), effective batch size 12 (batch 6 × grad-accum 2), up to 50 epochs, |
| seed 42. |
|
|
| ## Usage |
|
|
| ```bash |
| git clone https://github.com/royal-dargon/DeepGPS_3D |
| cd DeepGPS_3D |
| pip install -r requirements.txt |
| python scripts/download_weights.py --output weights/best_model.pth |
| |
| python -m src.demo_infer \ |
| --checkpoint weights/best_model.pth \ |
| --nuclei example/WASL/WASL_nuclei.npy \ |
| --embedding example/WASL/WASL_esm_650M_embedding.npz \ |
| --reference example/WASL/WASL_reference.npz \ |
| --output outputs/WASL_pred --device cuda |
| ``` |
|
|
| Or load directly: |
|
|
| ```python |
| import torch |
| from src.model.diffusion import get_diffusion_model |
| from src.config import MODEL_CONFIG |
| |
| model = get_diffusion_model(MODEL_CONFIG) |
| ckpt = torch.load("best_model.pth", map_location="cpu", weights_only=False) |
| state = ckpt["models"]["diffusion"] if isinstance(ckpt, dict) and "models" in ckpt else ckpt |
| model.load_state_dict(state.state_dict() if hasattr(state, "state_dict") else state) |
| model.eval() |
| # pred = model.sample_given_nuclei(nuclei[B,1,64,144,144], seq_emb[B,1280], guidance_scale=3.0) |
| ``` |
|
|
| ## Intended use & limitations |
|
|
| Intended for research on protein subcellular localization and as the generative |
| backbone for the downstream mutation-localization and image-based PPI tasks in |
| the repository. The model reproduces coarse axial organization well and moderate |
| voxel/radial structure; it captures localization signal rather than exact |
| single-cell reconstruction. It is trained on one OpenCell split and reports |
| point estimates without confidence intervals. Not for clinical use. |
|
|
| ## License |
|
|
| MIT. See the repository for details. |
|
|
