Reinforced active learning for image segmentation

Learning-based approaches for semantic segmentation have two inherent challenges . First , acquiring pixel-wise labels is expensive and time-consuming . Second , realistic segmentation datasets are highly unbalanced : some categories are much more abundant than others , biasing the performance to the most represented ones . In this paper , we are interested in focusing human labelling effort on a small subset of a larger pool of data , minimizing this effort while maximizing performance of a segmentation model on a hold-out set . We present a new active learning strategy for semantic segmentation based on deep reinforcement learning ( RL ) . An agent learns a policy to select a subset of small informative image regions – opposed to entire images – to be labeled , from a pool of unlabeled data . The region selection decision is made based on predictions and uncertainties of the segmentation model being trained . Our method proposes a new modification of the deep Q-network ( DQN ) formulation for active learning , adapting it to the large-scale nature of semantic segmentation problems . We test the proof of concept in CamVid and provide results in the large-scale dataset Cityscapes . On Cityscapes , our deep RL region-based DQN approach requires roughly 30 % less additional labeled data than our most competitive baseline to reach the same performance . Moreover , we find that our method asks for more labels of under-represented categories compared to the baselines , improving their performance and helping to mitigate class imbalance . 1 INTRODUCTION . Semantic segmentation , the task of labelling an image pixel-by-pixel with the category it belongs to , is critical for a variety of applications such as autonomous driving ( Müller et al. , 2018 ; Wang & Pan , 2018 ) , robot manipulation ( Schwarz et al. , 2018 ) , embodied question answering ( Yu et al. , 2019 ) and biomedical image analysis ( Ronneberger et al. , 2015 ) . Convolutional neural networks ( Lecun et al. , 1998 ) -based methods have achieved excellent results on large-scale supervised semantic segmentation , in which we assume pixel-level annotations are available ( Farabet et al. , 2013 ; Pinheiro & Collobert , 2014 ; Long et al. , 2015 ) . For such models to work , however , they need a large amount of pixel-level annotations that may require costly human labor ( Cordts et al. , 2016 ; Bearman et al. , 2016 ) . Current semantic segmentation datasets have pixel-wise annotations for each image . This standard approach has two important issues : ( i ) pixel-level labelling is extremely time consuming . For example , annotation and quality control required more than 1.5h per image ( on average ) on Cityscapes ( Cordts et al. , 2016 ) , a popular dataset used for benchmarking semantic segmentation methods . ( ii ) Class imbalance in the data is typically extreme . Certain categories ( such as ‘ building ’ or ‘ sky ’ ) can appear ∗Work done while interning at ElementAI with two orders of magnitude more frequently than others ( e.g . ‘ pedestrian ’ or ‘ bicycle ’ ) . This can lead to undesired biases and performance properties for learned models . This is specially relevant when we want to collect annotated data with a human in the loop to create a new dataset or to add more labeled data to an existing one . We can tackle the aforementioned problems by selecting , in an efficient and effective way , which regions of the images should be labeled next . Active learning ( AL ) is a well-established field that studies precisely this : selecting the most informative samples to label so that a learning algorithm will perform better with less data than a non-selective approach , such as labelling the entire collection of data . Active learning methods can be roughly divided in two groups : ( i ) methods that combine different manually-designed AL strategies ( Roy & McCallum , 2001 ; Osugi et al. , 2005 ; Gal et al. , 2017 ; Baram et al. , 2004 ; Chu & Lin , 2016 ; Hsu & Lin , 2015 ; Ebert et al. , 2012 ; Long & Hua , 2015 ) and ( ii ) data-driven AL approaches ( Bachman et al. , 2017 ; Fang et al. , 2017 ; Konyushkova et al. , 2017 ; Woodward & Finn , 2016 ; Ravi & Larochelle , 2018 ; Konyushkova et al. , 2018 ) , that learn which samples are most informative to train a model using information of the model itself . Although label acquisition for semantic segmentation is more costly and time consuming than image classification , there has been considerably less work in active learning for semantic segmentation ( Dutt Jain & Grauman , 2016 ; Mackowiak et al. , 2018 ; Vezhnevets et al. , 2012 ; Konyushkova et al. , 2015 ; Gorriz et al. , 2017 ; Yang et al. , 2017 ) , and they focus on hand-crafted strategies . Current AL techniques that use reinforcement learning ( Konyushkova et al. , 2018 ; Fang et al. , 2017 ; Woodward & Finn , 2016 ; Pang et al. , 2018 ; Padmakumar et al. , 2018 ; Bachman et al. , 2017 ) focus on labelling one sample per step until a budget of labels is met . In semantic segmentation , this would translate into labelling a single region per step . This is highly inefficient , since each step involves updating the segmentation network and computing the rewards . In this work , we propose an end-to-end method to learn an active learning strategy for semantic segmentation with reinforcement learning by directly maximizing the performance metric we care about , Intersection over Union ( IoU ) . We aim at learning a policy from the data that finds the most informative regions on a set of unlabeled images and asks for its labels , such that a segmentation network can achieve high-quality performance with a minimum number of labeled pixels . Selecting regions , instead of entire images , allows the algorithm to focus on the most relevant parts of the images , as shown in Figure 1 . Although class imbalance in segmentation datasets has been previously addressed in ( Badrinarayanan et al. , 2017 ; Chan et al. , 2019 ; Sudre et al. , 2017 ) , among others , they try to solve a problem that arises from the data collection process . We show that our proposed method can help mitigate the problem at its source , i.e . in the data annotation itself . Because our method maximizes the mean IoU per class , it indirectly learns to ask for more labels of regions with under-represented classes , compared to the baselines . Moreover , we propose and explore a batch-mode active learning approach that uses an adapted DQN to efficiently chose batches of regions for labelling at each step . To the best of our knowledge , all current approaches for active learning in semantic segmentation rely on hand-crafted active learning heuristics . However , learning a labelling policy from the data could allow the query agent to ask for labeled data as a function of the data characteristics and class imbalances , that may vary between datasets . Our main contributions can be summarized as follows : ( i ) we learn a RL-based acquisition function for region-based active learning for segmentation , ( ii ) we formulate our active learning framework with a batch-mode DQN , which labels multiple regions in parallel at each active learning iteration ( a more efficient strategy for large-scale datasets that is compatible with standard mini-batch gradient descent ) , and ( iii ) we test the proof of concept in CamVid ( Brostow et al. , 2008 ) dataset and provide results in Cityscapes ( Cordts et al. , 2016 ) dataset , beating a recent state-of-the-art technique known as BALD ( Gal et al. , 2017 ) , a widely used entropy-based selection criterion and uniform sampling baselines . 2 RELATED WORK . Active learning . Traditional active learning techniques focus on estimating the sample informativeness using hand-crafted heuristics derived from sample uncertainty : employing entropy ( Shannon , 1948 ) , query-by-committee ( Dagan & Engelson , 1995 ; Shannon , 1948 ; Freund et al. , 1993 ) , maximizing the error reduction ( Roy & McCallum , 2001 ) , disagreement between experts ( Dagan & Engelson , 1995 ; Freund et al. , 1993 ) or Bayesian methods that need to estimate the posterior distribution ( Houlsby et al. , 2011a ; Gal et al. , 2017 ) . Some approaches combine different techniques to improve AL performance . For instance , relying on exploration-exploitation trade-offs ( Osugi et al. , 2005 ) , on a bandit formulation ( Baram et al. , 2004 ; Chu & Lin , 2016 ; Hsu & Lin , 2015 ) and on reinforcement learning ( Ebert et al. , 2012 ; Long & Hua , 2015 ) . However , these methods are still limited in the sense that they combine hand-crafted strategies instead of learning new ones . More recent active learning methods rely on an acquisition function that estimates the sample informativeness with a learned metric . Konyushkova et al . ( 2017 ) estimate the error reduction of labelling a particular sample , choosing the ones that maximize the error reduction . Wang et al . ( 2017 ) introduce a cost-effective approach that also uses confident predictions as pseudo ground truth labels . AL with reinforcement learning . Recently , reinforcement learning has gained attention as a method to learn a labelling policy that directly maximizes the learning algorithm performance . For instance , Liu et al . ( 2018 ) ; Bachman et al . ( 2017 ) leverage expert knowledge from oracle policies to learn a labelling policy , and Pang et al . ( 2018 ) ; Padmakumar et al . ( 2018 ) rely on policy gradient methods to learn the acquisition function . In a different approach , some methods gather all labeled data in one big step . In Contardo et al . ( 2017 ) , all samples are chosen in one step with a bi-directional RNN for the task of one-shot learning . In Sener & Savarese ( 2018 ) , they propose to select a batch of representative samples that maximize the coverage of the entire unlabeled set . However , the bounded core-set loss used tends to perform worse when the number of classes grows . More similar to our approach , some prior works propose to learn the acquisition function with a Deep Q-Network ( DQN ) ( Mnih et al. , 2013 ) formulation . These works have examined both stream-based active learning ( Fang et al. , 2017 ; Woodward & Finn , 2016 ) , where unlabeled samples are provided one by one , and the decision is to label it or not , and pool-based active learning ( Konyushkova et al. , 2018 ) , where all the unlabeled data is provided beforehand , and the decision is later taken on which samples to choose . The work of Konyushkova et al . ( 2018 ) is the closest to ours . Similar to them , our method also leverages the benefits of Q-learning ( Watkins & Dayan , 1992 ) to tackle pool-based AL . Contrary to them , we deal with a much more complex problem : semantic segmentation versus simple classification on UCI repository ( Dua & Graff , 2017 ) . The large-scale nature of the problem requires us to use a very different definition of actions , states and rewards . Moreover , we need to adapt the DQN formulation to allow the problem to be computationally feasible . AL for semantic segmentation . Active learning for semantic segmentation has been relatively less explored than other tasks , potentially due to its large-scale nature . For instance , Dutt Jain & Grauman ( 2016 ) combine metrics ( defined on hand-crafted heuristics ) that encourage the diversity and representativeness of labeled samples . Some rely on unsupervised superpixel-based oversegmentation ( Vezhnevets et al. , 2012 ; Konyushkova et al. , 2015 ) – and highly depend on the quality of the super-pixel segmentation . Others focus on foreground-background segmentation of biomedical images ( Gorriz et al. , 2017 ; Yang et al. , 2017 ) , also using hand-crafted heuristics . Settles et al . ( 2008 ) ; Vijayanarasimhan & Grauman ( 2009 ) ; Mackowiak et al . ( 2018 ) focus on cost-effective approaches , proposing manually-designed acquisition functions based on the cost of labeling images or regions of images . However , this information is not always given , restricting their applicability . Mackowiak et al . ( 2018 ) focus on cost-effective approaches , where the cost of labeling an image is not considered equal for all images . Similar to our work , they use a region-based approach to cope with the large number of samples on a segmentation dataset . Contrary to us , their labelling strategy is based on manually defined heuristics , limiting the representability of the acquisition function . To the best of our knowledge , our work is the first to apply data-driven RL-based approach to the problem of active learning for semantic segmentation .

Summary: The paper works on active learning for semantic segmentation, aiming to annotate as few "blocks/patches" as possible while training a strong model. The authors proposed to learn a query policy via Q learning, and design states and actions specifically for segmentation. The experimental results show that the learned policy can attend to informative patches and rare classes to learn a model faster and efficiently.