Spatial-context-aware deep neural network for multi-class image classification

The overall framework of the proposed model is presented in Fig. 2. It consists of three main modules: 1) A feature extractor which integrates ResNeXt-101 [29] with feature pyramid networks (FPN) [30]. 2) Two patch generators. To utilize both spatial and contextual information, the object branch follows the structure in [31] to localize the objects and the context branch considers the image context information. These fragments and feature maps are combined to generate the final visual features. 3) Four patch processors. One regressor for positioning, one classifier for determining the existence of objects in the patch, and two predictors for predicting the label of the patch using a formulation of image multi-labelling.

The patch labelling results from the object branch and context branch are combined by an element-wise classifier fusion to generate the final predicted labels.

ResNeXt-101 [29] followed by FPN [30] is utilized as the feature extractor. The former aggregates a set of transformations to improve the classification capabilities of deep neural networks, and the latter employs the pyramid representations to extract a rich visual semantic abstract. The last pooling and classification layers in ResNeXt-101 are removed and the feature maps from the last convolutional layer are used as the input of FPN. A 4-stage semantic feature pyramid is built in FPN from high to low resolution. Denote the two successive feature extracting networks as $f_{R}\left(\mathbf{U};\theta_{R}\right)$ and $f_{F}\left(\mathbf{Z};\theta_{F}\right)$ respectively. Given an input image $\mathbf{I}$, the final generated pyramidal feature $\mathbf{X}$ can be obtained as:

Based on the generated feature maps, a set of bounding boxes are generated to locate objects, covering both the contextual and spatial information of the objects. However, a tightly cropped bounding box may only contain the information of the object, but ignore the contextual information surrounding the object.

The tightly cropped bounding box in the object branch may neglect the contextual information, so the context branch is utilized to expand it into a larger one with potential object and its surrounding image content. Note that the expanded context patch is only used to determine the label of the object, but not for locating the object. As Fig. 1 shows, confidence in predicting ‘tennis racket’ increases as information about the surrounding environment, e.g., ‘person’, is included.

The tightly cropped object patch contains the most discriminant information for classify the object, while the expanded image patch containing both the object and additional contextual information. Hence, two branches with different objectives are designed to generate image patches. Specifically, as shown in Fig. 2, the generator in the object branch (left) takes the feature maps as the input and produces a set of rectangular object patches with confidence scores. Only objects within the rectangles are evaluated in this branch, and it is often difficult to detect small objects. Therefore, a patch expansion mechanism is applied to the generator in the context branch. This design allows more contextual information to be captured, and hence improves the accuracy of the patch prediction. Two branches are used together to predict the labels of the input image.

The generated object patches and context patches are then fed into four dense networks for location regression, patch classification and label prediction, respectively.

The first dense network, location regressor, is designed to accurately locate the objects to explore and utilize the image-level spatial information of different labels. It is guided by the location regression loss $L_{r}$. The objective is to maximize the intersection over union (IOU) between the generated bounding boxes and the ground-truth bounding boxes.

where $\hat{t}_{i}$ and $t_{i}^{*}$ represent the predicted patches and ground-truth bounding boxes, respectively and x, y, w, h correspond to the centre coordinates, height and width of the proposed patch. $\phi\left(x\right)$ is a non-sensitive outlier-removal function for enhancing the training robustness.

The second dense network, patch classifier, determines the confidence whether an object exists in the bounding box. Hence it is a binary classification problem. The cross-entropy loss is used to guide the training process and the loss is defined as follows:

where $\hat{p}$ and $p^{*}$ represent the output confidence from the network and ground-truth, respectively. $p^{*}$ would be assigned to $1$ if the maximum IOU between the current patch and the ground-truth one exceeds a certain threshold, and 0 otherwise.

The remaining two dense networks are label predictors to determine which object the bounding box contains. These two dense networks are trained by using the binary cross-entropy (BCE) loss, which can exploit label dependencies in the multi-labelling tasks [10, 7, 8].

Given a one-hot vector of ground-truth labels $\mathbf{y^{*}}=\left[y^{1},y^{2},\ldots,y^{C}\right]^{T}$, where $C$ represents the number of all possible labels in the dataset. $y^{i}$ is a binary indicator that $y^{l}=1$ if the image contains the label $l$, and $y^{l}=0$ otherwise. The label-prediction loss is denoted as follows:

Similar to $\mathbf{y^{*}}$ and $y^{i}$, $\mathbf{\hat{y}}$ and $\hat{y}^{i}$ denote the predicted confidence over all possible labels and the $i$-th category. $\sigma$ in $L_{l}$ is a learnable weighting factor, we fix it as 0.5 given by the empirically study.

Hence the total loss function $L$ is represented as follows:

where $\alpha,\beta$ and $\gamma$ is scale values that utilized to balance the losses and $L_{l}^{O}$ and $L_{l}^{C}$ represent the loss of the object branch and context branch respectively.

The proposed model is trained on PyTorch [35]. Stochastic gradient descend strategy is employed for training, with a mini-batch size of 2, a momentum of 0.5 and a decay rate of 0.01. The initial learning rate is 0.002. All images are resized to 416$\times$416 before passing to the model. The proposed model is optimized by using the transfer-learning techniques, i.e., load and freeze the pre-trained weights of the backbone, and then train the remaining modules. Once the training converges, the whole model is jointly optimized.

The MS-COCO dataset [27] and PASCAL VOC [28] are two widely used public benchmark datasets for image labelling. The former contains 80 object categories and the latter covers 20 different categories. All instances in the two datasets are associated with at least one label. The standard evaluation protocol is used as in [11, 13, 18, 15, 19, 26, 14].

Standard evaluation metrics such as mean average precision (mAP), F1-score, precision and recall are used in experiments. The prefix ‘macro-’ indicates the average over all categories, while ‘micro-’ indicates the average over all samples. Therefore, ‘macro-’ is susceptible to the rare categories and ‘micro-’ is easily dominated by the major classes [36]. Denote ‘macro-’ as ‘-C’ and ‘micro-’ as ‘-O’.

The proposed method is compared against state-of-the-art models including CNN-RNN [7], ResNet101 [5], RNN-Attention [32], ResNet101-SRN [10], RNN-frequency [33], DELTA [34], ResNet101-ACfs [11], DecoupleNet [12], ML-GCN [13], ResNet101-CRL [18], KSSNet [15], MS-CMA [19], WSL-GCN [26] and C-Tran [14], among which CNN-RNN [7], DecoupleNet [12], ML-GCN [13], KSSNet [15] and WSL-GCN [26] exploit semantic relations by capturing global label dependencies. ResNet101 [5] is initially designed for single-label image classification while also performs well in multi-labelling by using appropriate loss functions. RNN-Attention [32], ResNet101-SRN [10], RNN-frequency [33], DELTA [34], ResNet101-ACfs [11] and ResNet101-ACfs [11] exploit spatial information by extending attention ideas, while ResNet101-CRL constructs explicit context relations by feature-label co-projection. MS-CMA [19] adapts the attention mechanism to a cross-modality version with graph embedding. C-Tran [14] exploits the dependencies among visual features and labels to tackle the image multi-labeling.

The comparison results to the state-of-the-art approaches on the MS-COCO dataset are summarized in Table 1. It is clear to see that the proposed model significantly outperforms all the state-of-the-art models in terms of the key evaluation metrics such as mAP and F1-score.

As shown in bold in Table 1, the proposed method achieves an overall mAP of 86.3$\%$, a per-class F1 score of 80.6$\%$ and an overall F1-score of 83.1$\%$, which greatly surpasses the previous best solution, C-Tran [14] by 1.2$\%$, 0.7$\%$ and 1.4$\%$, respectively. It shows that the proposed approach performs well on exploiting the semantic information, spatial and object context information and label dependencies.

Since the results of the state-of-the-art models are collected from the corresponding origional paper, only eight previous solutions are selected for the performance comparison on PASCAL VOC [28] dataset, and only comparing on the key evaluation metrics, mAP. The comparison results are summarized in Table 2. The proposed model achieves an mAP of 95.3 $\%$, surpassing all the state-of-the-art models.

To visually demonstrate the effectiveness of the proposed model, several sample detection results for are shown in Fig. 3. The marked labels and bounding boxes prove that the proposed method can effectively utilize the context information. Objects in red are the ones that can not be detected by previous approaches but now can be detected by the proposed model by utilizing the spatial and contextual information. From Fig. 3, it can be seen that the proposed model plays a key role in the annotation of the easily neglected small objects such as the knife on the table and the person sailing the boat.