Edema Estimation from Facial Images Taken
Before and After Dialysis via Contrastive
Multi-Patient Pre-Training

Traditional methods for measuring peripheral edema use a pitting test [25] or a medical tape to measure ankle circumference [26]. Both of these methods have significant intra- and inter-measurer variability [27] and require a medical staff to take the measurements. Another traditional method is measuring limb volume using water displacement [28], which can measure the total foot volume and thus edema. Although water displacement volumetry is accurate and sensitive [27], it takes a lot of time (25 minutes per foot [18]) and labor for the patient. Some recent work [18, 29, 30, 31] has utilized 3D imaging to measure foot volume as alternative technologies to water displacement. However, they require expensive scanners [29], a 3D camera installed on the wall [18], and scanning around the feet with a hand-held depth camera [30, 31], which are still difficult for self-management of fluid intake.

To address these issues, we aim to develop a method that helps dialysis patients easily self-manage their fluid intake. We focus on facial images taken before and after dialysis and attempt to estimate the degree of edema from the facial images. Facial images can be captured more easily by patients themselves than foot images. They can also use a tablet or smartphone to capture facial images without introducing additional special equipment. Furthermore, our image-based method can be incorporated into telemedicine, which is becoming more widespread these days.

In recent years, several studies have identified developmental disorders [32], genetic disorders [33], acute illness [34], and Alzheimer’s disease [35] from facial images. They employed deep learning approaches to identify diseases from facial images. Gurovich et al. [33] predicted hundreds of genetic syndromes from facial images by fine-tuning the CNN model pre-trained on a large-scale face identity database [36]. Forte et al. [34] trained a CNN model to classify facial images into healthy or acutely ill by using sets of simulated sick faces and synthetically generated data. These disease identifications recognize facial features such as facial dysmorphism [32, 33], specific symptoms appearing on a part of the face [34], or a specific complexion [35]. A common CNN model across all patients was constructed in these studies.

On the other hand, edema appears and decreases subtly on each patient’s face. Therefore, it is difficult for a common model across all patients to capture such subtle changes in each patient. Here, we construct a CNN model for each individual patient, rather than a common model across all patients. We pre-train a model to acquire feature representations of edema from multi-patient data before fine-tuning it to each patient.

Contrastive learning has recently become the mainstream of self-supervised [37, 38, 39, 40, 41] or supervised [42, 43] representation learning. In order to reduce manual annotations, self-supervised learning approaches produce pseudo-labels by applying specific augmentations or transformations to the input data. Subsequently, the model pre-trained or trained with the augmented/transformed data and associated pseudo-labels can be used in downstream tasks, e.g., classification and regression. A simple framework for contrastive learning of visual representations (SimCLR) [38] is one of the most popular self-supervised contrastive learning approaches. SimCLR first creates a pair of positively correlated views by applying augmentations including cropping, flipping, and color distortion to the input image. Then, both views are passed into a pair of convolutional encoders (e.g., ResNet50 [44]) and projection heads (e.g., two dense layers with ReLU activation [45]) to obtain their embedding vectors. SimCLR optimizes the whole architecture by maximizing the agreement between the positive pair of embedding vectors (i.e., augmented views of the same image) while minimizing the negative pairs (i.e., different images) in the same mini-batch. Finally, the projection heads are discarded while the convolutional encoders are kept to be utilized in downstream tasks. Momentum contrast (MoCo) [39] is another popular self-supervised contrastive learning approach. Instead of a pair of encoders as in SimCLR, MoCo uses an encoder and momentum encoder pair. The momentum encoder generates a dictionary as a queue of encoded keys by enqueuing the current mini-batch and dequeuing the oldest mini-batch. The advantage of using the momentum encoder is that the dictionary size is not restricted to the mini-batch size and can be larger. Whereas the number of negative samples available in SimCLR is limited to the mini-batch size, MoCo can cover a rich set of negative samples and learn better feature representations. Contrastive learning has been extended to supervised learning that makes use of the output labels [42, 43]. In supervised contrastive (SupCon) learning [42], the positive samples are generated not only from augmentations of the same input, but also from augmented views of other samples in the same class. While SupCon assumes the use of discrete labels as the output labels, contrastive learning utilizing continuous labels has also been proposed [43]. Dufumier et al. [43] introduced the y-Aware InfoNCE loss for 3D brain MRI classification to make feature representations of MRI images closer in similarity of patient age. The utilization of continuous labels (e.g., patient age) helps the model to learn a more generalizable representation of the data.

Our WeightSupMoCo loss leverages both discrete and continuous labels, i.e., pre- and post-dialysis labels and patient weight. Both of these two label types have important information about edema. Specifically, the pre- and post-dialysis labels provide coarse-grained information about edema while patient weight provides fine-grained information. Therefore, these two label types lead to better representation learning in the multi-patient pre-training step. Furthermore, we employ MoCo-based contrastive learning, which allows us to cover a richer set of positive and negative samples compared with the SimCLR-based SupCon [42] and y-Aware InfoNCE [43] losses. Figure 2 illustrates the differences between the SupCon, y-Aware InfoNCE, and WeightSupMoCo losses in our pre-training setting.

First, we obtain facial images of renal failure patients taken before and after dialysis. Then, we apply a facial detection method [46] to the obtained facial images and extract the centers of the facial images. Specifically, as shown in Fig. 4, we detect facial landmarks (green points) and extract the center part of the face by using four red points. By detecting facial landmarks for each patient and extracting the center part based on these landmarks, we compensate for rotation due to head or body tilt. Also, to compensate for the facial size, all extracted images are scaled to the same size. The center of the face includes the eyes and nose, where edema may appear and does not include areas unrelated to edema such as hair or clothes. Subsequently, the facial images are used in our MoCo-based contrastive learning framework (see Fig. 3. Step1). MoCo generates a query and key from the input images via the encoder and momentum encoder, respectively. Then, the encoder is trained by matching an encoded query to a dictionary of encoded keys (called a queue) by using the WeightSupMoCo loss (see III-B). The main components of our framework are the following data augmentation module, encoder network, and projection head:

Data Augmentation Module, $Aug(\cdot)$. Various combinations of data augmentations are randomly applied to each input image $\bm{x}$, as $\bm{\hat{x}}=Aug(\bm{x})$. The input images $\bm{x}$ are resized to 224$\times$224. The data augmentations consist of random horizontal flip ($p=0.5$), random color jitter ($p=0.8$, brightness $=0.4$, contrast $=0.4$, saturation $=0.4$, hue $=0.1$), and random grayscale conversion ($p=0.2$), where $p$ denotes a probability. These data augmentation settings are derived from the previous work[39, 47]. We generate two augmented views of each input image $\bm{x}$.

Encoder Network, $Enc(\cdot)$, is used for both the encoder and the momentum encoder. The encoder network maps $\bm{\hat{x}}$ to a representation vector, $\bm{r}=Enc(\bm{\hat{x}})\in\mathbb{R}^{D_{r}}$. The augmented images are separately input to the encoder and the momentum encoder, resulting in a pair of representation vectors. In our experiment, we use ResNet18 [44] as the encoder network; the dimension of representation vectors $D_{r}$ is 512. Previous studies [38, 39, 42, 43] have shown that encoder pre-trained in contrastive learning can be used for several related downstream tasks. As in the previous studies, we perform transfer learning by adding a linear layer for each downstream task to the same encoder.

Projection Head, $Proj(\cdot)$, maps the representation vector $\bm{r}$ onto an embedding vector, $\bm{z}=Proj(\bm{r})\in\mathbb{R}^{D_{z}}$. As in the previous work [38, 42], we instantiate $Proj(\cdot)$ as two dense layers with ReLU activation. The output dimension of the first layer is 512 and that of the second layer $D_{z}$ is 128. We normalize the output of the second layer to lie on the unit hypersphere, so that inner products can be used to measure cosine similarities in the embedding space. Although the projection head is an important component to project feature representations onto the embedding space, it is not required for the downstream edema estimation tasks. Therefore, we discard the projection head after the contrastive multi-patient pre-training and use only the encoder for the downstream tasks.

The embedding vector $\bm{z}$ is used to optimize the encoder and projection head. The common loss function for self-supervised contrastive learning in SimCLR and MoCo is the following InfoNCE loss [37]:

where index $i$ is called the anchor and index $p$ is called the positive (i.e., the augmented view of the anchor). Here, $A(i)$ denotes the indices of all augmented samples except anchor $i$, consisting of one positive $p$ and all other negatives (i.e., different images from the anchor). The $\cdot$ symbol denotes the inner product and $\tau$ is a temperature parameter.

The WeightSupMoCo loss aims to make the embedding vectors closer in similarity of patient weight when the pre- and post-dialysis labels are the same. This idea is based on two assumptions: (i) the pre- and post-dialysis labels indicate the presence or absence of edema, and (ii) patient weight indicates the degree of edema. These two label types thus provide coarse- and fine-grained information about edema. Let $l$ denote the pre- or post-dialysis label and $y$ denote the patient weight; the WeightSupMoCo loss is defined as follows:

where $w_{\sigma}(\cdot,\cdot)$ is a function representing the similarity of patient weights, in which we use the radial basis function (RBF) kernel, as in the y-Aware InfoNCE loss [43]. $\mbox{1}\mbox{l}_{i=j}\in\{0,1\}$ is an indicator function, which outputs 1 when $i=j$, and 0 otherwise. The hyperparameter of the RBF kernel $\sigma$ is set to 3.0 in our case, and the temperature parameter $\tau$ is set to 0.1 as in  [38, 43]. In Eq. (2), when the pre- and post-dialysis label is the same and the similarity of weights is large, the value of $w_{\sigma}(y_{i},y_{k})\cdot\mbox{1}\mbox{l}_{l_{i}=l_{k}}$ increases and thus the embedding vectors are trained to be closer. On the other hand, when the pre- and post-dialysis label is different, $\mbox{1}\mbox{l}_{l_{i}=l_{k}}$ becomes 0 and thus the distance of the embedding vectors is maximized. In our MoCo-based framework, the anchor $i$ is the set of all query and queue samples, and $A(i)$ is the set of query samples except anchor $i$ in the same mini-batch and queue samples. The dictionary size of the queue (1024 in our case) can be much larger than the mini-batch size of the query and key (16 in our case). In the queue, the encoded key samples of the current mini-batch are enqueued and those of the oldest mini-batch are dequeued. Hence, the queue can hold newer keys and remove older ones.

Although the use of queues enables the dictionary to be large, updating the momentum encoder by back-propagation is intractable because the gradient should propagate to all samples in the queue. This issue can be addressed by using the following momentum update [39]:

where $\theta_{q}$ and $\theta_{k}$ respectively denote the parameters of the encoder and momentum encoder, and $m\in[0,1)$ is the momentum coefficient. Only the $\theta_{q}$ is updated by back-propagation, while the $\theta_{k}$ evolves more smoothly than $\theta_{q}$ through the momentum update. Following the suggestion in [39], the momentum coefficient is set relatively large (0.9999 in our case) so that the $\theta_{k}$ will evolve slowly.

After the contrastive multi-patient pre-training based on WeightSupMoCo, we perform the downstream edema estimation tasks by transfer learning on per-patient data (see Fig. 3. Step2). The downstream edema estimation tasks comprise pre- and post-dialysis classification and weight prediction. In these tasks, we discard the projection head and add a linear layer to the pre-trained encoder. The projection head is used to project the embedding vector in contrastive learning, while the linear layer is used for the pre- and post-dialysis classification and weight prediction in downstream tasks, as in [42, 43]. As the linear layer, we use the classification layer with two output dimensions for the pre- and post-dialysis classification. We also use the regression layer with one output dimension for the weight prediction. Finally, the entire model including the encoder and linear layer is fine-tuned with per-patient data. As the loss functions for the fine-tuning, we use the cross-entropy loss for the pre- and post-dialysis classification and the mean squared error (MSE) loss for the weight prediction.

To collect facial images from patients before and after dialysis, University of Tsukuba recruited dialysis patients in the Takemura Medical Nephro Clinic, Japan. We explained the purpose, methods, and freedoms in participating in the study to each patient and obtained consent from all patients.

A total of 40 patients participated in the data collection, and data ranging from 1 to 10 dialysis days were obtained for each patient. The number of participation days depended on the patient’s freedom. These dialysis patients have renal failures such as diabetic nephropathy and chronic glomerulonephritis. On each dialysis day, we acquired data twice, before and after dialysis. Data were not obtained for several days after dialysis due to the patient’s physical condition. Facial videos and medical records including age, body weight (before and after dialysis), and water removal were collected from each patient. One patient was excluded from the dataset because a valid medical record could not be obtained; hence, data from 39 patients were used in the experiment. The characteristics of the dataset and dialysis patients are listed in Table I. The relationship between the number of days and the number of patients is shown in Appendix A.

Facial videos were recorded using the built-in camera of a Surface Pro 7 ²²2https://www.microsoft.com/en-us/surface/, a popular tablet device. Three minutes of video was captured, including frontal face and side views from the left and right while the patient was performing several tasks (e.g., rest, opening mouth, raising hands). In this experiment, we used only the facial videos from the frontal view at rest. Then, we randomly selected arbitrary frames from the video and obtained 100 images from each video, resulting in a total of 39,200 facial images. We excluded facial images when looking left or right based on the distance between the eye and nose landmarks (see Fig 4). We also excluded facial images when opening the mouth (e.g., while talking) based on the distance between the mouth landmarks (see Fig 4). We recorded the facial videos indoors with no direct sunlight. However, since lighting conditions could still vary from day to day, we apply data augmentation that changes the color of the image (i.e., color jitter and grayscale conversion). By changing the color of the image, we can mitigate over-fitting to lighting conditions and improve the performance of edema estimation (see Appendix C). The data of the 39 patients were divided into a contrastive multi-patient pre-training set (see Fig. 3, Step1) and downstream edema estimation set (see Fig. 3, Step2). In this experiment, data from 24 patients with less than 7 days of pre- and post-dialysis were used for the pre-training, and data from 15 patients with 7 or more days were used for the downstream edema estimation (see Appendix A).

For the multi-patient pre-training, 80% of the images were randomly selected as training data and the remaining 20% were used as validation data. For the downstream edema estimation, we performed a leave-one-day-out cross-validation, where the data from one day were used as the test data and the data from the other days were used as the training and validation data. The training and validation data ratios in this cross-validation were 80% and 20%, respectively. The method of splitting training and validation data in the multi-patient pre-training is discussed in Appendix B.

The framework was implemented with PyTorch [48] on an NVIDIA GeForce RTX 3060 GPU. During training, we used the Adam optimizer [49] with a learning rate of 1e-4 and weight decay of 5e-5 for the pre-training and a learning rate of 1e-3 for the downstream tasks. We trained for 100 epochs in the pre-training and 20 epochs in the downstream tasks. The best models using the validation data were used in the pre-training and downstream tasks. As in the data augmentation module in the pre-training (see III-A), data augmentations of random horizontal flip, random color jitter, and random grayscale conversion were used in the downstream tasks to mitigate the overfitting problem. An ablation study on data augmentation is shown in Appendix C. All images were resized to 224$\times$224, and the RGB channels were normalized using the mean and standard deviation of ImageNet [50]. In addition, for the weight prediction in the downstream task (not pre-training), each patient’s weight was normalized to mean 0 and variance 1. The predicted weights were then transformed back to the original scale using the mean and variance of the training data.

To verify the effectiveness of our WeightSupMoCo, we compared it to several methods (see Table II). In all comparative methods, we used the same CNN model (i.e., ResNet18) and data augmentation module as the proposed method. Performance comparisons using commonly used CNN networks other than ResNet18 are shown in Appendix D. To verify the effectiveness of the contrastive multi-patient pre-training, we compared our method with a model trained from scratch using only per-patient data (i.e., without multi-patient pre-training). To verify the effectiveness of using supervised labels (i.e., pre- and post-dialysis labels and patient weight), we compared our representation learning with the self-supervised learning methods, SimCLR [38], MoCo [39], and SimSiam [51]. SimSiam is one of the state-of-the-art self-supervised learning methods in computer vision. We also compared it with standard supervised pre-training methods, classification using pre- and post-dialysis labels based on the cross-entropy loss and weight prediction based on the MSE loss. In the weight prediction-based pre-training method, all patients’ weights were normalized. Furthermore, we compared our WeightSupMoCo loss with the SupCon loss [42] using pre- and post-dialysis labels and the y-Aware InfoNCE loss [43] using patient weight. In addition, to verify the effectiveness of using MoCo-based contrastive learning in our framework, we compared WeightSupMoCo loss with a SimCLR-based WeightSupCon loss (i.e., a SimCLR-based version of WeightSupMoCo). For a fair comparison, the downstream task settings were the same for all methods. In the pre-training, the learning rates of cross-entropy and MSE losses were set to 1e-3, while the other learning rates and epoch numbers were the same as in our method. The batch size was set to 16 for all methods. The temperature parameter $\tau$ in contrastive learning SimCLR, MoCo, SupCon, y-Aware InfoNCE, WeightSupCon, and WeightSupMoCo was set to 0.1. The dictionary size and momentum coefficient of MoCo and WeightSupMoCo were set to 1024 and 0.9999. The hyperparameter of the RBF kernel $\sigma$ in y-Aware InfoNCE, WeightSupCon, and WeightSupMoCo was set to 3.0.

As evaluation metrics for the pre- and post-dialysis classification, we used accuracy, area under the receiver operating characteristic curve (ROC-AUC), and area under the precision-recall curve (PR-AUC). As evaluation metrics for the weight prediction, we used mean absolute error (MAE), root mean squared error (RMSE), and correlation coefficient (CorrCoef) between the predicted and ground-truth data.