BiSeg-SAM: Weakly-Supervised Post-Processing Framework for Boosting Binary Segmentation in Segment Anything Models

2024-12-03

IEEE International Conference on Bioinformatics and Biomedicine

AI Summary

Pixel-level annotations are scarce in medical image binary segmentation, and generic foundation models (e.g., SAM) suffer from domain-knowledge deficiency, leading to suboptimal accuracy. Method: We propose a weakly supervised post-processing framework that integrates fine-tuned SAM with CNN-based local features, weak-supervision prompt generation, and multi-mask conversion. Contributions/Results: (1) WeakBox—first automatic bounding-box prompt generator; (2) MM2B—a mechanism converting coarse masks into reliable bounding-box prompts; (3) scale-consistency loss—mitigating multi-scale segmentation misalignment; (4) DetailRefine—a lightweight module refining boundaries using minimal ground-truth labels. Evaluated on five colon polyp and one skin lesion datasets, our method significantly outperforms state-of-the-art approaches under weak supervision, achieving high-precision segmentation with only bounding-box or scribble annotations. Code is publicly available.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Multi-instance/Multi-view Learning
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Web Mining and Content Analysis: Large pretrained models with web data
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Accurate segmentation of polyps and skin lesions is essential for diagnosing colorectal and skin cancers. While various segmentation methods for polyps and skin lesions using fully supervised deep learning techniques have been developed, the pixel-level annotation of medical images by doctors is both time-consuming and costly. Foundational vision models like the Segment Anything Model (SAM) have demonstrated superior performance; however, directly applying SAM to medical segmentation may not yield satisfactory results due to the lack of domain-specific medical knowledge. In this paper, we propose BiSeg-SAM, a SAM-guided weakly supervised prompting and boundary refinement network for the segmentation of polyps and skin lesions. Specifically, we fine-tune SAM combined with a CNN module to learn local features. We introduce a WeakBox with two functions: automatically generating box prompts for the SAM model and using our proposed Multi-choice Mask-to-Box (MM2B) transformation for rough mask-to-box conversion, addressing the mismatch between coarse labels and precise predictions. Additionally, we apply scale consistency (SC) loss for prediction scale alignment. Our DetailRefine module enhances boundary precision and segmentation accuracy by refining coarse predictions using a limited amount of ground truth labels. This comprehensive approach enables BiSeg-SAM to achieve excellent multi-task segmentation performance. Our method demonstrates significant superiority over state-of-the-art (SOTA) methods when tested on five polyp datasets and one skin cancer dataset. The code for this work is open-sourced and available at https://github.com/suencgo/BiSeg-SAM.

Problem

• Enhancing binary segmentation accuracy for medical images
• Reducing reliance on costly pixel-level annotations
• Integrating SAM with domain-specific medical knowledge

Innovation

• Fine-tune SAM with CNN for local features
• WeakBox auto-generates prompts and MM2B
• DetailRefine module enhances boundary precision