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

TL;DR

This paper introduces a privacy-preserving federated learning framework that fine-tunes a large pre-trained image segmentation model for foam detection in wastewater treatment plants, enhancing real-time monitoring with limited data.

Contribution

It combines Federated Learning with the Segment Anything Model 2 to enable collaborative, privacy-aware foam segmentation across multiple plants, addressing data scarcity and heterogeneity.

Findings

Improved segmentation accuracy with limited local data

Enhanced training convergence through federated approach

Demonstrated effectiveness on real-world and synthetic datasets

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.