Bounding Box Tightness Prior for Weakly Supervised Image Segmentation

Suppose $I$ denotes an input image, and $Y\in\{1,2,\cdots,C\}$ is its corresponding pixel-level category label, in which $C$ is the number of categories of the objects. The image segmentation problem is to obtain the prediction of $Y$, denoted as $P$, for the input image $I$.

In the fully supervised image segmentation setting, for an image $I$, its pixel-wise category label $Y$ is available during training. Instead, in this study we are only provided its bounding box label $B$. Suppose there are $M$ bounding boxes, then its bounding box label is $B=\{b_{m},y_{m}\},m=1,2,\cdots,M$, where the location label $b_{m}$ is a 4-dimensional vector representing the top left and bottom right points of the bounding box, and $y_{m}\in\{1,2,\cdots,C\}$ is its category label.

This study considers deep neural networks which output the pixel-wise prediction of the input image, such as Unet [16], FCN [12], etc. Due to the possible overlaps of objects of different categories in images, especially in medical images, the image segmentation problem is formulated as a multi-label classification problem in this study. That is, for a location $k$ in the input image, it outputs a vector ${\mathbf{p}}_{k}$ with $C$ elements, one element for a category; each element is converted to the range of $[0,1]$ by the sigmoid function.

Multiple instance learning (MIL) is a type of supervised learning. Different from the traditional supervised learning which receives a set of training samples which are individually labeled, MIL receives a set of labeled bags, each containing many training samples. In MIL, a bag is labeled as negative if all of its samples are negative, a bag is positive if it has at least one sample which is positive.

Tightness prior of bounding box indicates that the location label of bounding box is the smallest rectangle enclosing the whole object, thus the object must touch the four sides of its bounding box, and does not overlap with the region outside its bounding box. The crossing line of a bounding box is defined as a line with its two endpoints located on the opposite sides of the box. In an image $I$ under consideration, for an object with category $c$, any crossing line in the bounding box has at least one pixel belonging to the object in the box; any pixels outside of any bounding boxes of category $c$ do not belong to category $c$. Hence pixels on a cross line compose a positive bag for category $c$, while pixels outside of any bounding boxes of category $c$ are used for negative bags.

For category $c$ in an image, the baseline approach simply considers all of the horizontal and vertical crossing lines inside the boxes as positive bags, and all of the horizontal and vertical crossing lines that do not overlap any bounding boxes of category $c$ as negative bags. This definition is shown in Fig. 1(a) and has been widely employed in the literature [4, 5].

To optimize the network parameters, MIL loss with two terms are considered [4]. For category $c$, suppose its positive and negative bags are denoted by $\mathcal{B}_{c}^{+}$ and $\mathcal{B}_{c}^{-}$, respectively, then loss $\mathcal{L}_{c}$ can be defined as follows:

where $\phi_{c}$ is the unary loss term, $\varphi_{c}$ is the pairwise loss term, and $\lambda$ is a constant value controlling the trade off between the unary loss and the pairwise loss.

The unary loss $\phi_{c}$ enforces the tightness constraint of bounding boxes on the network prediction $P$ by considering both positive bags $\mathcal{B}_{c}^{+}$ and negative bags $\mathcal{B}_{c}^{-}$. A positive bag contains at least one pixel inside the object, hence the pixel with highest prediction tends to be inside the object, thus belonging to category $c$. In contrast, no pixels in negative bags belong to any objects, hence even the pixel with highest prediction does not belong to category $c$. Based on these observations, the unary loss $\phi_{c}$ can be expressed as:

where $P_{c}(b)=\max_{k\in b}(p_{kc})$ is the prediction of the bag $b$ being positive for category $c$, $p_{kc}$ is the network output of the pixel location k for category $c$, and $|\mathcal{B}|$ is the cardinality of $\mathcal{B}$.

The unary loss is binary cross entropy loss for bag prediction $P_{c}(b)$, it achieves minimum when $P_{c}(b)=1$ for positive bags and $P_{c}(b)=0$ for negative bags. More importantly, during training the unary loss adaptively selects a positive sample per positive bag and a negative sample per negative bag based on the network prediction for optimization, thus yielding an adaptive sampling effect.

However, using the unary loss alone is prone to segment merely the discriminative parts of an object. To address this issue, the pairwise loss as follows is introduced to pose the piece-wise smoothness on the network prediction.

where $\varepsilon$ is the set containing all neighboring pixel pairs, $p_{kc}$ is the network output of the pixel location k for category $c$.

Finally, considering all $C$ categories, the MIL loss $\mathcal{L}$ is:

For an object of height $H$ pixels and width $W$ pixels, the positive bag definition in the MIL baseline yields only $H+W$ positive bags, a much smaller number when compared with the size of the object. Hence it limits the selected positive samples during training, resulting in a bottleneck for image segmentation.

To eliminate this issue, this study proposes to generalize positive bag definition by considering all parallel crossing lines with a set of different angles. An parallel crossing line is parameterized by an angle $\theta\in(-90^{\circ},90^{\circ})$ with respect to the edges of the box where its two endpoints located. For an angle $\theta$, two sets of parallel crossing lines can be obtained, one crosses up and bottom edges of the box, and the other crosses left and right edges of the box. As examples, in Fig. 1(b), we show positive bags of two different angles, in which those marked by yellow dashed colors have $\theta=25^{\circ}$, and those marked by green dashed lines are with $\theta=0^{\circ}$. Note the positive bag definition in MIL baseline is a special case of the generalized positive bags with $\theta=0^{\circ}$.

The similar issue also exists for the negative bag definition in the MIL baseline. To tackle this issue, for a category $c$, we propose to define each individual pixel outside of any bounding boxes of category $c$ as a negative bag. This definition greatly increases the number of negative bags, and forces the network to see every pixel outside of bounding boxes during training.

The generalized MIL definitions above will inevitably lead to imbalance between positive and negative bags. To deal with this issue, we borrow the concept of focal loss [17] and use the improved unary loss as follows:

where $N^{+}=\max(1,|\mathcal{B}_{c}^{+}|)$, $\beta\in[0,1]$ is the weighting factor, and $\gamma\geq 0$ is the focusing parameter. The improved unary loss is focal loss for bag prediction, it achieves minimum when $P_{c}(b)=1$ for positive bags and $P_{c}(b)=0$ for negative bags.

To demonstrate the effectiveness of the proposed generalized MIL formulation, in Table 1 we show the performance of the generalized MIL for the two datasets. For comparison, we also show the results of the baseline method. It can be observed that the generalized MIL approach consistently outperforms the baseline method at different angle settings. In particular, the PROMISE12 dataset got best Dice coefficient of 0.878 at $\theta_{best}=(-40^{\circ},40^{\circ},20^{\circ})$ for the generalized MIL, compared with 0.859 for the baseline method. The ATLAS dataset achieved best Dice coefficient of 0.474 at $\theta_{best}=(-60^{\circ},60^{\circ},30^{\circ})$ for the generalized MIL, much higher than 0.408 for the baseline method.

To demonstrate the benefits of the smooth maximum approximation, in Table 2 we show the performance of the MIL baseline method when the smooth maximum approximation was applied. As can be seen, for the PROMISE12 dataset, the better performance is obtained for $\alpha$-softmax function with $\alpha=4$ and 6 and for $\alpha$-quasimax function with $\alpha=6$ and 8. For the ATLAS dataset, the improved performance can also be observed for $\alpha$-softmax function with $\alpha=6$ and 8 and for $\alpha$-quasimax function with $\alpha=8$.