SegFormer is a lightweight transformer-based model for semantic segmentation. It can be fine-tuned on image-and-mask datasets to locate objects automatically, without requiring the point prompts used by Segment Anything. section 15.7.1
A practical SegFormer workflow includes:
- Prepare CT slices and corresponding segmentation masks, then create training and validation dataloaders. section 15.7.1
- Load SegformerForSemanticSegmentation from the transformers library using the pretrained nvidia/mit-b0 checkpoint. section 15.7.1
- Configure two labels, background and nodule, and fine-tune the model for the target dataset. section 15.7.1
- Train with an optimizer such as AdamW, monitor training and validation losses, save parameters, and run inference on unseen scans. section 15.9
Nodule segmentation
A pretrained SegFormer can be adapted to CT slices so it produces masks for suspected nodules. Those masks can then identify regions to pass into a classification model, forming a segmentation-followed-by-classification pipeline. section 15.0 section 15.9
# In[]:
from transformers import SegformerForSemanticSegmentation
id2label = {"0": "background", "1": "nodule"}
label2id = {v: k for k, v in id2label.items()}
model = SegformerForSemanticSegmentation.from_pretrained(
"nvidia/mit-b0",
num_labels=2,
id2label=id2label,
label2id=label2id,
)
model.to(device)
# Out[]:
Some weights of SegformerForSemanticSegmentation were not initialized from
the model checkpoint at nvidia/mit-b0 and are newly initialized:
['decode_head.batch_norm.bias', 'decode_head.batch_norm.
num_batches_tracked', 'decode_head.batch_norm.running_mean',
'decode_head.batch_norm.running_var', ...]
You should probably TRAIN this model on a down-stream task to be able to
use it for predictions and inference.
Inspect the model setup, label mappings, pretrained checkpoint, and newly initialized decoder weights. section 15.7.1
