Real-time Water Area Segmentation for USV using Enhanced U-Net

As an image segmentation network widely used in the medical community, U-Net has proved the feasibility of networks with encoding-decoding paradigms. However, U-Net is rarely used for more complex RGB images due to its limited parameters. Combined with the scenes of water area segmentation, we make a targeted modification of U-Net to make it more suitable for RGB image segmentation, with only a small increase in parameters. In addition, a general image local gradient penalty is proposed to introduce local gradient information into the loss function to balance punishment for difficult and easy example points in output masks. The penalty is helpful to improve classification performance during training. To further reduce computational burden of unmanned surface vessels, we introduce model channel pruning in a more efficient way. In particular, the channel pruning and fine-tuning are combined into a repeatable step, and multiple small pruning ratios instead of single large pruning ratio are used to obtain a series of models with different accuracy-speed trade-offs. The experimental results show the practicability and applicability of our proposed method for water area segmentation.