|
|
--- |
|
|
library_name: transformers |
|
|
tags: |
|
|
- radiology |
|
|
- ct |
|
|
- organ |
|
|
- classification |
|
|
license: apache-2.0 |
|
|
base_model: |
|
|
- timm/tf_efficientnetv2_b0.in1k |
|
|
pipeline_tag: image-classification |
|
|
--- |
|
|
|
|
|
# TotalClassifier: Slice-Level Organ Classification for CT Examinations |
|
|
|
|
|
TotalClassifier is a classification model which predicts the presence of various organs on a 2D slice from a CT volume. |
|
|
It supports axial, sagittal, and coronal images, and a variety of windowing parameters. |
|
|
This model uses a `tf_efficientnetv2_b0` backbone with a gated recurrent unit (GRU) head which performs sequence modeling across extracted slice-level features. |
|
|
The model also works with single 2D images. |
|
|
|
|
|
The model is trained on the publicly available [TotalSegmentator dataset](https://zenodo.org/records/10047292), version 2.0.1. It predicts 117 labels corresponding to the |
|
|
available labels from TotalSegmentator. The classification labels were generated from the provided segmentation labels. |
|
|
|
|
|
Note that the model expects one channel. If you create a multi-channel image using multiple CT windows, simply take the mean across channels. |
|
|
The model also expects 8-bit input (converted to float). Thus if your CT volume is in Hounsfield units, you can apply a standard window, |
|
|
such as soft tissue (level=50, width=400), before inputting it into the model. |
|
|
|
|
|
## Example Usage |
|
|
|
|
|
``` |
|
|
import torch |
|
|
from transformers import AutoModel |
|
|
|
|
|
device = "cuda" |
|
|
organ_model = AutoModel.from_pretrained("ianpan/total-classifier", trust_remote_code=True).eval().to(device) |
|
|
|
|
|
# can use model to load CT from folder with DICOM files, if pydicom is installed |
|
|
# here we apply soft tissue window |
|
|
ct_volume = organ_model.load_stack_from_dicom_folder("/path/to/dicom/folder", windows=[[50, 400]], dicom_extension=".dcm") |
|
|
|
|
|
# ct_volume.shape is (num_slices, height, width, num_channels) if applying windows |
|
|
# otherwise is (num_slices, height, width) if using original Hounsfield units |
|
|
|
|
|
# preprocess |
|
|
x = model.preprocess(ct_volume, mode="3d", torchify=True, add_batch_dim=True, device=device) |
|
|
|
|
|
# here, ct_volume is a numpy array |
|
|
# if you are loading volumes as torch.Tensors, then you can skip the preprocess function |
|
|
# and just resize the volume to height and width of 256 x 256 |
|
|
|
|
|
# x is now torch.Tensor with shape (1, num_slices, num_channels, height, width) |
|
|
# note that these are the expected dims for the model's forward method |
|
|
|
|
|
with torch.inference_mode(): |
|
|
out = organ_model(x) |
|
|
out_df = organ_model(x, return_as_df=True) |
|
|
|
|
|
# out is a torch.Tensor of shape (1, num_slices, 117) containing scores [0-1] for each organ label |
|
|
# out_df is a list of pandas DataFrames with shape (num_slices, 117), where column names are the organ names |
|
|
# each element of the list corresponds to each sample in the batch |
|
|
# however if using batch sizes >1, then all samples need to be padded to the same number of slices |
|
|
|
|
|
# you can use out_df to only get slices with predicted organ labels greater than a certain threshold |
|
|
out_df = out_df[0] |
|
|
threshold = 0.5 |
|
|
liver_indices = np.where(out_df["liver"].values >= threshold)[0] |
|
|
|
|
|
# or slices where at least one of the specified organ labels is greater than threshold |
|
|
organs_of_interest = ["liver", "spleen", "pancreas"] |
|
|
threshold = 0.5 |
|
|
slice_indices = np.where((out_df[organs_of_interest].values >= threshold).max(1))[0] |
|
|
|
|
|
# organ_model.label2index can be used to convert organ label names to the indices 0-116 |
|
|
# organ_model.index2label is the inverse |
|
|
``` |
|
|
|
|
|
If you have a large number of slices and limited GPU memory, you can either process the volume in chunks, |
|
|
or downsample the volume along the slice dimension and interpolate the predictions back to the original number of slices. |
