Automatic Instrument Segmentation in Robot-Assisted Surgery Using Deep Learning

In plain words

A robot-assisted surgery system shows the surgeon a video of the operation. Software that can find the tools in each frame is useful for tracking them, and this paper trains a neural network to do that. The network colours every pixel as "tool" or "background". It can also tell the parts of a tool apart (shaft, wrist, grippers) and the tool types apart.

The authors start from a well-known image-to-mask design, U-Net, and add two variants with pretrained encoders. They train with a loss that rewards both correct pixels and good overlap with the true tool shape. The two-class task works well. The many-class task is much harder, because some tool types appear only a few times in the training data.

TL;DR

• The task is pixel-wise segmentation of surgical instruments in video frames, at three levels: binary (tool or background), 3 tool parts, and 7 instrument types §2.1.
• The models are U-Net, TernausNet-11 and TernausNet-16 (VGG encoders), and LinkNet-34 (ResNet-34 encoder). They are trained with cross-entropy plus a soft-Jaccard term §2.2 §2.3.
• Author-reported best results are binary IoU 83.60% and Dice 90.01% (TernausNet-16), parts IoU 65.50% (TernausNet-16), and instrument-type IoU 34.61% (TernausNet-11) tab. 1.
• The official MIT-licensed repository differs from the paper in the loss weights, the binary threshold and how metrics are aggregated. Read "Paper vs code" before implementing.
• The challenge training data is not in the repository, and the paper does not state a license for it.

Problem

Robot-assisted surgery needs accurate pixel masks of instruments so that tracking and pose estimation can work §1. Shadows, specular reflections, blood, lens fogging and moving tissue make the task hard §1. Earlier deep-learning work reported competitive binary results and promising multi-class results §1. The paper uses the MICCAI 2017 Endoscopic Vision SubChallenge data, where training data is 8 stereo sequences from a da Vinci Xi system §2.1. The abstract calls the authors' submission a winning solution, while §1 says it reached one of the top results [abstract] §1.

Method

Data §2.1

1. Training sequences: 8 sequences of 225 frames, sampled at 2 Hz, from stereo video at 1920×1080 RGB. Only left-camera frames are used, because only they have labels §2.1. The README says 255 frames per sequence [README.rst, Data section].
2. Crop: the black border is removed by cropping a 1280×1024 window whose top-left corner is at x = 320, y = 28 §2.1 prepare_data.py.
3. Labels are hand-drawn masks of the shaft, wrist and claspers (parts), plus instrument type labels. Ground-truth values are 10, 20, 30, 40 and 0 §2.1. The code maps values 10, 20 and 30 to parts 1–3 and ignores 40 prepare_data.py.
4. Instrument types: §2.1 lists four categories (left and right prograsp forceps, monopolar curved scissors, large needle driver, other). §3 uses 7 classes. The code maps six named types plus "Other" to classes 1–7 prepare_data.py.
5. Test data: 8 sequences of 75 frames after each training sequence, plus 2 full 300-frame sequences. When a 75-frame sequence is evaluated, the matching training sequence must be excluded §2.1. The repository does not contain the test data.

Network architectures §2.2

6. U-Net: a contracting path, an expanding path, and skip connections from each encoder level to the decoder §2.2. The paper says the network is "slightly modified" but does not give the modifications. The code's vanilla UNet uses 32 base filters and five levels, with 4×4 pooling at the bottom models.py. The code trains it from scratch, with no pretrained option train.py.
7. TernausNet-11 and TernausNet-16: a U-Net whose encoder is a pretrained VGG11 or VGG16 network §2.2. The code uses torchvision's pretrained=True weights train.py models.py models.py.
8. LinkNet-34: a ResNet-34 encoder with a 7×7 stride-2 first convolution and max-pooling §2.2. Each decoder block has a 1×1 convolution that reduces the channel count by 4, batch normalisation and a transposed convolution §2.2 models.py. The encoder output is added to the decoder output, rather than concatenated models.py.
9. Output: one logit channel for the binary task, and a log-softmax over classes for the multi-class task models.py.

Loss and metric §2.3

10. The Jaccard index is the overlap measure. For sets A and B it is the intersection over the union eq. 1:

J(A, B) = \frac{|A\cap B|}{|A\cup B|} = \frac{|A\cap B|}{|A|+|B|-|A\cap B|} eq. 1

11. For pixels, the paper adapts this to a per-pixel expression, where y_i is the binary label and ŷ_i the predicted probability of pixel i eq. 2:

J=\frac{1}{n}\sum_{i=1}^n\frac{y_i\hat{y}_i}{y_i+\hat{y}_i-y_i\hat{y}_i} eq. 2

12. The paper adds a classification loss H: binary cross-entropy for the binary task, categorical cross-entropy for the multi-class task §2.3. The total loss is H − log J eq. 3:

L = H - \log J eq. 3

13. The paper says that minimising this loss raises both the probabilities of correct pixels and the overlap J §2.3. The equation gives both terms weight 1.

Output and thresholding §2.3

14. The network returns one probability per pixel (binary) or per class (multi-class). The output has the input's size §2.3.
15. For the binary task, a threshold of 0.3 is applied. The paper says it was chosen on the validation set. Pixels above it become 255 and the rest become 0 §2.3.
16. For the multi-class task, the paper says it follows a similar procedure, with an integer code per class §2.3. The code takes the argmax over classes generate_masks.py.

# Binary loss: cross-entropy plus soft-Jaccard  [eq. 3] [§2.3]
# Input:  logits (B, 1, H, W); mask (B, 1, H, W), values in {0, 1}
# Output: scalar loss

def binary_loss(logits, mask, jaccard_weight):
    # Pixel-wise binary cross-entropy H  [§2.3] [code: loss.py#L18]
    cross_entropy = binary_cross_entropy(logits, mask)
    # Predicted probabilities for the overlap term  [code: loss.py#L23]
    probs = sigmoid(logits)
    # Soft intersection and union over the whole batch
    # [eq. 2] [code: loss.py#L25-L26]
    intersection = sum(probs * mask)
    union = sum(probs) + sum(mask)
    # Small constant against division by zero  [code: loss.py#L21]
    eps = 1e-15
    jaccard = (intersection + eps) / (union - intersection + eps)
    # Paper: H - log J, weight 1 each  [eq. 3]
    # Code: (1 - a) * H - a * log J, a = jaccard_weight
    # [code: loss.py#L18]
    return (1 - jaccard_weight) * cross_entropy - jaccard_weight * log(jaccard)

# Binary inference: probability map to mask
# [§2.3] [code: generate_masks.py#L86]
# Input:  logits (H, W) for one tile
# Output: mask (H, W) with values 0 or 255

def binary_mask(logits, threshold):
    # Per-pixel probability that the pixel belongs to an instrument
    probs = sigmoid(logits)
    # Paper: threshold 0.3, chosen on validation data  [§2.3]
    # Code: evaluation thresholds at 0.5  [code: evaluate.py#L77]
    return where(probs > threshold, 255, 0)

Results §3

17. The paper compares U-Net, TernausNet-11, TernausNet-16 and LinkNet-34 on the three tasks, and reports IoU and Dice in percent and inference time in milliseconds on a single GTX 1080 Ti §3 tab. 1.
18. The paper attributes the lower multi-class scores to the small training set, and says several classes appear only a few times in it §3. It suggests more data would help §3.
19. Visual comparisons are in fig. 3. Example frames with binary, parts and instrument-type masks are in fig. 1.
20. The paper says the binary and parts results are the best in the literature so far, citing Pakhomov et al. (arXiv:1703.08580) and Garcia et al. (arXiv:1706.08126) §3.

Hyperparameters

Results

Author-reported. Dataset: MICCAI 2017 Endoscopic Vision SubChallenge robotic instrument training data, described in §2.1. The paper does not say which split produced the numbers. Hardware: one GTX 1080 Ti for timing §3. Baselines: the other three networks in this paper, and the literature values cited in §3.

Full table from the paper tab. 1. IoU and Dice in %, time in ms. Bold marks the best per column, as in the paper.

tab. 1

The paper says LinkNet-34 is the fastest and "more than twice as fast as TernausNet" on the binary task §3. Against TernausNet-11 (142 ms) that is about 1.6×; against TernausNet-16 (184 ms) it is about 2.1× tab. 1.

Paper vs code

Code: ternaus/robot-surgery-segmentation at commit 1d18157, MIT. The owner is an author (Vladimir Iglovikov), and the LICENSE copyright lists all four authors LICENSE README.rst, Team members. The abstract links it. No other repository was searched for.

• Loss weights: paper uses H − log J with weight 1 each eq. 3. Code uses (1 − a)·H − a·log J, with a = 0.5 by default and a = 0.3 in the script loss.py loss.py train.py train.sh. With a = 0.5 the code is the paper's loss scaled by 0.5. With a = 0.3 the weights are 0.7 and 0.3.
• Jaccard term: paper's eq. 2 averages per-pixel terms. Code takes one ratio of sums over the whole batch, computed on sigmoid probabilities loss.py. Its validation metric is computed per image and averaged validation.py.
• Binary threshold: paper uses 0.3, chosen on validation §2.3. Code validation uses logit > 0 (probability 0.5) validation.py. Evaluation uses a 0.5 threshold on the written PNG evaluate.py. No code selects a threshold.
• Multi-class loss: code applies the Jaccard term to every class, including background, and sums the terms loss.py.
• Optimiser: paper does not state one. Code uses Adam train.py.
• VGG11 encoder: paper says VGG11 has seven convolutional layers §2.2. The code uses eight convolutions, at indices 0, 3, 6, 8, 11, 13, 16 and 18 of torchvision's VGG11 feature stack models.py.
• LinkNet decoder: paper describes a 1×1 reduction, batch normalisation and a transposed convolution §2.2. Code adds a 1×1 expansion back to the output channel count, with batch normalisation and ReLU after each step models.py.
• U-Net-family decoders: paper describes "upsampling followed by a convolution" §2.2. Code's default decoder blocks use a 4×4, stride-2 transposed convolution models.py models.py.
• Frame count: paper says 225 frames per training sequence §2.1. README says 255 [README.rst, Data section].
• Evaluation protocol: paper does not say which frames the table uses. The code predicts each fold's held-out instrument datasets, stitches the predictions back to full frames, and scores them against the training ground truth generate_masks.py evaluate.py. That is out-of-fold prediction on training data.
• Training recipe: README gives an older script (batch 16, Jaccard weight 1, 4 GPUs) that differs from train.sh (batch 6, weight 0.3, 2 GPUs) [README.rst, How to run] train.sh.
• Multi-class training: the repository's scripts cover only the binary task. train.py accepts parts and instruments train.py, but no learning rate, epoch count or batch size is given for those tasks.
• Label 40: the paper lists label 40 §2.1. The code does not map it to a part prepare_data.py.
• Timing: the paper reports inference times on a GTX 1080 Ti §3. The repository has no timing script.

Pitfalls

• Resumed training. The script's second run loads the checkpoint saved by the first run (epoch + 1, so it resumes at epoch 11) and runs epochs 11–20 at the new learning rate utils.py utils.py. Only the model weights are restored; the Adam state is rebuilt from scratch train.py. The paper does not describe this.
• Last epoch, not best epoch. The saved model is the last one, not the one with the best validation score utils.py. The paper does not say which model it reports.
• Threshold chosen on validation. The 0.3 threshold is described as chosen on validation data §2.3, but no code does this. Re-tune the threshold on your own validation split.
• Metric aggregation differs by script. Validation for multi-class pools a confusion matrix over the whole set and excludes background validation.py. evaluate.py averages each class per image over the classes present in the ground truth, then averages over images evaluate.py. For binary, the README says the score is averaged per video, but evaluate.py averages over all files [README.rst, Evaluation] evaluate.py. These can give different numbers on the same predictions.
• Batch-norm in prediction (code reading, not run). On a CUDA machine, generate_masks.py returns the model before calling model.eval() generate_masks.py. Models with batch normalisation would then run in training mode during prediction: LinkNet-34 models.py and AlbuNet, which has a ResNet-34 encoder models.py. Call eval() yourself before predicting.
• "Other" instruments. The binary and parts masks exclude instruments labelled "Other" prepare_data.py prepare_data.py. The instrument-type mask keeps them as class 7 prepare_data.py. The paper does not describe this rule.
• Crop does nothing at train time. The training crop equals the cropped image size, so random crop and padding do nothing. The only augmentation is the two flips train.py prepare_data.py.
• Unstated split. The paper does not say whether the table averages over the four folds or reports one. The code trains one model per fold train.py utils.py.
• Data availability. The training data and test data are the challenge's, and the repository does not include them. Check the data's terms before any use.
• Pretraining source. The paper says the encoders are "pre-trained" but does not name the dataset §2.2. The code relies on torchvision's pretrained=True train.py and does not name a dataset. Check which weights your torchvision version downloads.
• Instrument categories. §2.1 lists four categories, but §3 and the code use seven classes §2.1 prepare_data.py. Use the code's classes.
• Speed claim. The "more than twice as fast" claim holds only against TernausNet-16 tab. 1.

Lineage

• Builds on: TernausNet, arXiv:1801.05746 §2.2; the U-Net, arXiv:1505.04597 §2.2; VGG, arXiv:1409.1556 §2.2; ResNet, arXiv:1512.03385 §2.2; LinkNet, arXiv:1707.03718 §2.2. Only the TernausNet card exists.
• Compared against: arXiv:1703.08580 and arXiv:1706.08126, which the paper cites for earlier binary and multi-class values §3 README.rst, Results.
• Superseded by: not known from content already in this wiki.