16 research outputs found
Importance of Aligning Training Strategy with Evaluation for Diffusion Models in 3D Multiclass Segmentation
Recently, denoising diffusion probabilistic models (DDPM) have been applied
to image segmentation by generating segmentation masks conditioned on images,
while the applications were mainly limited to 2D networks without exploiting
potential benefits from the 3D formulation. In this work, we studied the
DDPM-based segmentation model for 3D multiclass segmentation on two large
multiclass data sets (prostate MR and abdominal CT). We observed that the
difference between training and test methods led to inferior performance for
existing DDPM methods. To mitigate the inconsistency, we proposed a recycling
method which generated corrupted masks based on the model's prediction at a
previous time step instead of using ground truth. The proposed method achieved
statistically significantly improved performance compared to existing DDPMs,
independent of a number of other techniques for reducing train-test
discrepancy, including performing mask prediction, using Dice loss, and
reducing the number of diffusion time steps during training. The performance of
diffusion models was also competitive and visually similar to
non-diffusion-based U-net, within the same compute budget. The JAX-based
diffusion framework has been released at
https://github.com/mathpluscode/ImgX-DiffSeg.Comment: Accepted at Deep Generative Models workshop at MICCAI 202
A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models
Denoising diffusion models have found applications in image segmentation by
generating segmented masks conditioned on images. Existing studies
predominantly focus on adjusting model architecture or improving inference,
such as test-time sampling strategies. In this work, we focus on improving the
training strategy and propose a novel recycling method. During each training
step, a segmentation mask is first predicted given an image and a random noise.
This predicted mask, which replaces the conventional ground truth mask, is used
for denoising task during training. This approach can be interpreted as
aligning the training strategy with inference by eliminating the dependence on
ground truth masks for generating noisy samples. Our proposed method
significantly outperforms standard diffusion training, self-conditioning, and
existing recycling strategies across multiple medical imaging data sets: muscle
ultrasound, abdominal CT, prostate MR, and brain MR. This holds for two widely
adopted sampling strategies: denoising diffusion probabilistic model and
denoising diffusion implicit model. Importantly, existing diffusion models
often display a declining or unstable performance during inference, whereas our
novel recycling consistently enhances or maintains performance. We show that,
under a fair comparison with the same network architectures and computing
budget, the proposed recycling-based diffusion models achieved on-par
performance with non-diffusion-based supervised training. By ensembling the
proposed diffusion and the non-diffusion models, significant improvements to
the non-diffusion models have been observed across all applications,
demonstrating the value of this novel training method. This paper summarizes
these quantitative results and discusses their values, with a fully
reproducible JAX-based implementation, released at
https://github.com/mathpluscode/ImgX-DiffSeg.Comment: Accepted for publication at the Journal of Machine Learning for
Biomedical Imaging (MELBA) https://melba-journal.org/2023:01
Strategising template-guided needle placement for MR-targeted prostate biopsy
Clinically significant prostate cancer has a better chance to be sampled
during ultrasound-guided biopsy procedures, if suspected lesions found in
pre-operative magnetic resonance (MR) images are used as targets. However, the
diagnostic accuracy of the biopsy procedure is limited by the
operator-dependent skills and experience in sampling the targets, a sequential
decision making process that involves navigating an ultrasound probe and
placing a series of sampling needles for potentially multiple targets. This
work aims to learn a reinforcement learning (RL) policy that optimises the
actions of continuous positioning of 2D ultrasound views and biopsy needles
with respect to a guiding template, such that the MR targets can be sampled
efficiently and sufficiently. We first formulate the task as a Markov decision
process (MDP) and construct an environment that allows the targeting actions to
be performed virtually for individual patients, based on their anatomy and
lesions derived from MR images. A patient-specific policy can thus be
optimised, before each biopsy procedure, by rewarding positive sampling in the
MDP environment. Experiment results from fifty four prostate cancer patients
show that the proposed RL-learned policies obtained a mean hit rate of 93% and
an average cancer core length of 11 mm, which compared favourably to two
alternative baseline strategies designed by humans, without hand-engineered
rewards that directly maximise these clinically relevant metrics. Perhaps more
interestingly, it is found that the RL agents learned strategies that were
adaptive to the lesion size, where spread of the needles was prioritised for
smaller lesions. Such a strategy has not been previously reported or commonly
adopted in clinical practice, but led to an overall superior targeting
performance when compared with intuitively designed strategies.Comment: Paper submitted and accepted to CaPTion (Cancer Prevention through
early detecTion) @ MICCAI 2022 Worksho
Boundary-RL: Reinforcement Learning for Weakly-Supervised Prostate Segmentation in TRUS Images
We propose Boundary-RL, a novel weakly supervised segmentation method that
utilises only patch-level labels for training. We envision the segmentation as
a boundary detection problem, rather than a pixel-level classification as in
previous works. This outlook on segmentation may allow for boundary delineation
under challenging scenarios such as where noise artefacts may be present within
the region-of-interest (ROI) boundaries, where traditional pixel-level
classification-based weakly supervised methods may not be able to effectively
segment the ROI. Particularly of interest, ultrasound images, where intensity
values represent acoustic impedance differences between boundaries, may also
benefit from the boundary delineation approach. Our method uses reinforcement
learning to train a controller function to localise boundaries of ROIs using a
reward derived from a pre-trained boundary-presence classifier. The classifier
indicates when an object boundary is encountered within a patch, as the
controller modifies the patch location in a sequential Markov decision process.
The classifier itself is trained using only binary patch-level labels of object
presence, which are the only labels used during training of the entire boundary
delineation framework, and serves as a weak signal to inform the boundary
delineation. The use of a controller function ensures that a sliding window
over the entire image is not necessary. It also prevents possible
false-positive or -negative cases by minimising number of patches passed to the
boundary-presence classifier. We evaluate our proposed approach for a
clinically relevant task of prostate gland segmentation on trans-rectal
ultrasound images. We show improved performance compared to other tested weakly
supervised methods, using the same labels e.g., multiple instance learning.Comment: Accepted to MICCAI Workshop MLMI 2023 (14th International Conference
on Machine Learning in Medical Imaging
Deep hashing for global registration of untracked 2D laparoscopic ultrasound to CT
PURPOSE: The registration of Laparoscopic Ultrasound (LUS) to CT can enhance the safety of laparoscopic liver surgery by providing the surgeon with awareness on the relative positioning between critical vessels and a tumour. In an effort to provide a translatable solution for this poorly constrained problem, Content-based Image Retrieval (CBIR) based on vessel information has been suggested as a method for obtaining a global coarse registration without using tracking information. However, the performance of these frameworks is limited by the use of non-generalisable handcrafted vessel features. METHODS: We propose the use of a Deep Hashing (DH) network to directly convert vessel images from both LUS and CT into fixed size hash codes. During training, these codes are learnt from a patient-specific CT scan by supplying the network with triplets of vessel images which include both a registered and a mis-registered pair. Once hash codes have been learnt, they can be used to perform registration with CBIR methods. RESULTS: We test a CBIR pipeline on 11 sequences of untracked LUS distributed across 5 clinical cases. Compared to a handcrafted feature approach, our model improves the registration success rate significantly from 48% to 61%, considering a 20 mm error as the threshold for a successful coarse registration. CONCLUSIONS: We present the first DH framework for interventional multi-modal registration tasks. The presented approach is easily generalisable to other registration problems, does not require annotated data for training, and may promote the translation of these techniques
Image quality assessment for machine learning tasks using meta-reinforcement learning
In this paper, we consider image quality assessment (IQA) as a measure of how images are amenable with respect to a given downstream task, or task amenability. When the task is performed using machine learning algorithms, such as a neural-network-based task predictor for image classification or segmentation, the performance of the task predictor provides an objective estimate of task amenability. In this work, we use an IQA controller to predict the task amenability which, itself being parameterised by neural networks, can be trained simultaneously with the task predictor. We further develop a meta-reinforcement learning framework to improve the adaptability for both IQA controllers and task predictors, such that they can be fine-tuned efficiently on new datasets or meta-tasks. We demonstrate the efficacy of the proposed task-specific, adaptable IQA approach, using two clinical applications for ultrasound-guided prostate intervention and pneumonia detection on X-ray images
Image quality assessment by overlapping task-specific and task-agnostic measures: application to prostate multiparametric MR images for cancer segmentation
Image quality assessment (IQA) in medical imaging can be used to ensure that
downstream clinical tasks can be reliably performed. Quantifying the impact of
an image on the specific target tasks, also named as task amenability, is
needed. A task-specific IQA has recently been proposed to learn an
image-amenability-predicting controller simultaneously with a target task
predictor. This allows for the trained IQA controller to measure the impact an
image has on the target task performance, when this task is performed using the
predictor, e.g. segmentation and classification neural networks in modern
clinical applications. In this work, we propose an extension to this
task-specific IQA approach, by adding a task-agnostic IQA based on
auto-encoding as the target task. Analysing the intersection between
low-quality images, deemed by both the task-specific and task-agnostic IQA, may
help to differentiate the underpinning factors that caused the poor target task
performance. For example, common imaging artefacts may not adversely affect the
target task, which would lead to a low task-agnostic quality and a high
task-specific quality, whilst individual cases considered clinically
challenging, which can not be improved by better imaging equipment or
protocols, is likely to result in a high task-agnostic quality but a low
task-specific quality. We first describe a flexible reward shaping strategy
which allows for the adjustment of weighting between task-agnostic and
task-specific quality scoring. Furthermore, we evaluate the proposed algorithm
using a clinically challenging target task of prostate tumour segmentation on
multiparametric magnetic resonance (mpMR) images, from 850 patients. The
proposed reward shaping strategy, with appropriately weighted task-specific and
task-agnostic qualities, successfully identified samples that need
re-acquisition due to defected imaging process.Comment: Accepted for publication at the Journal of Machine Learning for
Biomedical Imaging (MELBA) https://www.melba-journal.or
Semi-weakly-supervised neural network training for medical image registration
For training registration networks, weak supervision from segmented corresponding regions-of-interest (ROIs) have been proven effective for (a) supplementing unsupervised methods, and (b) being used independently in registration tasks in which unsupervised losses are unavailable or ineffective. This correspondence-informing supervision entails cost in annotation that requires significant specialised effort. This paper describes a semi-weakly-supervised registration pipeline that improves the model performance, when only a small corresponding-ROI-labelled dataset is available, by exploiting unlabelled image pairs. We examine two types of augmentation methods by perturbation on network weights and image resampling, such that consistency-based unsupervised losses can be applied on unlabelled data. The novel WarpDDF and RegCut approaches are proposed to allow commutative perturbation between an image pair and the predicted spatial transformation (i.e. respective input and output of registration networks), distinct from existing perturbation methods for classification or segmentation. Experiments using 589 male pelvic MR images, labelled with eight anatomical ROIs, show the improvement in registration performance and the ablated contributions from the individual strategies. Furthermore, this study attempts to construct one of the first computational atlases for pelvic structures, enabled by registering inter-subject MRs, and quantifies the significant differences due to the proposed semi-weak supervision with a discussion on the potential clinical use of example atlas-derived statistics
Semi-weakly-supervised neural network training for medical image registration
For training registration networks, weak supervision from segmented
corresponding regions-of-interest (ROIs) have been proven effective for (a)
supplementing unsupervised methods, and (b) being used independently in
registration tasks in which unsupervised losses are unavailable or ineffective.
This correspondence-informing supervision entails cost in annotation that
requires significant specialised effort. This paper describes a
semi-weakly-supervised registration pipeline that improves the model
performance, when only a small corresponding-ROI-labelled dataset is available,
by exploiting unlabelled image pairs. We examine two types of augmentation
methods by perturbation on network weights and image resampling, such that
consistency-based unsupervised losses can be applied on unlabelled data. The
novel WarpDDF and RegCut approaches are proposed to allow commutative
perturbation between an image pair and the predicted spatial transformation
(i.e. respective input and output of registration networks), distinct from
existing perturbation methods for classification or segmentation. Experiments
using 589 male pelvic MR images, labelled with eight anatomical ROIs, show the
improvement in registration performance and the ablated contributions from the
individual strategies. Furthermore, this study attempts to construct one of the
first computational atlases for pelvic structures, enabled by registering
inter-subject MRs, and quantifies the significant differences due to the
proposed semi-weak supervision with a discussion on the potential clinical use
of example atlas-derived statistics