|
|
--- |
|
|
license: apache-2.0 |
|
|
tags: |
|
|
- pytorch |
|
|
- unet |
|
|
- chest-ct |
|
|
- survival-analysis |
|
|
- time-to-event |
|
|
- model-3d |
|
|
model-index: |
|
|
- name: UNet-TTE |
|
|
results: [] |
|
|
--- |
|
|
|
|
|
|
|
|
# SwinUNETR Checkpoint |
|
|
|
|
|
This is a PyTorch Lightning `.ckpt` checkpoint for a SwinUNETR model trained on chest CT images with TTE objective. |
|
|
|
|
|
## Usage |
|
|
|
|
|
A quickstart script is below. |
|
|
|
|
|
```python |
|
|
import torch |
|
|
from src.networks import SwinUNETRForClassification |
|
|
swin_unetr_params = { |
|
|
"img_size": (224, 224, 224), |
|
|
"in_channels": 1, |
|
|
"out_channels": 2, |
|
|
"feature_size": 48, |
|
|
"drop_rate": 0.0, |
|
|
"attn_drop_rate": 0.0, |
|
|
"dropout_path_rate": 0.0, |
|
|
"use_checkpoint": True, |
|
|
} |
|
|
model = SwinUNETRForClassification( |
|
|
swin_unetr_params=swin_unetr_params, num_classes=2 |
|
|
).to(device) |
|
|
state_dict = torch.load( |
|
|
loadmodel_path, map_location=f"cuda:{torch.cuda.current_device()}" |
|
|
) |
|
|
model.load_state_dict(state_dict) |
|
|
``` |
|
|
|
|
|
For detailed instructions please follow the [README in Github repo](https://github.com/som-shahlab/tte-pretraining/tree/main?tab=readme-ov-file#evaluation). |
