GAN based Unsupervised Segmentation:
Should We Match the Exact Number of Objects

The CycleGAN [7] is used to generate our synthetic training data. As the standard CycleGAN implementation, generators and discriminators are used to transfer the styles between two image modalities. The role of the generators is to convert the real images to another domain, which are typically called ”fake” images. The discriminators then judge if a given image is real or fake. The CycleGAN model design creatively forms the entire learning process as a cycle-consistent loop, where the reconstructed fake images after two generators should be close to the original real images. In each branch, the generator tries to generate realistic images, while the discriminator try to distinguish the fake images from the real ones.

In our unsupervised image segmentation framework (Fig. 1), the CycleGAN is employed to synthesize segmentation masks (or called ”annotations”) from the real images $I$, and synthesize realistic looking images from simulated segmentation masks $M$. In the ideal case, the trained generator $G_{I->M}$ can be directly used as a segmentation network to segment new images. However, the quality of synthesis between real images and clean binary masks is typically not optimal since the underlying Poisson distribution of the binary masks is not a realistic distribution in real images [13]. Moreover, the optimization of the KL divergence for training discriminators is more difficult to converge [14] using clean binary masks. Therefore, the Gaussian smoothing, random noise, and brightness variations are used to generate augmented masks $M_{A}$ in additional to the simulated mask images $M$ for better synthetic performance (Fig. 2). Then, the trained generator $G_{M_{A}->I}$ will provide us unlimited fake but realistic looking images $I_{F}$ from the simulated and augmented masks $M_{A}$. Eventually, the fake images $I_{F}$ and the clean binary masks $M_{A}$ (before augmentation) are used to train another independent segmentation network (see Image Segmentation section).

Compared with traditional pixel-to-pixel conditional GAN design which needs pixel-level correspondence between two modalities, the CycleGAN does not need paired images for training. However, the superior synthetic segmentation performance is typically achieved if the distributions of the number of objects in real image modality and annotation modality are roughly matched [13, 14], named as ”macro-level” matching. However, no studies have explored the level of matching in the middle of pixel-level and macro-level. In this study, we proposed the idea called ”micro-level” matching, which matches the number of objects in each mini-batch (Fig. 1). For example, if a real image has roughly 21 microvilli, we will provide a simulated mask with the same 21 sticks, when forming the mini-batch. Then, the next question is that how can we get rough number of objects from the real images. In this study, we utilize the multi-channel nature of fluorescence microscopy to split the microvilli marker mCherry-Espin (magenta color objects) and microvilli tip marker EGFP-EPS8 (green color objects). Using the simple intensity thresholding based cell counting algorithm [19], the rough number of protein objects are easily achieved. The numbers are then used as the rough number of microvilli to simulate the corresponding mask files with the same number of objects, as the micro-level matching. Note that we only match the number of the objects in the micro-level matching, where the spatial distribution of the objects is still random.

U-Net [20] is employed as the segmentation backbone network, which is a fully convolutional neural network and widely used in image segmentation tasks. The segmentation part of our framework is shown in Fig. 3. In our proposed unsupervised segmentation framework, the input images of U-Net are the fake microvilli images, which are generated from the simulated masks using $G_{M->I}$ or $G_{M_{A}->I}$ from trained CycleGAN. Then the Dice loss function is calculated by comparing the predicted segmentation with the simulated binary masks. Traditional deep neural network typically needs a large number of annotated images to train a segmentation network. Using our design, however, we can generate unlimited number of training data to train the segmentation network without any manual annotation efforts.

Twelve microvilli images acquired using fluorescence microscopy were used as training data, where each image had $\approx 900\times 900$ pixels with pixel resolution 1.1 $\mu$m. Then, 500 image patches with $128\times 128$ pixels were randomly sampled from the twelve images to train the CycleGAN as the real images. Then, another independent microvilli video with 20 frames was used as testing data to evaluate the performance of the proposed unsupervised segmentation methods. Each frame has $256\times 256$ pixels with pixel resolution 1.1 $\mu$m. All microvilli in each frame were densely annotated manually by an experienced biologist as the gold standard segmentation. The video of microvilli frames and manual annotations are presented in the supplementary materials: https://github.com/iamliuquan/GAN_based_segmentation.

In order to test if micro-level matching can improve the unsupervised segmentation performance, we performed experiments using both macro-level and micro-level matching. For micro-level matching, the number of sticks of each images was obtained by automatically counting the number of green proteins. For macro-level matching, the number of sticks for each image was randomly sampled from a uniform distribution (range from 11 to 63), according to distribution of proteins.

As shown in Fig 2, we have four different augmentation settings to generate the simulated masks in annotation domain:
Binary masks: The binary masks were directly simulated as the images in the annotation domain, without any augmentation. Based on [17] and the prior biological knowledge of microvilli, the width of each microvilli was simulated between 2 to 5 $\mu$m, while the length was simulated between 10 to 50 $\mu$m. As the pixel resolution of all our images was 1.1 $\mu$m, we randomly generated sticks with 2 to 4 pixels width and 9 to 45 pixels length from uniform distribution.
Gaussian Smoothing: The first augmentation was Gasussian smoothing, where a Gaussian filter with kernal size of $5\times 5$ was applied to the binary masks.
Random noise: Upon the Gaussian smoothing, the random Gaussian noise was further applied to the entire mask image. The values of random noise ranged from 0 to 255 follows Gaussian distribution.
Different brightness: To further introduce the global intensity variations, random intensity values (200 to 255) were assigned to each stick in binary masks, where maximum foreground intensity value was 255.

To improve the segmentation performance, CycleGAN was employed to synthesis cell images for U-Net model training. In our experiment, CycleGAN is used to learn the mapping from simulated masks to real microvilli cell images. We built up dataset in these two domains as CycleGAN model’s input. Our CycleGAN model was trained for 60 epochs. According to the training loss, generator trained for 50 epochs shows the best performance. Generator trained in CycleGAN will be used to synthesis microvilli cell images based on simulated mask images.

CycleGAN model cannot cover all details using original frames as input which has too many cells. For both CycleGAN and U-Net, the input images were all with $128\times 128$ resolution cropped from original frames, and then resized to $256\times 256$ during training. When applying trained U-Net on testing microvilli images, each testing image was first split to four $128\times 128$ images, and the final segmentation were achieved by concatenating the corresponding four predictions back to the original resolution. The Dice results were calculated in the original $256\times 256$ resolution for testing images. The CycleGAN and U-Net were deployed on a computer with GeForce GTX 1060 Graphic Card with 6 GB memory. To get better synthesised data and avoid over-fitting, the CycleGAN was trained with 50 epochs and the U-Net was trained with 10 epochs for all experiments. According to the prediction performance, U-Net has the best performance after 10 epochs. The results from the last epochs were reported in this paper.