ianpan
/

total-classifier

@@ -1,199 +1,70 @@
 ---
 library_name: transformers
-tags: []
 ---
-# Model Card for Model ID
-<!-- Provide a quick summary of what the model is/does. -->
-## Model Details
-### Model Description
-<!-- Provide a longer summary of what this model is. -->
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
-### Model Sources [optional]
-<!-- Provide the basic links for the model. -->
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-## Uses
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-### Direct Use
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
-[More Information Needed]
-### Downstream Use [optional]
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
-[More Information Needed]
-### Out-of-Scope Use
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
-[More Information Needed]
-## Bias, Risks, and Limitations
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-[More Information Needed]
-### Recommendations
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
-## How to Get Started with the Model
-Use the code below to get started with the model.
-[More Information Needed]
-## Training Details
-### Training Data
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
-[More Information Needed]
-### Training Procedure
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-#### Preprocessing [optional]
-[More Information Needed]
-#### Training Hyperparameters
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
-#### Speeds, Sizes, Times [optional]
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-[More Information Needed]
-## Evaluation
-<!-- This section describes the evaluation protocols and provides the results. -->
-### Testing Data, Factors & Metrics
-#### Testing Data
-<!-- This should link to a Dataset Card if possible. -->
-[More Information Needed]
-#### Factors
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-[More Information Needed]
-#### Metrics
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-[More Information Needed]
-### Results
-[More Information Needed]
-#### Summary
-## Model Examination [optional]
-<!-- Relevant interpretability work for the model goes here -->
-[More Information Needed]
-## Environmental Impact
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-## Technical Specifications [optional]
-### Model Architecture and Objective
-[More Information Needed]
-### Compute Infrastructure
-[More Information Needed]
-#### Hardware
-[More Information Needed]
-#### Software
-[More Information Needed]
-## Citation [optional]
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-**BibTeX:**
-[More Information Needed]
-**APA:**
-[More Information Needed]
-## Glossary [optional]
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-[More Information Needed]
-## More Information [optional]
-[More Information Needed]
-## Model Card Authors [optional]
-[More Information Needed]
-## Model Card Contact
-[More Information Needed]

 ---
 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.
+## 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)
+# 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.