Semantic Segmentation of Underwater Images

Dataset and Preprocessing

• —SUIM dataset, which provides pixel-level annotations for 10 class labels (IDs 0 to 9, including background)
• —Images and masks resized to 320×240 pixels
• —Pixel values scaled to the range [0, 1]
• —RGB masks converted to pixel-wise class indices
• —Random horizontal flips with a 50% chance as data augmentation

Classes

• —0: Background (black)
• —1: Human divers (blue)
• —2: Aquatic plants (green)
• —3: Wrecks/ruins (gray)
• —4: Robots/ROVs (cyan)
• —5: Fish (yellow)
• —6: Crustaceans (red)
• —7: Mollusks (magenta)
• —8: Coral (orange)
• —9: Other animals (white)

The colors are the mask colors used for each class.

Model Architecture

• —U-Net encoder-decoder with attention blocks, to improve focus on important features in very noisy images
• —ConvBlock: Conv2d, BatchNorm, and ReLU, repeated twice for deep feature extraction
• —MaxPooling reduces the spatial dimensions of the feature maps in the encoder
• —Attention blocks let the decoder focus on important encoder features by suppressing irrelevant background regions
• —ConvTranspose2d layers upsample the feature maps in the decoder, where they are combined with encoder features to reconstruct the segmentation map
• —One final convolution layer produces the segmentation output
• —95 layers in total: 34 in the encoder, 32 in the attention blocks, 28 in the decoder, and 1 output layer

Loss Function and Metrics

• —Cross-entropy loss: the model outputs per-pixel class scores (logits), and the loss applies softmax and penalizes incorrect class probabilities during training
• —Pixel accuracy: the fraction of correctly classified pixels
• —Dice coefficient: measures the overlap between predicted and true masks, and is less sensitive to class imbalance than pixel accuracy. Dice = (2 × intersection + smooth) / (prediction sum + label sum + smooth), where smooth is a small constant that avoids dividing by zero
• —Intersection over Union (IoU): the overlap ratio between the predicted and ground-truth regions. IoU = (intersection + smooth) / (union + smooth), where union is the total number of pixels predicted or labeled as the class

Training Setup

• —Batch size of 4 and 120 epochs
• —Adam optimizer with a learning rate of 1e-4
• —Cross-entropy loss, suited to multi-class segmentation
• —Data loader with 2 worker processes and shuffling enabled
• —Trained on a GPU when available, and on the CPU otherwise
• —Best model weights saved whenever the average training loss improves
• —Each epoch: forward pass, loss computation, backpropagation, then pixel accuracy, Dice, and IoU calculation, with running loss and accuracy displayed during training

Results

• —Loss decreased from about 0.75 to about 0.10 over the course of training
• —Pixel accuracy increased from about 30% to about 93%
• —Dice score increased in step with pixel accuracy
• —IoU grew continuously throughout training

The Dice score is reported as 1 minus the Dice loss.

Future Work

• —Train with a combined Dice loss and cross-entropy loss
• —Experiment with different preprocessing and augmentation strategies