MAS-SAM: Segment Any Marine Animal with Aggregated Features
Tianyu Yan Zifu Wan Xinhao Deng Pingping Zhang Yang Liu Huchuan Lu


Two one-parameter families of nonconforming enrichments of the Crouzeix–Raviart finite element
    [
    April 29, 2024
==============================================================================================




Recently, Segment Anything Model (SAM) shows exceptional performance in generating high-quality object masks and achieving zero-shot image segmentation. However, as a versatile vision model, SAM is primarily trained with large-scale natural light images. In underwater scenes, it exhibits substantial performance degradation due to the light scattering and absorption. Meanwhile, the simplicity of the SAM's decoder might lead to the loss of fine-grained object details. To address the above issues, we propose a novel feature learning framework named MAS-SAM for marine animal segmentation, which involves integrating effective adapters into the SAM's encoder and constructing a pyramidal decoder.

More specifically, we first build a new SAM's encoder with effective adapters for underwater scenes. Then, we introduce a Hypermap Extraction Module (HEM) to generate multi-scale features for a
comprehensive guidance. Finally, we propose a Progressive Prediction Decoder (PPD) to aggregate the multi-scale features and predict the final segmentation results. When grafting with the Fusion Attention Module (FAM), our method enables to extract richer marine information from global contextual cues to fine-grained local details. Extensive experiments on four public MAS datasets demonstrate that our MAS-SAM can obtain better results than other typical segmentation methods.

The source code is available at https://github.com/Drchip61/MAS-SAM.



§ INTRODUCTION

Marine Animal Segmentation (MAS) is a critical and fundamental task in the field of visual intelligence and underwater robotics <cit.>. It aims at identifying and segmenting marine animals from underwater images or videos. Functionally, the accurate segmentation of marine animals is of great importance for various research areas, including marine biology, ecology and conservation. However, MAS presents unique challenges compared to typical terrestrial image segmentation <cit.>. In fact, underwater environments are characterized by complex light scattering and absorption effects, leading to degraded image quality, reduced contrast and blurred objects. In addition, marine animals often exhibit camouflage properties, which has further complicated the segmentation task. To address these challenges, advanced perception techniques are required.

With deep Convolutional Neural Networks (CNN), many methods <cit.> achieve significant improvements in the segmentation accuracy. However, the complex environment demands a more comprehensive understanding of underwater images to achieve accurate MAS. Meanwhile, CNN's local receptive fields are not well-suited to capture long-range dependencies and contextual information. Thus, it struggles to discern marine animals in underwater scenes.

Recently, Segment Anything Model (SAM) <cit.> is proposed and has shown great potential in general segmentation tasks.

However, SAM's training scenarios primarily involve natural lighting conditions, which limits its performance in underwater environments. Besides, the simplistic decoder structure in SAM lacks the capability for generating fine-grained segmentation results.

Considering the above facts, in this work we propose a novel SAM-based feature learning framework named MAS-SAM for marine animal segmentation.

More specifically, by freezing the pre-trained parameters of the SAM’s encoder and introducing effective adapters, we build an Adapter-informed SAM Encoder (ASE) for extracting the unique features from marine animal images.

In addition, we construct a Hypermap Extraction Module (HEM) to extract multi-scale feature maps from the new SAM's encoder.

It serves as a comprehensive guidance for the subsequent mask prediction process.

To improve the SAM's decoder, we introduce a Progressive Prediction Decoder (PPD) to aggregate features from the original prompt, ASE and HEM.

When grafting with the Fusion Attention Module (FAM), our PPD can prioritize the importance of multi-grained feature maps and extract richer marine information from global contextual
cues to fine-grained local details.

Extensive experiments on four MAS datasets show that our method can consistently obtain outstanding results.

In summary, our contributions are as follows:


  * We introduce a novel feature learning framework named NAS-SAM for Marine Animal Segmentation (MAS).
    It can custom SAM with aggregated multi-scale features for high-performance MAS.

  * We propose a Hypermap Extraction Module (HEM), which generates multi-scale features based on the SAM's encoder. It serves as a comprehensive guidance for the subsequent mask prediction process.

  * We propose a Progressive Prediction Decoder (PPD) to effectively improve the representation ability of SAM's decoder, capturing a wide range of marine information from global contextual cues to fine-grained local details.

  * We perform extensive experiments to verify the effectiveness of the proposed modules. Our method achieves state-of-the-art performances on four MAS datasets.



§ RELATED WORK




 §.§ Marine Animal Segmentation

MAS is a specific and challenging sub-domain of image segmentation focused on marine animals.

With the advance of deep learning, CNNs have become a popular choice for MAS.

For example, Li et al. <cit.> propose an Enhanced Cascade Decoder Network (ECDNet) for accurately segmenting marine animals from complex underwater environments.

Chen et al. <cit.> propose to reduce the impact of water degradation diversity via a Siamese network.

Recently, Cheng et al. <cit.> built a bidirectional collaborative mentoring network for leveraging structural texture and contextual clues of marine images.

Fu et al. <cit.> design a fusion network to learn semantic features of camouflage animals.

However, the lack of global understanding ability limits the performance of these CNN-based methods.

Meanwhile, vision Transformers <cit.> have shown promising capabilities in capturing long-range dependencies and global contextual information.

They deliver much better performance in general image segmentation tasks <cit.>.

Inspired by these works, Hong et al. <cit.> propose a hybrid network for underwater salient object segmentation, which jointly leverages the advantage of CNNs and Transformers.

However, Transformer-based methods still require massive training data to achieve satisfactory results.

Different from previous works, our method resorts to the vision fundamental model to achieve better MAS results in challenging scenes.


 §.§ SAM for Customized Tasks

Recently, SAM <cit.> has drawn great attention for its powerful image understanding abilities.

With appropriate prompts, it enables transfer learning across various image segmentation tasks <cit.>.

It also emerges the ability of few-shot/zero-shot segmentation, eliminating the need for task-specific retraining.



However, SAM is trained on large-scale natural images.

Thus, it fails to adequately capture the unique characteristics of other scenes and can hardly achieve outstanding performance on customized tasks.

Meanwhile, due to the decoder's simplicity, SAM only utilizes the terminal features of vision Transformers and might lead to the loss of fine-grained object details for accurate segmentation.

To relieve the above drawbacks, some works <cit.> have integrated domain-specific cues into SAM through the utilization of simple adapters.

Furthermore, Lai et al. <cit.> verify that sophisticated adapters can highlight task-specific information.

However, the exploration of applying SAM to MAS is rather limited.

Xu et al. <cit.> directly fine-tune SAM for underwater image segmentation.

Zhang et al. <cit.> propose a dual-SAM structure with auto-prompts for MAS.

More importantly, previous modifications on the SAM's decoder can not capture intricate details and structures of marine animals.

Therefore, in this work we step further and custom SAM with aggregated multi-scale features for high-performance MAS.





§ PROPOSED METHOD

As shown in Fig. <ref>, our proposed method includes three main components: Adapter-informed SAM Encoder (ASE), Hypermap Extraction Module (HEM) and Progressive Prediction Decoder (PPD).

These key components will be elaborated in the following sections.




 §.§ Adapter-informed SAM Encoder

As stated in previous sections, although SAM has shown great potential for universal segmentation tasks, it is suboptimal to directly deploy the pre-trained SAM to MAS.

To address this issue, we propose an Adapter-informed SAM Encoder (ASE), which customizes SAM for marine animal images.

As shown in Fig. <ref>, we retain the core components of the original SAM and utilize two parameter-efficient fine-tuning mechanisms for improving the pre-trained encoder.

As shown in Fig. <ref>, we incorporate the LoRA <cit.> and Adapter  <cit.> into the Multi-Head Self-Attention (MHSA) and Feed-Forward
Network (FFN) of each Transformer block, respectively.

More specifically, let X_i∈ℝ^N× D be the input of the i-th Transformer block.

Here, N is the number of tokens and D denotes the embedding dimension.

The MHSA layer modified by LoRA can be represented as follows:

    Q_i = W_q(X_i)+W^up_q(W^down_q(X_i),




    K_i = W_k(X_i),




    V_i = W_v(X_i)+W^up_v(W^down_v(X_i),




    X_i =MHSA ( Q_i,K_i,V_i ) +  X_i,

where W_q, W_k and W_v are weights of three linear projection layers to generate the original Query, Key and Value matrices, respectively.

W_q,v^down∈ℝ^M× D and W_q,v^up∈ℝ^M× D are weights of two linear projection layers to reduce and restore the feature dimension, respectively.

M is the down-projection dimension.

In this way, we can freeze the pre-trained weights (W_q, W_k and W_v) and utilize rank decomposition matrices to greatly reduce the number of trainable parameters.

Furthermore, we insert an Adapter into the FFN as:

    X_i= MLP (LN ( X_i )    ),


    X_i+1=W_adpt^up ( σ ( W_adpt^down ( X_i )   )    )+X _i,

where LN and MLP stand for the Layer Normalization (LN) <cit.> and Multilayer Perceptron (MLP), respectively.

σ is the Rectified Linear Unit (ReLU).

W_adpt^down∈ℝ^P× D and W_adpt^up∈R^P× D are weights of two linear projections to reduce and restore the feature dimension, respectively.

P is the down-projection dimension.

Similar to LoRA, by employing an extremely low value of the parameter P, we can achieve a parameter-efficient fine-tuning for adapting the pre-trained SAM's encoder to marine scenes.


 §.§ Hypermap Extraction Module

Due to the complex underwater environment, it is of vital needs to exploit both local details and global contexts for robust and accurate MAS.

As previous works point out, different Transformer layers capture different-level semantics <cit.>.

Generally, shallow layers retain more local details and deep layers express more contextual
information.

Therefore, to enable our proposed model leverage richer marine information, we propose a Hypermap Extraction Module (HEM) to consider multi-scale feature maps from ASE.

It then serves as a comprehensive guidance for the subsequent mask prediction process.

More specifically, we first feed an image I∈ℝ^H× W× 3 into ASE and obtain the outputs of different Transformer layers.

In this work, we select the 3-6-9-12 layers and get the multi-scale token features, i.e., X_i(i=3,6,9,12).

Then, we reshape them to spatial feature maps F_i∈ℝ^H/16× W/16× D.

To simultaneously consider these multi-scale feature maps, we perform the following feature aggregation,

    R_i = ϕ_1×1(F_i),


    H_1= ϕ_3×3 [R_3 ,R_6,R_9,R_12 ],

where ϕ_1×1 and ϕ_3×3 are convolutional layers with 1×1 and 3×3 kernels, respectively.
To improve the training stability, a Batch Normalization (BN) and a ReLU activation function are also introduced after the convolutional layers.

[·] is the concatenation in channel.

Afterwards, we introduce a channel-attention layer to generate the hypermap H_j as follows:

    H_1= H_1×δ (ϕ_1 × 1(GAP(H_1)))+H_1,


    H_j+1=ϕ_3×3(ψ_2×2(H_j)), j=1,2,3,

where GAP is the Global Average Pooling (GAP), δ is the Sigmoid function, and ψ_2×2 is a deconvolutional layer with a 2×2 kernel.

As shown in the top-left part of Fig. <ref>, we can obtain multi-scale hypermaps.

These hypermaps play a crucial role in improving the performance of MAS.


 §.§ Progressive Prediction Decoder

Due to the significant appearance variations of marine animals, the simple decoder design in SAM struggles to achieve accurate segmentation masks.

To this end, inspired by <cit.>, we propose a Progressive Prediction Decoder (PPD) to effectively improve the prediction ability.

As shown in the right part of Fig. <ref>, it has a pyramidal structure to progressively aggregate multi-source features from the original prompt, ASE and HEM, and obtain the final segmentation predictions.


As shown in Fig. <ref>, we propose a Fusion Attention Module (FAM) to fully aggregate the multi-source features.

More specifically, we begin by up-sampling the features from ASE and resizing the input features into the same size.

Then, we fuse them as follows:

    U_i=φ(F_i),


    D_j+1=FAM(U_i, D_j,H_j)),

where U_i is the up-sampled feature by utilizing a bilinear interpolation φ.

D_j is the output of the j-th pyramidal stage in the proposed PPD.

For the FAM, we utilize a channel-attention to prioritize the importance of multi-source features.

Residual structures are also deployed to enhance the representation ability.

The procedure can be formulated as:


    Y_j=ϕ_1×1([U_i,D_j,H_j]),


    W_j=δ(ϕ_1 × 1(GAP(Y_j))+ϕ_1 × 1(GMP(Y_j))),


    D_j+1=W_j(Y_j)+Y_j,


    D_j+1=ϕ_1×1(D_j+1),

where GMP is the Global Max Pooling (GMP).

The channel-wise weights can highlight relevant features and suppress irrelevant ones <cit.>.

Meanwhile, the attention mechanism employed by our FAM helps to capture intricate relationships between features across various scales, resulting in more coherent and informative feature representations.

Thus, FAM can effectively integrate and refine the multi-source features.

Finally, to achieve the progressive prediction, we build the PPD grafted with the FAM as follows:


    P_j= ϕ_1×1(φ(FAM(U_i, D_j,H_j))),

where P_j is the prediction mask of the j-th pyramidal stage.

The proposed PPD facilitates the seamless aggregation of multi-source features from the original
prompt, ASE and HEM, resulting in richer marine information from global contextual cues to fine-grained local details.

To further improve the prediction results, we take all the predictions at different stages and generate the final prediction as follows:

    P= ϕ_1×1([P_1,P_2,P_3,P_4]).

With the synergistic use of the pyramid structure and FAM, our MAS-SAM can effectively leverage the diverse information and produce highly refined and detailed segmentation masks for a wide range of marine animal shapes and sizes.


 §.§ Model Training

During training, we follow previous methods <cit.> and use deep supervision from three levels, i.e., pixel-level supervision (Binary Cross-Entropy Loss), region-level supervision (SSIM Loss) and overall-level supervision (IoU Loss).

Thus, we define ℒ^f or ℒ^j as a combined loss with three terms:

    ℒ^f/j = ℒ_BCE+ℒ_SSIM+ℒ_IoU,

where ℒ^f and ℒ^j are losses of the final prediction and the j-th stage-output, respectively.

Due to the space limitation, we refer readers to <cit.> for the exact formulas.



§ EXPERIMENTS



 §.§ Datasets and Evaluation Metrics

In this work, we adopt four public MAS benchmarks to evaluate the model performance.

The MAS3K dataset <cit.> comprises 3,103 marine images, in which 193 are background images. Following the default split, we use 1,769 images for training and 1,141 images for testing.

The RMAS dataset <cit.> consists of 3,014 marine animal images with 2,514 images used for training and 500 images for testing.

The UFO120 dataset <cit.> comprises 1,620 underwater images with various scenes. Following the default split, we use 1,500 images for training and 120 images for testing.

The RUWI dataset <cit.> is a real underwater image dataset captured under complex light conditions. The dataset consists of 700 images, splitting into 525 images for training and 175 images for testing.

Meanwhile, we adopt five metrics to evaluate the segmentation performance of different models.

They are Mean Intersection over Union (mIoU), Structural Similarity Measure (S_α) <cit.>, Weighted F-measure (F_β^w) <cit.>, Mean Enhanced-Alignment Measure (mE_ϕ) <cit.> and mean absolute error (MAE).


 §.§ Implementation Details

Our model is implemented with the PyTorch toolbox and one RTX 3090 GPU.

For our model, the SAM's encoder and prompt encoder are initialized from the pre-trained SAM-B <cit.>, while other proposed modules are randomly initialized.

During the training process, we freeze the SAM's encoder and only fine-tune the remaining modules.

The widely-used AdamW optimizer <cit.> is adopted to update the model parameters.

The initial learning rate and weight decay are set to 0.001 and 0.1, respectively.

We reduce the learning rate by a factor of 10 at every 20 epochs.

The total number of training epochs is set to 50.

The mini-batch size is set to 8.

The input images are uniformly resized to 512×512×3.





 §.§ Comparison with State-of-the-arts

In this section, we compare our method with other approaches on four public MAS datasets.

Specifically, we follow previous works, i.e., MASNet <cit.> and ECDNet <cit.>, to include compared methods/tools.

MASNet and ECDNet are typical MAS methods.

Other compared methods are modified for underwater scene segmentation.

Especially, some methods (e.g., SINet <cit.>, C2FNet <cit.>) focus on camouflage object detection for underwater images.

Thus, it is very reasonable to compare with them.

I-MedSAM <cit.>, MedSAM <cit.>, SAM-Adapter <cit.> and SAM-DADF <cit.> are similar to our method in addressing the adaptability issue of SAM for downstream tasks.

Quantitative Comparisons.
Tab. <ref> and Tab. <ref> illustrate the quantitative results of compared methods.

They clearly demonstrate that our method outperforms other methods on all five metrics across four datasets.

These results show the superiority of our method in terms of overall completeness, structural accuracy, and pixel-wise precision.

Compared with CNN-based methods, our method delivers significant advantages in multi-scale and long-range modeling capabilities.

On the large-scale MAS3K dataset, our method achieves the best mIoU, S_α , F_β^w and mE_ϕ values.

It demonstrates an improvement of approximately 3-4% across various metrics.

Compared with Transformer-based methods, our method shows a 2-3% gain across multiple metrics on the MAS3K dataset.

Moreover, significant improvements are also observed on other three datasets.

Compared with SAM-based methods, our method efficiently injects specific marine scene information into SAM through two concise adapter modules.

In short, our design notably enhances the model's segmentation ability and achieves a considerable performance gain over other approaches.

Qualitative Comparisons.
To demonstrate the superiority of our method more intuitively, we present visual comparisons of different methods in Fig. <ref>.

The visualizations provide a clearer advantage of our method compared with the previous approaches.

There are challenging images with high background clutters and abundant details.

However, our method still generates better segmentation masks.




 §.§ Ablation Studies

On the MAS3K dataset, we conduct experiments to verify the effectiveness of each module in our model.

More results are in the supplementary material.

The model (A) is directly deploying the pre-trained SAM.

Visual results are in Fig. <ref>.




Effects of Adapters in ASE.
In the 1-3 rows of Tab. <ref>, we show the impact of adopters in the ASE. By incorporating two effective adapters, our method has a significant improvement in terms of all evaluation metrics.

These results indicate that it is very necessary and powerful to inject marine environment information into the fundamental model.

Effects of Key Modules.
In the 3-6 rows of Tab. <ref>, we examine the impact of other key modules.

By constructing a PPD, the model can capture more local details, resulting in a 0.8% increase in mIoU.

Subsequently, HEM enables the full utilization of multi-level features extracted by the ASE, serving as informative guidance for the subsequent prediction structure. This enhancement leads to a noticeable improvement in performance, better than typical FPN structures. Lastly, we leverage the fused information through the FAM. Compared with the original SAM decoder, our method achieves great improvements in all metrics.

Effects of Different Losses.
In the 6-7 rows of Tab. <ref>, we demonstrate the significance of utilizing multi-level supervision during the training process. By employing three levels of supervision, our method shows a notable improvement.



 §.§ More Discussions

In this work, our aim is to adapt SAM to MAS.

However, there are three key issues. 1) Domain differences. SAM is pre-trained on natural light scenes, while
MAS always suffers from the refraction and absorption of light.

Thus, directly employing SAM is hard to achieve good performance for MAS.

2) Camouflage behaviours. Marine animals usually exhibit camouflage behaviours, and large variations in shapes and sizes.

The original SAM is not good at modeling these appearances.

3) Missing details. With cross-attentions and only two deconvolutions, the original SAM’s decoder leads to a substantial loss of object details.

To address above issues, we present three contributions as follows:
1) We inject marine domain knowledge into the SAM backbone through efficient adapters, thus making SAM more adaptable to marine scene tasks.

2) We employ a HEM to fully leverage multi-scale information.

This helps our model to accurately locate marine animals against complex appearances.

3) We propose a PPD structure to progressively integrate feature maps from different levels.

With three-grained supervision losses, it can capture the intricate details of marine life extensively.

The above facts make our work valuable for underwater intelligence.

Meanwhile, our work is not limited to the MAS task and can be easily translated into many binary segmentation tasks, such as salient object detection, polyp segmentation, text segmentation, image anomaly detection, etc., showing a strong transferability.






§ CONCLUSION

In this work, we propose a novel feature learning framework named MAS-SAM for MAS.

Specifically, we first propose an Adapter-informed SAM Encoder (ASE) for underwater scenes.

Then, we introduce a Hypermap Extraction Module (HEM) to generate multi-scale features for a
comprehensive guidance.

Finally, we propose a Progressive Prediction Decoder (PPD) to aggregate the multi-scale features and predict the final segmentation results.

When grafting with the Fusion Attention Module (FAM), our approach extracts richer marine information from global contextual cues to fine-grained local details.

Extensive experiments on four MAS datasets show the effectiveness of our MAS-SAM.


Acknowledgements.
This work was supported in part by the National Natural Science Foundation of China (No.62101092), the Fundamental Research Funds for the Central Universities (No.DUT23YG232) and the Open Project Program of State Key Laboratory of Virtual Reality Technology and Systems, Beihang University (No.VRLAB2022C02).



named


