DR-UNet: A robust deep learning segmentation method for hematoma volumetric detection in intracerebral hemorrhage

Structure of DR-UNet

This is an implementation of DR-UNet on Python 3.6, Keras, and TensorFlow2.0. DR-UNet consists of an encoding (down-sampling) path and a decoding (up-sampling) path. The entire workflow of our computational analysis.
To increase the segmentation performance of the model, three reduced dimensional residual convolution units (RDRCUs) were developed to replace the traditional convolution layer. The three convolution blocks are illustrated in the following figure. The three RDRCUs have two branches (main branch and side branch) to process the input characteristics continuously.

Experimental results of hematoma segmentation

We first trained DR-UNet to recognize the hematoma region in patients. The performance was evaluated on two testing datasets (internal and external) using the following criteria: i) sensitivity, ii) specificity, iii) precision, iv) Dice, v) Jaccard and vi) VOE (details in the Methods section). Moreover, we compared DR-UNet with UNet, FCM and active contours. In all four methods, segmentation labeling was considered the ground truth standard (details in the Methods section). The main calculation results are shown in the figure and table below.

As shown in the table below, results of sensitivity, specificity, precision, Dice, Jaccard and VOE by four methods in the internal testing and the external testing dataset.

The internal testing dataset in the retrospective dataset was enriched to include all ICH subtypes. In Figure B, four different types of hematomas were included, and we visually presented a performance comparison among the DR-UNet, UNet, FCM and active contour methods.

Irregularly shaped and epidural hematoma

Results of hematoma volumetric analysis in (A) irregularly shaped hematoma group and (B) subdural and epidural hematoma group. Input ICH images with manual segmentation were denoted by red line. The segmented outputs of the DR-UNet model were denoted by blue lines, the segmented outputs of the UNet were denoted by green lines.

Four examples of ICH segmentation in the subdural and epidural hematomas with original images and partially enlarged image in the prospective dataset.

Four examples of ICH segmentation in the irregular-shaped hematomas with original images and partially enlarged image in the prospective dataset.

Calculate the hematoma volume experiment

The hematoma volumetric analysis by DR-UNet, UNet and Coniglobus method. A. The diagnoses of HVs were presented by ground truth and three different methods. B. The correlation plots among ground truth and three different methods. C. The error curves of three methods were plotted. The error curves of three methods were plotted. They were the error of {ground truth – the measurement by DR-UNet}, the error of {ground truth – the measurement by UNet} and the error of {ground truth – the measurement by Coniglobus formula}, respectively. D. The results of RMSE, SD, MAE and averaged time (second/scan).

Getting Started

data.py Used to make your own dataset. Making your own dataset needs to satisfy having original images and the ground truth images. The completed dataset is a unique data format of tensorflow.

```python
from data import make_data

# make tfrecords datasets
make_data(image_shape, image_dir, mask_dir, out_name, out_dir)

# get tfrecords datasets
dataset = get_tfrecord_data(
tf_record_path, tf_record_name, data_shape, batch_size=32, repeat=1, shuffle=True)
```

loss.py According to the characteristics of the cerebral hematoma dataset, in order to obtain higher segmentation accuracy. We use binary cross entropy with dice as the loss function of DR-Unet.
module.py This file contains several auxiliary functions for image processing.
utils.py This python file contains several auxiliary functions for file operations.
performance.py In order to evaluate the segmentation performance of the model, this file contains auxiliary functions for the calculation of several common segmentation indicators.
drunet.py This file contains the specific implementation of DR-UNet and three reduced dimensional residual convolution units (RDRCUs).

```python
from model import dr_unet

model = dr_unet.dr_unet(input_shape=(256, 256, 1))
model.summary()
```

segment.py This file shows how to train, test and verification DR-UNet on your own dataset. Including hematoma segmentation and hematoma volume estimation.

```python
import pathlib

# Parameter configuration
parser = argparse.ArgumentParser(description="Segment Use Args")
parser.add_argument('--model-name', default='DR_UNet', type=str)
parser.add_argument('--dims', default=32, type=int)
parser.add_argument('--epochs', default=50, type=int)
parser.add_argument('--batch-size', default=16, type=int)
parser.add_argument('--lr', default=2e-4, type=float)

# Training data, testing, verification parameter settings
parser.add_argument('--height', default=256, type=int)
parser.add_argument('--width', default=256, type=int)
parser.add_argument('--channel', default=1, type=int)
parser.add_argument('--pred-height', default=4 * 256, type=int)
parser.add_argument('--pred-width', default=4 * 256, type=int)
parser.add_argument('--total-samples', default=5000, type=int)
parser.add_argument('--invalid-samples', default=1000, type=int)
parser.add_argument('--regularize', default=False, type=bool)
parser.add_argument('--record-dir', default=r'', type=str, help='the save dir of tfrecord')
parser.add_argument('--train-record-name', type=str, default=r'train_data',
help='the train record save name')
parser.add_argument('--test-image-dir', default=r'', type=str,
help='the path of test images dir')
parser.add_argument('--invalid-record-name', type=str, default=r'test_data',
help='the invalid record save name')
parser.add_argument('--gt-mask-dir', default=r'', type=str,
help='the ground truth dir of validation set')
parser.add_argument('--invalid-volume-dir', default=r'', type=str,
help='estimation bleeding volume')
args = parser.parse_args()

segment = Segmentation(args)
# start training
segment.train()
# predict hematoma volume
segment.predict_blood_volume(input_dir, save_dir, calc_nums=-1, dpi=96, thickness=0.45)
```

train_segment.py If you want to train the segmentation model, you can run this file directly after filling in the data path.

```python
import segment

if name == 'main':
Seg = segment.Segmentation()
# start training
Seg.train()
```

predict_segment.py If you want to predict the segmentation result, you can run this file directly after filling in the ct images path.

```python
import segment

if name == 'main':
Seg = segment.Segmentation()
# start predict
input_dir = r'' # Fill in the image path
save_dir = r'' # fill in save path
Seg.predict_and_save(input_dir, save_dir)

```

predict_volume.py If you want to predict the complete hematoma volume of a patient, after filling in the path, then you can run this file,

```python
import segment

if name == 'main':
Seg = segment.Segmentation()
# start predict
input_dir = r'' # Fill in the image path
save_dir = r'' # fill in save path
Seg.predict_blood_volume(input_dir, save_dir, thickness=0.45)
```

test_performance.py If you want to understand the segmentation performance of the model, you need to fill in the relevant path first, and then run this file.

```python
import performance

if name == 'main':
# test model segmentation performance
pred_path = r'' # predict result path
gt_path = r'' # ground truth path
calc_performance(pred_path, gt_path, img_resize=(1400, 1400))
```

Requirements

Python 3.6, TensorFlow 2.1 and other common packages listed in requirements.txt.