| --- |
| library_name: transformers |
| tags: |
| - ct |
| - computed_tomography |
| - crop |
| - dicom |
| - radiology |
| license: apache-2.0 |
| base_model: |
| - timm/mobilenetv3_small_100.lamb_in1k |
| pipeline_tag: object-detection |
| --- |
| |
| This model crops the foreground from the background in CT slices. It is a lightweight `mobilenetv3_small_100` |
| model trained on CT examinations from the [public TotalSegmentator dataset](https://zenodo.org/records/10047292), version.2.0.1. |
|
|
| The following function was used to generate masks for each CT: |
| ``` |
| import nibabel as nib |
| import numpy as np |
| from scipy.ndimage import binary_closing, binary_fill_holes, minimum_filter |
| from skimage.measure import label |
| |
| def generate_mask(array): |
| mask = (array > 0).astype("uint8") |
| mask_label = label(mask) |
| labels, counts = np.unique(mask_label, return_counts=True) |
| labels, counts = labels[1:], counts[1:] |
| max_label = labels[np.argmax(counts)] |
| mask = mask_label == max_label |
| mask = np.stack([ |
| binary_fill_holes(binary_closing(mask[:, :, i])) |
| for i in range(mask.shape[2]) |
| ], axis=2).astype("uint8") |
| mask = np.stack([ |
| minimum_filter(mask[:, :, i], size=3) |
| for i in range(mask.shape[2]) |
| ], axis=2) |
| return mask |
| |
| array = nib.load("ct.nii.gz").get_fdata() |
| # apply soft tissue window |
| array = apply_ct_window(array, window_level=50, window_width=400) |
| mask = generate_mask(array) |
| ``` |
|
|
| Bounding box coordinates were generated from the masks for individual slices. |
| The model was then trained to predict normalized (0-1) `xwyh` coordinates, given an individual CT slice. |
| If the mask was empty, the coordinates were set to all zero. |
| Images were converted from Hounsfield units (HU) to 4 CT windows: |
| 1. Soft tissue (level=50, width=400) |
| 2. Brain (level=40, width=80) |
| 3. Lung (level=-600, width=1500) |
| 4. Bone (level=400, width=1800) |
|
|
| During training, random combinations of channels were selected. If more than 1 channel was selected, |
| the images were averaged channel-wise to create a single-channel output. Strong data augmentation was also applied. |
| Thus, this model should be robust to different CT windows and combinations thereof. |
|
|
| Example usage below: |
| ``` |
| import cv2 |
| import torch |
| from transformers import AutoModel |
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| cropper = AutoModel.from_pretrained("ianpan/ct-crop", trust_remote_code=True).eval().to(device) |
| |
| # single image |
| img = cv2.imread("ct_slice.png", cv2.IMREAD_GRAYSCALE) |
| cropped_img = cropper.crop(img, mode="2d", device=device, raw_hu=False, add_buffer=None) |
| |
| # expand all 4 sides by 2.5% each |
| cropped_img = cropper.crop(img, mode="2d", device=device, raw_hu=False, add_buffer=0.025) |
| |
| # expand box height by 2.5% in each direction |
| # and box width by 5% in each direction |
| buffer = (0.05, 0.025) |
| cropped_img = cropper.crop(img, mode="2d", device=device, raw_hu=False, add_buffer=buffer) |
| |
| # stack of images |
| img_list = ["ct_slice_1.png", "ct_slice_2.png", ...] |
| stack = np.stack([cv2.imread(img, cv2.IMREAD_GRAYSCALE) for img in img_list], axis=0) |
| cropped_stack = cropper.crop(img, mode="3d", device=device, raw_hu=False, add_buffer=None) |
| ``` |
|
|
| You can also get the coordinates directly and do the cropping yourself. |
| You must separately preprocess the input. Example below: |
| ``` |
| # single image |
| img0 = cv2.imread("ct_slice.png", cv2.IMREAD_GRAYSCALE) |
| img_shapes = torch.tensor([_.shape[:2] for _ in [img0]]).to(device) |
| img = cropper.preprocess(img0, mode="2d") |
| # if multi-channel, need to convert from channels-last -> channels-first |
| img = torch.from_numpy(img).expand(1, 1, -1, -1).float().to(device) |
| with torch.inference_mode(): |
| coords = cropper(img, img_shape=img_shapes, add_buffer=None) |
| |
| # if you do not provide img_shapes, output will be normalized (0-1) coordinates |
| # otherwise will be scaled to img_shape |
| ``` |
|
|
| The model also contains methods to load DICOM images, if you have `pydicom` installed: |
| ``` |
| img = cropper.load_image_from_dicom(path_to_dicom_file, windows=None) |
| |
| # note: RescaleSlope and RescaleIntercept already applied in the method |
| # apply CT window |
| brain_window = (40, 80) |
| img = cropper.load_image_from_dicom(path_to_dicom_file, windows=brain_window) |
| |
| # or multiple windows |
| soft_tissue_window = (50, 400) |
| img = cropper.load_image_from_dicom(path_to_dicom_file, |
| windows=[brain_window, soft_tissue_window]) |
| # each window is a separate channel, img will be channels-last format |
| ``` |
|
|
| You can also load a stack of DICOM images from a folder: |
| ``` |
| dicom_folder = "/path/to/ct/head/images/" |
| |
| # dicom_extension is used to filter files, default is ".dcm" |
| # can pass "" if you do not want to filter files |
| # default sort is by ImagePositionPatient using automatically determined |
| # orientation, can also sort by InstanceNumber |
| # can also apply CT windows, as above |
| stack = cropper.load_stack_from_dicom_folder(dicom_folder, |
| windows=None, |
| dicom_extension=".dcm", |
| sort_by_instance_number=False) |
| |
| # can input raw Hounsfield units into cropper |
| cropped_stack = cropper.crop(stack, mode="3d", device=device, raw_hu=True) |
| ``` |
|
|
| By default, the cropper will not remove slices in a stack, even if they are predicted to be empty. |
| You can enable this by specifying `remove_empty_slices=True`, which will also return |
| the indices in the original input of the removed empty slices. |
| ``` |
| cropped_stack, empty_slice_indices = cropper.crop(stack, mode="3d", remove_empty_slices=True) |
| ``` |
