Cut out the annotator, keep the cutout: better segmentation with weak supervision

1 INTRODUCTION . Automated image segmentation has seen rapid improvements with recent developments in deep learning ( Li et al. , 2018 ; Chen et al. , 2017 ; Milletari et al. , 2016 ) . Convolutional neural networks ( CNNs ) achieve high segmentation performance—but can require large , labeled training datasets . Acquiring training labels is laborious and slow , particularly for medical images where expert segmentation is often required in 3- or 4-dimensions . While large datasets and pre-trained networks exist for natural images , medical image segmentation is a more targeted task , typically requiring new training sets for every imaging modality , scanner type , anatomical structure , and patient population . Such difficulties abound , but the significant impact of improved medical image segmentation motivates tackling these challenges ( Hesamian et al. , 2019 ) . Many few-shot learning ( FSL ) approaches have been proposed to mitigate these difficulties by training networks using only a few labeled examples . For example , data augmentation can reduce the needed amount of labeled data by introducing additional variation into small , labeled training sets through operations such as affine transforms or learned transformations ( e.g . deformations learned by GANs ) ( Zhao et al. , 2019 ; Eaton-Rosen et al. , 2018 ; Çiçek et al. , 2016 ) . Many semisupervised approaches aim to learn a useful representation from unlabeled data then fine-tune on a few manually-annotated images ( Chen et al. , 2020 ; Bai et al. , 2019 ; Chen et al. , 2019 ; Chaitanya et al. , 2020a ) . Finally , other approaches aim to transfer knowledge learned from segmenting one class in order to segment a previously unseen class ( Shaban et al. , 2017 ; Rakelly et al. , 2018 ; Roy et al. , 2020 ; Ouyang et al. , 2020 ) . However , these FSL approaches have limitations . For example , effective data augmentation often requires model- and task-specific transformations that are difficult ∗For correspondence , please contact smhooper @ stanford.edu to generalize to all tasks , and transferring knowledge from one segmentation target to another relies on the segmentations of other classes being present in training datasets . More generally , bringing all of the known domain expertise or structural information to bear is prohibitive . Weak supervision ( WS ) is a different approach to training without manual labels that takes advantage of multiple weak labeling sources ( Mintz et al. , 2009 ; Craven & Kumlien , 1999 ; Takamatsu et al. , 2012 ; Gupta & Manning , 2014 ; Ratner et al. , 2019 ) . In many WS workflows , labeling functions ( LFs ) are crafted to encode the decision-making processes of experts into small functions that programatically label training data . LFs can drastically reduce labeling time , but the resulting labels are often noisy and conflicting . The driving idea behind WS is that the accuracies of and correlations between LFs can be learned without access to ground truth labels , then used to aggregate the LF outputs . This WS workflow enables rapid construction of training sets without manual annotation . Due to the widespread success of weak supervision methods ( Dunnmon et al. , 2020 ; Sheng et al. , 2020 ; Bach et al. , 2019 ) , weak supervision is an appealing approach for generating labeled segmentation training sets without relying on hand-labeled data . However , existing WS frameworks do not readily extend to the segmentation setting for the following two reasons : • Injecting knowledge . It is unclear how users can construct multiple segmentation LFs that inject knowledge for a segmentation task at hand—and are suitable for any image . • Modeling LF accuracies . We need to fuse the LFs , but the probabilistic graphical models ( PGMs ) of existing WS frameworks rely on estimating the accuracies for each labeling function , which is a problematic metric for segmentation : it is not sensitive to small differences at segmentation borders nor robust to arbitrary image crops ( Figure 1 , right ) . We propose a best-of-both-worlds approach that fuses FSL and WS for segmentation tasks . Specifically , we use FSL models as LFs , obviating the need to craft expertise-injecting LFs by hand . To resolve the difficulty in tracking accuracies per LF , which are too coarse of a metric for segmentation , we proposing a new WS PGM . Our PGM introduces conditional terms that “ focus ” on areas of contention , where LFs disagree , more finely tracking the LF performance in such areas . We validate our approach on seven medical image segmentation tasks . We show that our method can achieve within 1.4 Dice points of a fully supervised network using only five labeled images . We compare our proposed method to five other few-shot segmentation methods , and find that our method exceeds the next-best by a mean 3.6 Dice points . Finally , we analyze how different FSL methods scale with additional unlabeled training data and compare our PGM to existing WS models . There are two primary contributions in this work . First , we present a weak supervision pipeline tailored to the segmentation use case which utilizes a small number of labeled images and many unlabeled images to train a high-performing segmentation network . Second , we develop a novel PGM to aggregate noisy segmentation masks and show that the parameters of this model can be learned without access to ground truth labels . We show that our new approach to few-shot image segmentation outperforms many competitive baselines . 2 RELATED WORK . Data augmentation Data augmentation has been shown to enable training on few labeled examples , particularly in medical imaging . These augmentations range from simple operations such as affine transforms , elastic transforms , and contrast jitter ( Çiçek et al. , 2016 ; Cireşan et al. , 2011 ) to more complex learned operations , such as the use of GANs to learn intensity and spatial deformation transforms to generate synthetic medical images ( Chaitanya et al. , 2020b ) . Transferring knowledge A popular FSL approach enables prediction on previously unseen classes by leveraging representations learned from prior experience ( Snell et al. , 2017 ; Sung et al. , 2018 ) . Applications to segmentation include Rakelly et al . ( 2018 ) ; Shaban et al . ( 2017 ) ; Dong & Xing ( 2018 ) . For example , PANet is a proposed approach for few-shot semantic segmentation that utilizes representations for each semantic class from the support set , then labels query images by propagating labels from the support set in embedding space ( Wang et al. , 2019 ) . Few-shot segmentation was extended specifically to the medical imaging case by Roy et al . ( 2020 ) , where “ squeeze and excite ” blocks are used to learn class representations without pretrained networks . However , these approaches rely on the availability of labels of other classes during training . Self- and semi-supervised learning Self- and semi-supervised methods have been proposed for image segmentation without relying on the availability of other classes during training ( Chen et al. , 2020 ; Bai et al. , 2019 ; Chen et al. , 2019 ) . Instead , semi-supervised segmentation methods often pretrain on large unlabeled datasets to learn useful representations without any annotations . For example , contrastive learning has been proposed to pretrain CNNs , which are then fine-tuned using a small number of manually segmented images ( Chaitanya et al. , 2020a ) . Additionally , superpixel segmentations have been proposed to self-supervise networks , which can be used to segment test images without fine tuning ( Ouyang et al. , 2020 ) . In the Appendix , we also discuss conditional random fields and using coarse labels ( e.g. , bounding boxes , scribbles , and image-level labels ) to train segmentation models . 3 SEGMENTATION WITH FSL AND WS FUSION . We first give background on WS and state the problem setting . Then , we follow with details of the primary challenges and solutions in each stage of our method . Specifically , we discuss LFs and a new PGM that fuses FSL with WS . We end with a summary of our proposed workflow . 3.1 WEAK SUPERVISION BACKGROUND . We build our approach off of principles from WS . We note that WS is an overloaded term in literature . In this work , we refer to WS as the class of approaches in which multiple noisy labeling sources are aggregated into probabilistic training labels without access to ground truth labels ( Ratner et al. , 2019 ) . The goal of WS is to best utilize multiple sources of noisy signal ( i.e . LFs ) to label training data without relying on manual annotation . Each LF takes as input an unlabeled training point and outputs a vote for the label of that training point . These LFs can replace manual training data annotation , however LF votes can be imperfect and conflict with one another . A latent variable model ( referred to as a label model ) can be used to model these observed LFs and their relationship to the latent ground truth label . In Figure 2a , we show a simple model where five LFs each vote on the label for a single unobserved ground truth label . Importantly , the parameters of this model need to be estimated without access to ground truth annotation . The learned label model can be then used to produce probabilistic training labels to train a downstream network , referred to as an end model . Building off of this past work , our weak supervision pipeline proceeds in three steps : ( 1 ) noisy labels are generated by multiple LFs , ( 2 ) the noisy labels are aggregated into weak labels by a probabilistic graphical model ( PGM ) , and ( 3 ) the weak labels are used to train a segmentation model ( Figure 1 ) . 3.2 PROBLEM SETTING . We are given a training set of n images , X = [ X1 , X2 , ... , Xn ] , where Xi ∈ RI×J×K×T is a tensor of P pixels . I , J , K , and T represent the dimensions of image Xi , where K > 1 for volumetric images and T > 1 for time series images . For each image Xi there is a segmentation mask Yi ∈ { c } I×J×K×T , where { c } is the set of class labels , yielding the set Y = [ Y1 , Y2 , ... , Yn ] of all unobserved ground truth segmentation masks . We have access to m labeling functions that provide noisy segmentation masks λli ∈ { c } I×J×K×T for each Xi , where l ∈ { 1 , 2 , ... , m } . From these observed sources , we seek to generate the set of weak labels Ŷ to estimate Y. Ŷ provides labels for a large dataset , and we can use these to train any end model ( e.g. , a segmentation CNN ) . We tackle the binary segmentation case , { c } = { −1 , 1 } ; simple multiclass extensions are possible . Model Structure The ( unobserved ) ground truth mask Y and the ( observed ) sources λ are jointly modeled by an undirected graphical model ( Koller & Friedman , 2009 ; Wainwright & Jordan , 2008 ) with graph G = ( V , E ) . Since Y is unobserved , this is a latent-variable model . We use conventional structures for each component and combine them in a natural way : • In a standard weak supervision approach , the m weak sources are connected to a groundtruth label Y ( Fig . 2a ) , • The segmentation mask Y is a multi-dimensional grid over the pixels ( Fig . 2b ) , • Our full model has one copy of the WS graph for each pixel in the grid ( Fig . 2c ) . Our goal is to construct effective LFs and to complete the model specification above by defining the potentials associated with the edges in E and the corresponding parameters .

Summary: This paper studies few-shot segmentation for medical images where labeled data is hard to obtain.  It proposed a stage-wise pipeline where pseudo labels are first proposed by noisy learning functions (LFs), then aggregated through a PGM, and finally used to train a segmentation model. The authors argue the two biggest challenges are "injecting knowledge for segmentation model" and "measuring accuracies of pseudo labels", and propose two modules to address them. The experiments have shown that the proposed method outperforming several few-shot learning baselines.