Disease Severity Regression with Continuous Data Augmentation

TL;DR

This paper introduces a continuous data augmentation method using a novel GAN to generate medical images at real-valued severity levels, improving disease severity regression accuracy.

Contribution

It proposes a continuous severity GAN and dataset-disjoint multi-objective optimization to enhance medical image data augmentation for severity estimation.

Findings

Achieved higher classification performance than conventional methods.

Effectively generated images at real-valued severity levels.

Improved disease severity regression accuracy.

Abstract

Disease severity regression by a convolutional neural network (CNN) for medical images requires a sufficient number of image samples labeled with severity levels. Conditional generative adversarial network (cGAN)-based data augmentation (DA) is a possible solution, but it encounters two issues. The first issue is that existing cGANs cannot deal with real-valued severity levels as their conditions, and the second is that the severity of the generated images is not fully reliable. We propose continuous DA as a solution to the two issues. Our method uses continuous severity GAN to generate images at real-valued severity levels and dataset-disjoint multi-objective optimization to deal with the second issue. Our method was evaluated for estimating ulcerative colitis (UC) severity of endoscopic images and achieved higher classification performance than conventional DA methods.

Tables

Table 1. Table 1 : Classification performance of each method. ‘*’ denotes a statistically significant difference ( p < 0.05 𝑝 0.05 p<0.05 in paired t-test) between the baseline and the other methods.

Method	Precision	Recall	F1-score
Regression (Baseline)	0.782${}^{}$	0.631${}^{}$	0.652${}^{}$
+ Classic DA	0.731${}^{}$	0.657${}^{}$	0.668${}^{}$
+ GAN-based DA	0.697∗	0.651${}^{}$	0.663${}^{}$
C-DA w/o GAN	0.744${}^{}$	0.629${}^{}$	0.648${}^{}$
C-DA ($ϵ=1$)	0.743${}^{}$	0.624${}^{}$	0.638${}^{}$
C-DA ($ϵ=0.5$)	0.717∗	0.690∗	0.696∗
C-DA ($ϵ=0.25$)	0.688∗	0.672∗	0.675${}^{}$

Equations

$${\mathcal{L}}_{\mathrm{order}}={\mathbb{E}}_{\mathbf{z}}\left[\sum _{y=2}^{L-1}|{\mathbf{s}}_y(\mathbf{z})-{\tilde{\mathbf{s}}}_y(\mathbf{z})|\right],$$

$${\mathbf{s}}_{y+1}(\mathbf{z})-{\mathbf{s}}_y(\mathbf{z})={\mathbf{s}}_y(\mathbf{z})-{\mathbf{s}}_{y-1}(\mathbf{z}).$$

$${\mathbf{s}}_{y+ϵ}=(1-ϵ){\mathbf{s}}_y+ϵ{\mathbf{s}}_{y+1}.$$

Full text

**Index Terms— ** Data augmentation, generative adversarial network, endoscopic images

1 Introduction

Disease severity regression is a task to determine a function $f(\bm{x})$ that satisfies $y_{n}\sim f(\bm{x}_{n})$ for a given dataset $\Omega=\{(\bm{x}_{n},y_{n}),n\in\{1,\ldots,N\}\}$, where $\bm{x}_{n}$ is a medical image, such as an endoscopic image, and $y_{n}$ is its severity level. Nowadays, it is common to use a convolutional neural network (CNN) as the model of $f(\bm{x})$ because CNN has a powerful representation ability to deal with the nonlinear relationship between image appearance and its severity. It is also common to use $L$ discrete severity levels as $y_{n}$. For example, Mayo scores of endoscopic images with ulcerative colitis (UC) have $L=4$ levels.

If the labeled dataset $\Omega$ is too small to train the CNN, data augmentation (DA) is often employed to generate synthetic data $\Omega’=\{(\bm{x}^{\prime}_{m},y^{\prime}_{m}),m\in\{1,\ldots,M\}\}$ from $\Omega$. A possible DA technique is a conditional generative adversarial network (cGAN). Given a discrete severity level $y^{\prime}_{m}\in\{1,\ldots,L\}$ as the condition, cGAN generates various images $\bm{x}^{\prime}_{m}$ at the severity level $y^{\prime}_{m}$. The generated images $\Omega^{\prime}$ are then used to train the CNN together with the original dataset $\Omega$.

This paper focuses on two issues of the above cGAN-based DA for disease severity regression. The first issue is that disease severity is inherently continuous, so we do not need to adhere to the discrete conditions as $y^{\prime}_{m}$. In other words, generating images at real-valued severity levels $y^{\prime}_{m}$ will help train the CNN $f$ appropriately. The second issue is that the severity of the generated image is not very reliable. Even if we generate an image $\bm{x}^{\prime}_{m}$ with the condition $y^{\prime}_{m}$, there might be a risk that the visual severity of $\bm{x}^{\prime}_{m}$ is precisely equal to $y^{\prime}_{m}$.

We propose a continuous DA scheme, where a new technique tackles each issue. For the first issue, we propose a continuous severity GAN (csGAN). Fig. 1 (a) shows the overview of csGAN. Our csGAN is trained with images with discrete levels ($y_{n}\in\{1,\ldots,L\}$) but can generate images at real-valued severity levels ($y^{\prime}_{m}\in[1,L]$).

For the second issue, we use a dataset-disjoint multi-objective optimization, where the original dataset $\Omega$ (with discrete levels) and the augmented dataset $\Omega^{\prime}$ (with real-valued levels) are used in different ways according to their different reliability. Specifically, as shown in Fig. 1 (b), we train a CNN $f$ with a regression loss for $\Omega$ and a ranking loss for $\Omega^{\prime}$. The former works to satisfy $y_{n}\sim f(\bm{x}_{n})$ and the latter $f(\bm{x}^{\prime}_{m})\lessgtr f(\bm{x}^{\prime}_{k})$ when $y^{\prime}_{m}\lessgtr y^{\prime}_{k}$. This means that the levels $\{y^{\prime}_{m}\}$ of the augmented data are not used as absolute ground truth but as relative conditions for training $f$.

The proposed techniques are evaluated by using a UC image dataset. As a qualitative evaluation, we observe the images by csGAN and confirm that we can continuously control the visual severity level of the generated images. As a quantitative evaluation, we confirm that our continuous DA helps to improve the severity regression performance.

Our main contributions are summarized as follows:

• •

We propose csGAN, which can generate images at real-valued severity levels.

• •

We also propose to use dataset-disjoint multi-objective optimization for the disease severity regression task with an augmented dataset.

• •

Experimental evaluations with a UC image dataset show the performance superiority of our continuous DA scheme using the above two techniques over a baseline and other cGAN-based DA.

2 Related Work

Conditional GANs: Various cGANs have been proposed so far, [1, 2, 3, 4] and they assume various types of conditions. For example, in the pix2pix [5] framework, an image is given as a condition. The most common condition is class labels – they can be given as a one-hot vector or discrete number. In other words, for specifying the target type of generated images, it is not common to give a condition by a real-valued number (such as 1.33 and 0.28). Exceptionally, CcGAN [6] accepts real-valued conditions; however, it relies on a hard assumption that real-valued annotation has already been attached to each training sample. In contrast, our csGAN can be trained with discrete conditions but still can generate images at real-valued conditions.

DA for medical images: Due to a high cost for annotation, medical image analysis tasks often suffer from a limited number of labeled data and thus employ DA methods. According to a survey paper in 2021 [7], basic augmentation techniques, such as linear and nonlinear geometric transformations and intensity level perturbations, are still the majority for medical image DA. However, the survey also shows that GAN-based DA methods have increased in recent papers (such as [8, 9, 10]). The above review for cGANs says that GAN-based DA for medical images has also not dealt with real-valued conditions. Moreover, to the authors’ best knowledge, the augmented dataset is simply merged with the original dataset without any special treatment.

3 Continuous data augmentation

This section describes two techniques in continuous DA, i.e., csGAN and dataset-disjoint multi-task optimization. The former is a new conditional GAN trained to generate images at real-valued severity levels. The latter is a technique to train the regression model $f$ by using the original dataset $\Omega$ and the augmented dataset $\Omega^{\prime}$ in different manners by considering the reliability of the severity levels of the augmented data.

3.1 Continuous Severity GAN (csGAN)

As noted in Section 1, we propose csGAN to generate images at real-valued severity levels. Inspired by StarGAN v2[11], csGAN comprises four modules: a mapping network $F$, a generator $G$, a style encoder $E$, and a discriminator $D$, as shown in Fig. 1(a). csGAN uses a style vector $\bm{s}_{y}(\bm{z})$ as a condition to generate images at the level $y$, where $\bm{z}\sim\mathcal{N}(\bm{0},\bm{I})$. A different $\bm{z}$ results in a different style vector $\bm{s}_{y}(\bm{z})$ and finally contributes to having a different generated image at the level $y$. Hereafter, we often denote $\bm{s}_{y}(\bm{z})$ as $\bm{s}_{y}$ for simplicity.

Due to the page limitation, we briefly summarize the roles of four modules $F$, $G$, $E$, and $D$:

• •

$F$ accepts a random vector $\bm{z}$ and then outputs $L$ style vectors $\bm{s}_{1},\ldots,\bm{s}_{L}$ at once.

• •

$G$ accepts a real or generated (i.e., fake) image $\bm{x}$ and the condition $\bm{s}_{y}$ and then generates an image $\bm{x}^{\prime}$ at $y$.

• •

$E$ accepts a real or fake image $\bm{x}$ with its level $y$ and then estimates its style vector while expecting the estimated vector is similar to $\bm{s}_{y}$ input to $G$.

• •

$D$ is a standard discriminator for real/fake decisions of $\bm{x}$.

Those modules are trained to achieve cycle consistency; a generated image $\bm{x}^{\prime}=G(\bm{x}_{n},\bm{s}_{y})$ for the level $y\neq y_{n}$ needs to satisfy the condition $\bm{x}_{n}\sim G(\bm{x}^{\prime},\bm{s}_{y_{n}})$. (Note that $\bm{x}_{n}$ is a real image at the level $y_{n}$.) By this cycle consistency, we have a level-$y$ version of $\bm{x}_{n}$ and a level-$y_{n}$ version of $\bm{x}^{\prime}$. Consequently, we have images at all $L$ levels, even from a single image at a certain level $y$.

For generating images at real-valued levels, csGAN introduces an additional loss function, called a order loss, for $F$:

[TABLE]

where $\tilde{\bm{s}}_{y}(\bm{z})=(\bm{s}_{y-1}(\bm{z})+\bm{s}_{y+1}(\bm{z}))/2$. With the order loss, we expect that the style vectors from the same $\bm{z}$ will have a linear property, that is,

[TABLE]

This linear property will allow us to consider a real-valued severity level $y+\epsilon$, where $\epsilon\in(0,1)$. More specifically, we can derive the style vector for the real-valued level $y+\epsilon$ by the linear interpolation,

[TABLE]

As shown in Fig. 1 (a), at the test phase, we use $G$ to generate an image at a real-valued severity level $y^{\prime}=y+\epsilon\in[1,L]$. First, $\{\bm{s}_{1}(\bm{z}),\ldots,\bm{s}_{L}(\bm{z})\}$ is obtained by the mapping network $F$ with a $\bm{z}$. Then, for a certain $\epsilon$, $\bm{s}_{y^{\prime}}$ is determined by Eq. (3). Finally, a level-$y^{\prime}$ version of an image $\bm{x}$ is generated by $\bm{x}^{\prime}=G(\bm{x},\bm{s}_{y^{\prime}})$.

3.2 Learning by Dataset-Disjoint Multi-Objective Optimization

As noted in Section 1, the severity level $y^{\prime}$ of the generated data is not very reliable. Especially, since we used a simple linear style vector interpolation of Eq. (3), we cannot guarantee that the generated data of the level $y^{\prime}=y+\epsilon$ has exact visual characteristics as the level $y^{\prime}$. In other words, the level $y^{\prime}$ is not fully reliable as the absolute level.

However, $y^{\prime}$ is still reliable as a relative level; for a pair of real-valued levels $y^{\prime}_{m}$ and $y^{\prime}_{k}$ (where $y^{\prime}_{m}>y^{\prime}_{k}$), the generated images $\bm{x}^{\prime}_{m}$ and $\bm{x}^{\prime}_{k}$ are expected to show the same relative order in their severity levels, that is, $f(\bm{x}^{\prime}_{m})>f(\bm{x}^{\prime}_{k})$. By training the model $f$ to satisfy this relative condition (instead of training $f$ to satisfy $y_{m}^{\prime}\sim f(\bm{x}^{\prime}_{m})$), we can utilize the augmented data by csGAN in an appropriate manner.

Considering the above property of the generated data, we use dataset-disjoint multi-objective optimization scheme to train the regression model $f$, as shown in Fig. 1 (b). Assume we have an original dataset $\Omega=\{(\bm{x}_{n},y_{n})\}$ with manually-annotated discrete severity levels $y_{n}\in\{1,L\}$ and a generated image dataset $\Omega^{\prime}=\{(\bm{x}^{\prime}_{m},y^{\prime}_{m})\}$ at various real-valued levels $y^{\prime}_{m}\in[1,L]$. Then, the CNN-based regression model $f$ is trained with both datasets $\Omega$ and $\Omega^{\prime}$ in different usages. Since $y_{n}$ is reliable as an absolute level, the image $\bm{x}_{n}$ in $\Omega$ are used to train $f$ to satisfy $y_{n}\sim f(\bm{x}_{n})$. Here, we use a mean squared error loss $\sum_{n=1}^{N}(f(\bm{x}_{n})-y_{n})^{2}/N$. On the other hand, since $y^{\prime}_{m}$ is reliable as a relative level, images $\bm{x}^{\prime}_{m}$ and $\bm{x}^{\prime}_{k}$ in $\Omega^{\prime}$ with the relative relationship $y_{m}^{\prime}>y_{k}^{\prime}$ are used to train $f$ to satisfy $f(\bm{x}^{\prime}_{m})>f(\bm{x}^{\prime}_{k})$. Here, we use the loss function of ListNet [12], which is one of the most popular methods for learning-to-rank. These two loss functions are balanced by a hyperparameter, which is optimized by a validation set.

4 Experimental Results

4.1 Experimental Setup

Dataset: To evaluate the proposed method (continuous DA, C-DA in short), we used a dataset of UC endoscopic images collected from the Kyoto Second Red Cross Hospital. The dataset contains 10,265 images from 388 patients. All images are annotated with discrete Mayo scores $y_{n}\in\{1,2,3,4\}$ by multiple experts and resized to 256 $\times$ 256 pixels. The distribution of Mayo scores is 6,678, 1,995, 1,395, and 197 images for Mayo 0, 1, 2, and 3, respectively. Note that Mayo 0 corresponds to the level $y_{n}=1$ and Mayo 3 to $y_{n}=4=L$.

Fig. 2 shows several examples of endoscopic images for each Mayo score. Schroeder et al.[13] categorized the endoscopic findings of UC as follows: Mayo 0 is a normal or inactive disease, Mayo 1 is a mild disease (erythema, decreased vascular pattern, etc.), Mayo 2 is a moderate disease (marked erythema, erosions, etc.), Mayo 3 is a severe disease (spontaneous bleeding, ulceration, etc.).

We performed five-fold cross-validation. The dataset was divided into training, validation, and test sets at 60, 20, and 20%, respectively. The splittings were performed by random patient-disjoint sampling, and the class ratios for each set were maintained. Moreover, random oversampling was used to mitigate class imbalance in the training set.

Implementation: For csGAN, we used the same network structure and hyperparameter values (except that the number of iterations was 50,000) as the official implementation of StarGAN v2 [11]. For the regression model $f$, we used DenseNet [14] pretrained on ImageNet [15] and Adam as the optimizer with the initial learning rate set to $1\times 10^{-4}$. The batch size was set to 64. The learning was stopped by the early stopping (no decrease in validation loss for 20 epochs).

Evaluation Metric: We quantitatively evaluated the effect of C-DA by the prediction performance of the Mayo score severity classification by $f$. The prediction class (i.e., discrete Mayo level) of the images is determined by quantizing the model outputs into these neighboring discrete levels (e.g., 1.3 $\rightarrow$ 1). Since the dataset is substantially imbalanced in the number of images in each class, we mainly used the F1 score for the performance evaluation.

Comparative Methods: We compared the performance of the proposed DA method (C-DA) with three comparative methods: 1) Baseline, which is conventional regression, 2) Classical DA, which is the baseline with DA by a random combination of horizontal/vertical flipping and rotation, and 3) GAN-based DA, which is used to generate images by the original implementation of a cGAN, called StyleGAN2-ADA [16]. For 2) and 3), we will show the results with 5,000 generated images per class (i.e., 20,000 in total) because their validation F1 score was saturated even though we used more generated images.

In addition, as an ablation study of C-DA, we evaluated the classification performance of a method that uses the original images $\Omega$ as $\Omega^{\prime}$ (C-DA w/o GAN). We also performed C-DA under different severity intervals. Specifically, we examined $\epsilon=1,0.5,$ and $0.25$ to generate 4, 7, and 13 images from a single $z$, respectively. We used 250 randomly selected $z$s and thus generated 1,000, 1,750, and 3,250 images for each $\epsilon$. Note that the validation F1 scores were almost saturated at 250 $z$s; this means that ours show faster saturations than the above conventional methods, which need 20,000 images ($>3,250$) to saturate.

4.2 Qualitative Evaluation of Generation Images

Fig. 4 shows the generated images with and without order loss by csGAN. Each image was generated with $\epsilon=0.5$ from an original image at Mayo 0. With the order loss, the severity shifts smoothly between the generated images as the erythema becomes intense, and the semilunar folds gradually disappear as the severity increases. In contrast, without the order loss, the image generated at Mayo 1.5 shows large noises, and the severity is unclear. This observation confirms that the order loss has a stabilization effect of generating images at real-valued levels.

4.3 Classification Performance

Table 1 shows the classification performance of each method. Baseline and conventional DA methods had similar F1 scores, while C-DA ($\epsilon=0.5$) had a higher F1 score than the three comparison methods. These results indicate that the generated images with real-valued severity levels are more effective than conventional DAs. The following facts also confirm this effect. First, the F1 score of C-DA ($\epsilon=0.5$) was higher than that of C-DA w/o GAN. Second, F1 scores of C-DA ($\epsilon=0.5$) were even higher than that of C-DA ($\epsilon=1$).

On the other hand, the results also show that image generation at $\epsilon=0.25$ is not very effective. As we noted before, the real-valued levels of the generated images are not completely reliable. Therefore, when $\epsilon$ becomes smaller, the difference between the neighboring levels (e.g., 0.25 and 0.5) becomes unreliable even as the relative levels. This fact indicates a limitation in generating images at real-valued levels, and at the same time, it proves the validity of our dataset-disjoint optimization strategy.

Fig. 4 shows box plots of the model output $f$ for test images of each Mayo score. Here, (a) is the regression (Baseline) and (b) the proposed DA (C-DA($\epsilon=0.5$)). The horizontal and vertical axes correspond to the correct Mayo score and the model outputs, respectively. The overall model outputs of C-DA are a narrower interquartile range for each Mayo score than Baseline. Especially, the overlap between the interquartile ranges of Mayo 2 and Mayo 3 is decreased. Consequently, C-DA had better classification performance, even for minor classes with fewer images.

5 Conclusion

We proposed a continuous data augmentation (DA) scheme comprising two techniques: continuous severity GAN (csGAN) to generate medical images with real-valued severities and dataset-disjoint multi-objective optimization to utilize the generated images. Through qualitative and quantitative evaluations on an endoscopic ulcerative colitis (UC) image dataset, we confirmed that our DA scheme achieves higher F1 scores by utilizing appropriately generated images.

The current limitations of this work are as follows. First, our method is applicable to various tasks with real-valued conditions, and therefore we need to examine our method with other datasets. Second, our UC datasets only have discrete levels and thus could make our quantitative evaluation in a discrete manner. We will examine different performance evaluations if we find a medical dataset with reliable real-valued annotations.

6 Compliance with Ethical Standards

This study was performed in line with the principles of the Declaration of Helsinki. Ethical approval for this study was granted by the Ethics Committee of the Kyoto Second Red Cross Hospital.

7 Acknowledgments

This work was supported by JSPS KAKENHI, JP21K18312, and JST SPRING, Grant Number JPMJSP2136.