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

AI summary

Overview

Research area: Applied computer vision and privacy-preserving machine learning — specifically federated image segmentation of foam in wastewater treatment plants using a foundation model.

Technical level: Intermediate. The paper assumes familiarity with image segmentation, fine-tuning, and the client-server vocabulary of federated learning, but its framing is applied engineering rather than new theory.

Scope: This paper proposes and evaluates a federated learning framework that fine-tunes Segment Anything Model 2 (SAM2) across distributed wastewater treatment plant clients, using SAM2's pre-trained weights as the initialization to accelerate convergence and protect plant data privacy.

What This Paper Is About

Foam in wastewater treatment plants reduces treatment efficiency, increases operating costs, and risks discharging partially treated effluent. Automating foam detection with deep learning is difficult because labeled foam images are scarce, foam looks different at every facility, and plants refuse to centralize their operational data for privacy reasons. The paper's goal is a foam segmentation system that trains across multiple plants without any raw image ever leaving a plant, and that converges quickly despite small local datasets.

Key Contributions

1. A federated learning framework tailored to wastewater treatment plant foam segmentation, built with the Flower framework, using FedAvg aggregation on a Fog server that orchestrates edge-node clients and never accesses raw images.
2. A novel initialization strategy using SAM2, where the global model starts from pre-trained SAM2 weights (sam2_hiera_small.pt) rather than random initialization, adapted to foam imagery characteristics.
3. A systematic evaluation under both IID and non-IID data splits, intended to reflect realistic heterogeneity between wastewater treatment plants.
4. An integrated, deployable implementation — described as an open-source project (https://github.com/ertis-research/zerovision) — spanning area scan cameras, on-premise edge nodes, an SFTP-based image retrieval pipeline, a PostgreSQL database, and on-premise Fog server deployment.

Main Findings

• Federated SAM2 achieves 84.38% Dice score with two clients. The paper reports this result for the SAM2-based approach under federated learning with two clients.
• SAM2 initialization accelerates convergence. Federated models initialized with SAM2 weights converge in a significantly smaller number of communication rounds than models with random initialization, attributed to the strength of SAM2's pre-trained weights.
• Five communication rounds proved sufficient in the experiments. The authors chose five rounds based on early tests indicating swift convergence, likely due to SAM2's strong initialization.
• Hierarchical features improve foam boundaries. SAM2's hierarchical vision transformer extracts multi-scale feature maps that refine segmentation boundaries, which the authors say aids accurate foam detection under difficult conditions.
• Centralized classical models can score higher on some metrics, but SAM2 remains competitive. The paper states that U-Net and DeepLabV3+ can achieve higher scores on certain metrics in centralized training, while SAM2 delivers highly competitive and robust performance within the practical constraints of a federated, privacy-preserving system.
• Prompt-sensitivity analysis shows comparable visual outputs. Despite differences in evaluation scores, the authors report that SAM2's visual outputs are comparable to other leading models.
• The initial SAM approach was discarded. Experiments with the original SAM, using a frozen ViT-B image encoder and bounding box prompts, revealed limitations in fine-tuning for this task, leading the authors to adopt SAM2 instead.
• Dataset composition. The study used 3,980 images total: 1,888 ADE20K water segmentations for zero-shot transfer, 1,792 real-world images from a wastewater treatment plant in Granada, Spain, and 300 synthetic foam images. For the ADE20K zero-shot dataset specifically, the text states 3,020 frames were used for training and 756 for testing.
• Data collection specifics. Granada images were captured at 10-minute intervals on selected days in October 2024 within the Zerovision project, and 1,792 images remained after excluding missing or low-quality images.
• Results not reported in the provided content. The truncated content does not include the full per-round metrics tables, ablation results, or explicit comparative numbers for the IID versus non-IID splits beyond stating that both were evaluated.

Methodology in Plain English

The authors set up a system with two sides. On one side, a central "Fog server" holds a global segmentation model and hands it out. On the other side, each wastewater plant has an edge computer that keeps its own images locally.

The global model starts not from scratch but from SAM2, a large pre-trained segmentation foundation model that already knows how to find object boundaries in general images. Each round, the server sends the current model to selected plant clients (chosen with a round-robin strategy). Each client fine-tunes it on its own private images using the AdamW optimizer with a learning rate of 1e-5 and weight decay of 4e-5, running 30 local epochs with a batch size of 32. Clients then send back only the updated model weights, and the server combines them with a weighted average (FedAvg), where each client's contribution is weighted by how many samples it has. Because only weights travel, no raw plant images ever leave the facility.

To make training work with little labeled data, the authors built datasets from three sources: an open-source water segmentation dataset (ADE20K) used for zero-shot transfer, 300 synthetic foam images generated with the Turbo SDXL diffusion model, and real plant images from Granada. For the real images, they avoided manual annotation by creating masks automatically with a day/night pipeline: images are converted to grayscale and, if average brightness falls below 100 they are treated as night images and brightened and denoised, otherwise a CLAHE contrast method is applied; then adaptive Gaussian thresholding, morphological opening with a 3x3 kernel for 2 iterations, and connected-component filtering that discards regions smaller than 75 pixels produce the final foam mask.

All images are resized to 1024x1024 pixels and augmented with the Albumentations library: horizontal and vertical flips (p=0.5 each), random brightness and contrast changes (p=0.2), and random shifts of ±5%, scales of ±10%, and rotations of ±15° (p=0.5). Instead of manually drawing prompts, the framework automatically generates point prompts from the ground-truth masks — a positive point picked from a random pixel inside each foam object, and a negative point near the image center for foam-free images. Training uses a combination of Dice and BCE losses, and SAM2 adds a small penalty term aligning predicted mask-quality scores with Intersection over Union. The experimental setup used two clients: one trained on the real Granada plant dataset, the other on the synthetic dataset.

Why This Matters

Impact on research. The paper connects two areas that are usually studied separately: large pre-trained foundation models for segmentation and federated learning under privacy constraints. Its central claim — that strong pre-trained initialization can compensate for small, non-identical local datasets in federated settings — offers a practical recipe for other domains where data is distributed and sensitive. It also extends prior federated wastewater work (Carballo Mato et al., 2024, which reported an 86% Dice score using DeepLabv3+ and One Shot Texture Segmentation in centralized training) into a privacy-preserving distributed setting.

Real-world applications:

• Wastewater treatment plant operations — continuous automated foam tracking in aeration tanks and secondary clarifiers, enabling timely control measures instead of relying on subjective visual inspection.
• Environmental compliance and effluent safety — early foam detection helps prevent discharge of partially treated effluent into the environment.
• Cost and efficiency management — foam reduces oxygen transfer and biomass concentration and increases maintenance expenses, so automated quantification supports cost reduction.
• Other privacy-constrained industrial settings — the authors suggest the approach extends to resource management and industrial applications where data sharing is restricted and data is produced in real time.

Industry relevance. The framework is designed around equipment plants already have: existing industrial CCTV cameras writing to an on-premise SFTP server, an on-premise edge computer, and an on-premise Fog server. The architecture is explicitly built to scale across different wastewater treatment plants, and the authors note that centralizing raw wastewater data can be difficult or legally fraught (citing Newhart et al., 2019 on legal and ethical implications). This makes the approach deployable without renegotiating data-sharing agreements.

Future Directions

• Extending beyond five communication rounds and two clients. The experiments used only two clients and five rounds; scaling to more plants and longer training regimes is the obvious next test of whether the convergence advantage holds.
• Making better use of the non-IID evaluation. The paper lists systematic IID and non-IID evaluation as a contribution and notes that non-IID conditions cause slow convergence or reduced performance in other federated work, but the truncated content does not report the comparative outcomes — those numbers are needed to judge the method's robustness to plant heterogeneity.
• Improving on centralized baselines. Because U-Net and DeepLabV3+ can score higher on certain metrics in centralized training, closing that gap while keeping privacy is an open question.
• Validating the synthetic and zero-shot data strategies. The authors describe the Turbo SDXL synthetic dataset as experimental and its efficacy as requiring further evaluation, and note that the ADE20K zero-shot approach led to segmentation errors — both are natural candidates for follow-up work.
• Broader domain transfer. The authors suggest the insights may apply to other industrial and resource-management settings with restricted data sharing, which would need separate validation.

Target Audience

This paper is most useful to applied machine learning engineers and researchers working on industrial computer vision, and to wastewater treatment plant operators and utility technologists evaluating automation. It also suits federated learning practitioners looking for a real-world deployment case study with foundation-model initialization, and researchers in privacy-preserving AI who want to see how initializing from large pre-trained models affects convergence under distributed, non-identical data.

Authors’ 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.