Underwater Image Semantic Segmentation Using a U-Net Deep Learning Framework

Abstract

Underwater image semantic segmentation has significant applications in marine robotics, ecological monitoring, and inspection of underwater infrastructure, but underwater images are difficult due to color attenuation, light scattering, turbidity and low contrast. We propose a compact U-Net encoder-decoder as a clear and repeatable baseline for underwater semantic segmentation. The dataset was first converted to 1,525 image-mask pairs, then a four-class experimental formulation was created by mapping RGB mask colors to integer class labels. Images were resized to 256 × 256 pixels, images were normalized to [0,1] and nearest-neighbor interpolation was used to maintain the discrete class labels of masks. With a split of approximately 80:20, the dataset was split into 1220 pairs of training data and 305 pairs of validation data. The proposed U-Net has about 1.95 million trainable parameters and comprises four encoder blocks with filter numbers of 16, 32, 64 and 128, a 256-filter bottleneck, four decoder blocks that introduce skip connections, and a final 1 × 1 convolution with softmax activation for the prediction of four classes. The model was built with TensorFlow/Keras, trained with the Adam optimizer with a learning rate of 0.0001, a loss of 'sparse categorical cross-entropy', a batch size of 8, and a max of 50 epochs using standard training callback functions. The model achieved around 82% validation accuracy of pixels and a Mean IoU of 0.8056. Such a number as 0.8056 was also reported for Mean Dice; it is the same as Mean IoU and thus needs to be verified before it can be considered to be a confirmed result. A relatively accurate segmentation was obtained in a diver scene and a high over-prediction of background in a low-contrast coral/rock scene using qualitative analysis. The results show that a small U-Net can be used as a baseline model for semantic segmentation in underwater settings, and that it is important to carefully consider the preprocessing, metrics, and assessment of the model in challenging visual conditions that are not always present. Future work should explore per-class performance, data augmentation, larger and more diverse data sets, independent test sets, and compare the performance of the segmentation architectures with attention-based and transformer-based architectures.