Real-Time Polyp Detection, Localization and Segmentation in Colonoscopy Using Deep Learning

In plain words

During a colonoscopy a doctor looks inside the bowel for small growths called polyps, because some of them can turn into cancer, and some are easy to miss. This paper asks how well today's computer-vision models can find polyps in colonoscopy images. There are two versions of the task: draw a box around each polyp, or mark exactly which pixels belong to it.

The authors run a fair side-by-side test of many existing models on one public image collection, and they measure both accuracy and speed. They also propose a small new network, ColonSegNet, designed to be light enough to run in real time. The main message is that it is fast and reasonably accurate, while the most accurate models in the comparison are much slower.

TL;DR

• Benchmark of several detection and segmentation networks on the public Kvasir-SEG set of 1000 colonoscopy images, reporting accuracy and frames per second §III §V.
• Proposes ColonSegNet, a small residual encoder-decoder with squeeze-and-excitation blocks, used for segmentation §IV-C.
• Author-reported: ColonSegNet reaches Dice 0.8206 at 182.38 FPS on segmentation tab. V, and AP 0.8000 at 180 FPS for box detection tab. III.
• The public repository implements ColonSegNet for segmentation only; the detection numbers for it are not reproducible from this repository README.md.

Problem

Polyps can develop into colorectal cancer, and colonoscopy is the main way to find and remove them. Reviewing recorded video is labour-intensive, and polyps can be missed during the procedure §I. Automated detection, localisation and segmentation could help, but the paper argues that the methods in the literature were rarely compared on a common dataset and tasks, and that speed was rarely reported for clinical use §I §II. The stated aim is a benchmark with existing methods plus a fast baseline §I.

Method

The paper is a benchmark plus one proposed architecture. Section numbers below use IEEE journal numbering (Roman numerals for sections, letters for subsections), as produced by the IEEEtran class in the source; they were not checked against a rendered PDF.

Data and input. Kvasir-SEG has 1000 colonoscopy images with pixel masks and bounding boxes §III. The split is 880 images for training and 120 for validation §V-B. Every image is resized to 512 × 512 §V-B train.py. The resize does not preserve aspect ratio data.py.

Detection and localisation baselines. The paper uses EfficientDet-D0, Faster R-CNN, RetinaNet (ResNet50 and ResNet101), YOLOv3+spp and YOLOv4 §IV-A. Their hyperparameters are in tab. II. The detector code is not in the public repository, so these settings cannot be checked against code.

Segmentation baselines. The paper uses FCN8, UNet, UNet with ResNet34, ResUNet, ResUNet++, PSPNet, DeepLabv3+ (ResNet50 and ResNet101), HRNet and DoubleUNet §IV-B tab. IV. Per §V-B, UNet, ResUNet, ResUNet++, DoubleUNet and HRNet were implemented in Keras/TensorFlow, and the others in PyTorch. These baselines are not in the repository either.

ColonSegNet (proposed). An encoder-decoder with two encoder blocks and two decoder blocks §IV-C fig. 3. Channel widths, strides and skip connections are as follows:

1. Encoder 1 maps 3 channels to 64. It applies a residual block at full resolution, a 3×3 stride-2 convolution, a second residual block at half resolution, and a 2×2 max-pool. Outputs: a full-resolution skip, a half-resolution skip, and a quarter-resolution tensor §IV-C model.py model.py.
2. Encoder 2 maps 64 to 256 channels with the same layout, on the quarter-resolution tensor. Outputs: a quarter-resolution skip, an eighth-resolution skip, and a sixteenth-resolution tensor model.py model.py.
3. Decoder 1 upsamples the eighth-resolution skip by 4 with a transposed convolution to 256 × 256, concatenates the half-resolution skip, and applies a residual block with 128 channels §IV-C model.py model.py.
4. Decoder 1 upsamples the quarter-resolution skip by 2 to 256 × 256, concatenates the output of step 3, and applies a residual block with 128 channels model.py model.py.
5. Decoder 2 upsamples by 2 to 512 × 512, concatenates the full-resolution skip, and applies a residual block with 64 channels model.py model.py.
6. Decoder 2 upsamples the half-resolution skip by 2 to 512 × 512, concatenates the output of step 5, and applies a residual block with 32 channels model.py model.py.
7. A 1×1 convolution gives one output channel. A sigmoid turns it into a mask probability §IV-C model.py test.py.

Each residual block has two 3×3 convolutions with batch norm and ReLU. Its shortcut is a 1×1 convolution with batch norm, followed by a squeeze-and-excitation gate. The two branches are summed and passed through ReLU model.py. The paper describes the residual block only as "residual block with squeeze and excitation" §IV-C, so these internal details come from the code.

# ColonSegNet forward pass  [§IV-C] [fig. 3] [code: model.py#L112-L133]
# Input: image (B, 3, 512, 512); output: logits (B, 1, 512, 512)

def colonsegnet(image):
    # Encoder 1: 3 -> 64 channels
    # [code: model.py#L94] [code: model.py#L113]
    s11, s12, p1 = encoder_block(image, out_channels=64)
    # Encoder 2: 64 -> 256 channels
    # [code: model.py#L95] [code: model.py#L114]
    s21, s22, p2 = encoder_block(p1, out_channels=256)
    # Decoder 1: upsample x4 of the deepest skip to 256 x 256
    # [code: model.py#L98] [code: model.py#L116]
    up1 = conv_transpose(s22, 128, kernel=4, stride=4)
    # Concatenate the half-resolution skip, then residual block
    # [code: model.py#L117-L118]
    r1 = residual_block(concat(up1, s12), out_channels=128)
    # Decoder 1: upsample x2 of the quarter-resolution skip
    # [code: model.py#L100] [code: model.py#L120]
    up2 = conv_transpose(s21, 128, kernel=4, stride=2)
    # Concatenate the previous output first, then residual block
    # [code: model.py#L121-L122]
    r2 = residual_block(concat(r1, up2), out_channels=128)
    # Decoder 2: upsample x2 to full resolution
    # [code: model.py#L104] [code: model.py#L124]
    up3 = conv_transpose(r2, 64, kernel=4, stride=2)
    # Concatenate the full-resolution skip, then residual block
    # [code: model.py#L125-L126]
    r3 = residual_block(concat(up3, s11), out_channels=64)
    # Decoder 2: upsample x2 of the half-resolution skip
    # [code: model.py#L106] [code: model.py#L128]
    up4 = conv_transpose(s12, 32, kernel=4, stride=2)
    # Concatenate the previous output, then residual block
    # [code: model.py#L129-L130]
    r4 = residual_block(concat(r3, up4), out_channels=32)
    # 1x1 convolution to one mask channel; sigmoid is applied outside
    # [code: model.py#L110]
    return conv1x1(r4, out_channels=1)

# Encoder block: residual, stride-2 conv, residual, 2x2 max-pool
# [code: model.py#L73-L87]

def encoder_block(inputs, out_channels):
    # Residual block at the input resolution (skip 1)  [code: model.py#L83]
    full = residual_block(inputs, out_channels)
    # 3x3 stride-2 conv, batch norm, ReLU: halves resolution
    # [code: model.py#L63-L71] [code: model.py#L84]
    half = strided_conv_block(full, out_channels)
    # Residual block at half resolution (skip 2)  [code: model.py#L85]
    half = residual_block(half, out_channels)
    # 2x2 max-pool gives the next encoder's input  [code: model.py#L86-L87]
    pooled = max_pool(half, size=2)
    return full, half, pooled

# Residual block with squeeze-and-excitation on the shortcut
# [code: model.py#L26-L57]

def residual_block(inputs, out_channels):
    # Main branch: conv, batch norm, ReLU  [code: model.py#L30-L45]
    main = relu(batch_norm(conv3x3(inputs, out_channels)))
    # Main branch: conv, batch norm  [code: model.py#L33-L48]
    main = batch_norm(conv3x3(main, out_channels))
    # Shortcut: 1x1 conv and batch norm, then squeeze-and-excitation
    # [code: model.py#L36-L52]
    shortcut = batch_norm(conv1x1(inputs, out_channels))
    shortcut = squeeze_excite(shortcut, reduction=16)
    # Sum the branches, then ReLU  [code: model.py#L54-L55]
    return relu(main + shortcut)

# Squeeze-and-excitation gate  [code: model.py#L6-L24]

def squeeze_excite(x, reduction):
    # Channel count C  [code: model.py#L20]
    channels = x.shape[1]
    # Global average pool: one value per channel  [code: model.py#L21]
    pooled = global_avg_pool(x)
    # Bottleneck: C -> C / reduction -> C, with ReLU in between
    # [code: model.py#L11-L15]
    gate = linear(relu(linear(pooled, channels // reduction)), channels)
    # Sigmoid gates scale each channel  [code: model.py#L23-L24]
    return x * sigmoid(gate)

Data augmentation (as described). The paper says basic augmentations are used: horizontal and vertical flips, random rotation, random scale and random cropping. It also says every image is normalised by subtracting the mean and dividing by the standard deviation §IV-C. See Paper vs code for what the released code does.

Loss. tab. IV lists cross-entropy plus Dice for ColonSegNet. The code uses a sum of binary cross-entropy and soft Dice loss, applied to sigmoid outputs over the flattened batch, with smoothing constant 1:

# Loss: binary cross-entropy + (1 - soft Dice) on sigmoid outputs
# [tab. IV] [code: loss.py#L23-L41]

def dice_bce_loss(logits, targets):
    # Sigmoid maps logits to probabilities  [code: loss.py#L30]
    probs = sigmoid(logits)
    # Flatten batch and spatial dimensions  [code: loss.py#L32-L34]
    p = flatten(probs)
    t = flatten(targets)
    # Soft Dice with smoothing 1  [code: loss.py#L36-L37]
    dice = (2 * sum(p * t) + 1) / (sum(p) + sum(t) + 1)
    # Binary cross-entropy, averaged over pixels  [code: loss.py#L38]
    bce = mean(-(t * log(p) + (1 - t) * log(1 - p)))
    return bce + (1 - dice)

Evaluation. Detection uses IoU and average precision (AP). AP is averaged over IoU thresholds 0.25 to 0.75 in steps of 0.05 §V-A eq. 1 eq. 2. Segmentation uses Jaccard, Dice, F2, precision, recall, accuracy and FPS §V-A eq. 3–eq. 9.

\text{IoU}(A,B) = \frac{A \cap B}{A \cup B} eq. 1

\mathrm{AP} = \sum_n \left\{ (r_{n+1} - r_n)\, p_{\mathrm{interp}}(r_{n+1}) \right\} eq. 2

with p_{\mathrm{interp}}(r_{n+1}) = \max_{\tilde{r} \ge r_{n+1}} p(\tilde{r}).

\text{DSC} = \frac{2\,tp}{2\,tp + fp + fn} eq. 3

The paper also gives IoU eq. 4, recall eq. 5, precision eq. 6, F2 eq. 7, accuracy eq. 8 and FPS eq. 9 in terms of pixel counts.

Hyperparameters

Values from the paper, with the released code where it has them. Anything marked "not stated" is absent from the paper. The baseline tables are reproduced in full below, because the paper gives no other source for them.

tab. II Detection baselines (Kvasir-SEG)

tab. IV Segmentation baselines and ColonSegNet (Kvasir-SEG)

Results

Author-reported. Setup: Kvasir-SEG, 880 training and 120 validation images, inputs at 512 × 512 §V-B. The paper's figure captions call the evaluation set "test dataset", while §V-B calls the 120 images "validation". The code uses the same 120 images for checkpoint selection and for testing (see Paper vs code). Speed is measured on an NVIDIA GTX 2080Ti §V-B.

Headline numbers. ColonSegNet reaches AP 0.8000 and IoU 0.8100 at 180 FPS for detection tab. III, and Dice 0.8206 at 182.38 FPS for segmentation tab. V. The paper's best Dice is UNet with ResNet34 at 0.8757, at 35.00 FPS tab. V. Its best detection AP is RetinaNet with ResNet101 at 0.8745, at 16.80 FPS tab. III. The authors name YOLOv4 as the best detector, because of its trade-off between AP, IoU and speed §VI.

tab. III Detection and localisation on Kvasir-SEG

tab. V Segmentation on Kvasir-SEG

The text says ColonSegNet is "more than four times" faster than UNet with ResNet34 §V-C; the table ratio is 182.38 / 35.00 ≈ 5.2. The text also says ColonSegNet is "around 0.75 times" the size of UNet §V-C; tab. IV gives 5,014,049 / 7,858,433 ≈ 0.64. See Pitfalls.

Qualitative results. Detection examples and best/worst segmentation cases are shown in fig. 4 and fig. 5. The authors report that flat or very small polyps, saturated regions and strongly inclined views were hard for the segmenters, and that polyps whose colour resembles surrounding mucosa were missed §V-D §VI.

Paper vs code

Official code: https://github.com/DebeshJha/ColonSegNet, commit d8f782f971884dcb93aa5e86940b409d13b67c7a. The paper does not link it. I used it because the GitHub owner is the first author, Debesh Jha. The GitHub search also returned two third-party repositories (HJ-LIM-DS, SeongminCC), which were not used. The README links a TensorFlow port by a third party, nikhilroxtomar, which is not the official code. The repository has no LICENSE file, so code_license is recorded as "none". No code is quoted above; only pseudocode with citations is used.

The repository implements the segmentation network and its training and test scripts. It does not implement the detection baselines or the other segmentation baselines.

Differences between paper and code:

• Batch size. tab. IV lists 8 for ColonSegNet. The code trains with batch size 1 train.py.
• Normalisation. §IV-C says images are normalised with mean subtraction and standard deviation division. The code divides pixel values by 255 only data.py test.py.
• Augmentation. §IV-C lists flips, rotation, scale and cropping as used. The released training reads the non-augmented folder, and the augmented-data lines are commented out train.py. The augmentation script calls its routine with augment=False data_aug.py. When enabled, it produces 12 variants per image, including elastic, grid and optical distortion, dropout and noise data_aug.py. Random scale is absent from the code.
• Validation and test. The paper calls the 120 held-out images "validation" §V-B. The code selects the checkpoint with the lowest loss on val.txt train.py. The test script then reports metrics on the same val.txt images test.py test.py. No separate test set is used.
• Post-processing. The paper does not mention a CRF. The test script applies a dense CRF to the saved masks test.py crf.py, with unary confidence 0.7, spatial sigma 3, compatibility 3 and 10 inference steps crf.py. Metrics are computed before the CRF test.py test.py.
• Binary threshold. Probabilities are thresholded at 0.5 for the metrics test.py test.py test.py. The paper does not state a threshold.
• F2 column. Equation 7 defines F2 with β = 2 eq. 7. The test script computes its "F2" with fbeta_score(..., beta=1.0), which is the same as the F1 / Dice score test.py test.py. For ColonSegNet, tab. V lists F2 = DSC = 0.8206, which fits beta = 1.
• Learning-rate schedule and epochs. The paper does not state these. The code runs 100 epochs train.py and reduces the learning rate when validation loss plateaus for 5 epochs train.py. The factor is not passed, so the library default applies.
• Loss. tab. IV says cross-entropy plus Dice. The code uses BCE on sigmoid probabilities plus (1 − Dice), with Dice computed over the whole batch loss.py.
• Detection boxes. The paper states that Kvasir-SEG provides bounding boxes §III, but does not say how ColonSegNet's boxes are produced. The README says ground-truth masks were converted to boxes, using a converter in another repository README.md. The code has no such step.
• Checkpoint. The repository includes checkpoint.pth. I did not load it, so I cannot say which epoch it corresponds to, or whether it reproduces the table.
• Parameter count. tab. IV gives 5,014,049 for ColonSegNet. I did not instantiate the model, so this is unverified.

Where the code and the paper agree: 512 × 512 input, Adam with learning rate 1e-4, the layer layout described above and the encoder/decoder widths model.py, and the 880/120 split files.

Pitfalls

• Metric leakage. The checkpoint is chosen on the 120 images that are then used for the reported metrics train.py test.py. Treat the numbers as validation results rather than held-out test results.
• Single run, single split. The paper reports one number per metric. It gives no variance over seeds or splits §V. The code sets seed 42 train.py utils.py.
• Metrics are averaged per image. Each image's score is computed, then the scores are averaged test.py test.py. This is not the same as pooling pixels over the dataset. The paper does not say which it uses.
• Evaluation at 512 × 512. Ground-truth masks are resized to 512 × 512 and metrics are computed at that size test.py. Scores at original resolution will differ, and the aspect ratio is not kept data.py.
• FPS measurement. Timing wraps the model forward and sigmoid, without a CUDA synchronisation call before reading the clock test.py. GPU work may not be finished when the timer stops, so the FPS may be overstated. The mean includes the first image, with no warm-up test.py. The paper's FPS is measured on a GTX 2080Ti §V-B.
• Mixed frameworks. Some baselines are Keras/TensorFlow and others PyTorch, and training hardware differs §V-B. FPS comparisons across rows therefore mix framework and hardware effects.
• Detection numbers. The ColonSegNet detection row in tab. III depends on a mask-to-box conversion that the paper does not describe. It cannot be reproduced from the repository.
• Residual block details. The paper does not give the residual block's internal layout. The code applies batch norm after each convolution, puts the squeeze-and-excitation gate on the 1×1 shortcut only, and uses reduction 16 model.py model.py. The SE reduction is not stated in the paper.
• Upsampling details. The paper says "stride 4" and "stride 2" for the decoder transposed convolutions §IV-C. The code also fixes which skip feeds each one, and in what order the concatenation happens model.py. Implement it as the code does, or the shapes will not match.
• Numbers that conflict within the paper.
  • The text gives 1000 images §III and also "1072 images" in total §III. The size bins it lists (700 + 323 + 48) sum to 1071.
  • ColonSegNet is "around 0.75 times" the size of UNet §V-C. tab. IV gives about 0.64.
  • The text says ColonSegNet is "nearly seven times faster than DeepLabv3 with ResNet34" §V-C. tab. V has no ResNet34 DeepLabv3+ row; the nearest, ResNet50, runs at 27.90 FPS, a ratio of about 6.5.
  • The text says YOLOv3's IoU is "nearly 8%" above RetinaNet §V-C. tab. III gives 0.8248 against 0.7313, a difference of about 9 points.
• Clinical claims. The paper calls the work a retrospective study and says clinical validation is still needed §VI. Speed numbers are offline frame rates on a single GPU, not measured in a clinical setting.

Lineage

• Builds on: residual blocks §IV-C (He et al.), squeeze-and-excitation §IV-C (Hu et al.), and the encoder-decoder idea from U-Net §IV-B. The dataset is Kvasir-SEG §III.
• Compared against: detectors EfficientDet-D0, Faster R-CNN, RetinaNet, YOLOv3 and YOLOv4 tab. III; segmenters FCN8, UNet, UNet-ResNet34, ResUNet, ResUNet++, PSPNet, DeepLabv3+, HRNet and DoubleUNet tab. V.
• Superseded by: not stated.