Foam Segmentation in Wastewater Treatment Plants: A Federated Learning Approach with Segment Anything Model 2

TL;DR

This paper addresses automatic foam segmentation in wastewater treatment plants under strict data privacy constraints. It combines a federated learning framework with the SAM2 foundation model, initializing from strong pre-trained weights to enable rapid convergence and robust performance across heterogeneous WTPs. By leveraging real, synthetic, and public datasets, the approach demonstrates competitive Dice and IoU scores while preserving data locality, with edge-edge coordination via a fog orchestrator. The work highlights the practical viability of integrating large-scale foundation models into privacy-preserving distributed learning for real-time industrial monitoring, and suggests directions for more advanced FL schemes and tighter edge optimization.

Abstract

Foam formation in Wastewater Treatment Plants (WTPs) is a major challenge that can reduce treatment efficiency and increase costs. The ability to automatically examine changes in real-time with respect to the percentage of foam can be of great benefit to the plant. However, large amounts of labeled data are required to train standard Machine Learning (ML) models. The development of these systems is slow due to the scarcity and heterogeneity of labeled data. Additionally, the development is often hindered by the fact that different WTPs do not share their data due to privacy concerns. This paper proposes a new framework to address these challenges by combining Federated Learning (FL) with the state-of-the-art base model for image segmentation, Segment Anything Model 2 (SAM2). The FL paradigm enables collaborative model training across multiple WTPs without centralizing sensitive operational data, thereby ensuring privacy. The framework accelerates training convergence and improves segmentation performance even with limited local datasets by leveraging SAM2's strong pre-trained weights for initialization. The methodology involves fine-tuning SAM2 on distributed clients (edge nodes) using the Flower framework, where a central Fog server orchestrates the process by aggregating model weights without accessing private data. The model was trained and validated using various data collections, including real-world images captured at a WTPs in Granada, Spain, a synthetically generated foam dataset, and images from publicly available datasets to improve generalization. This research offers a practical, scalable, and privacy-aware solution for automatic foam tracking in WTPs. The findings highlight the significant potential of integrating large-scale foundational models into FL systems to solve real-world industrial challenges characterized by distributed and sensitive data.

Figures (14)

• Figure 1: Sample images from the ADE20K dataset used for zero-shot segmentation
• Figure 2: Synthetic foam/water images generated via Turbo SDXL
• Figure 3: Overall workflow of day/night segmentation. The raw images are taken from the WTPs, pre-processed via OpenCV-based contrast adjustments (middle), and thresholded adaptively to yield a final foam mask
• Figure 4: Sample images from the real-world WTPs dataset collected in Granada
• Figure 5: FL architecture for foam segmentation in WTPs