ML-TN-009 — AI at the edge: IoT real-time endoscopes and Federated Learning

From DAVE Developer's Wiki
Jump to: navigation, search
Info Box
NeuralNetwork.png Applies to Machine Learning
TBD.png


History[edit | edit source]

Version Date Notes
0.1.0 March 2025 First public draft

Abstract[edit | edit source]

This article summarizes the work carried out by Niccolò Brusadin for his internship at DAVE Embedded System. He prototyped a Federated Learning system whose edge devices are machines emulating IoT smart endoscopes for automatic early detection of gastro-intestinal tract polyps.

Early detection of colorectal polyps is crucial for colorectal cancer prevention, however current endoscopic techniques have limitations regarding detection accuracy and efficiency, since most studies proved a significant number of polyps can be missed out during routine exam procedures. Deep Learning-based computer detection systems have been developed to assist endoscopists with real-time polyp detection but training these models on medical data raises concerns about privacy, data security, and regulatory compliance.

This thesis explores the integration of Federated Learning (FL), a Machine learning paradigm that allows multiple institutions to collaboratively train a distributed AI model without sharing sensitive patient data. The work develops a Federated Learning system for automatic GI polyp detection based on an IoT-based platform. In such a context, the role of edge devices will be played by a small fleet of prototypes of a "smart" endoscope able to process in real-time streams through a YOLOv5 deep learning model. The training infrastructure will be created using the NVFlare framework, which is an open-source FL platform for privacy-preserving AI applications. This approach ensures data privacy while improving accuracy and reliability for polyp detection without sharing personal data.

Experimental assessments of the centralized approach demonstrated the efficacy and high accuracy of the YOLOv5 model for polyp detection, achieving strong performance in terms of mAP. Evaluations of the federated training scenarios indicated that it produced comparable performance compared to the centralized learning environment while maintaining data privacy. Moreover, the applicability of the inference-federated hybrid system was demonstrated, which allowed real-time polyp detection. Limitations were identified, including hardware capabilities and communications latency. This suggests future optimizations could be made to the growing area of AI that focuses on patient privacy in imaging studies, and create opportunities for using federated AI, in practice.

The work in based on the achievements of two previous internships detailed here and here.

If you are interested in having Brusadin's entire thesis, please send a request to this address.

Introduction[edit | edit source]

Colorectal cancer (CRC) is a prominent malignant tumour in the digestive system, primarily affecting the colon or rectum, and it significantly contributes to global cancer mortality, accounting for around 10% of all cases. Early detection is crucial, as survival rates can rise from approximately 63% to 91% when diagnosed early. Colonoscopy is the gold standard for identifying CRC, allowing detailed examination of polyps. However, the endoscopist’s expertise greatly influences its sensitivity, and polyps can be easily missed due to various factors, highlighting the need for advanced digital solutions in detection.

Recent research has turned toward AI-assisted screening, employing Deep Learning techniques for object detection and semantic segmentation in polyp diagnosis. YOLO models, known for their efficacy in medical image analysis, offer real-time polyp detection capabilities, balancing speed and performance. Nevertheless, training AI models in healthcare raises significant data privacy and security concerns due to the centralized nature of data collection.

Federated Learning (FL) addresses these issues by enabling devices to collaboratively train a shared AI model while keeping patient data localized, thus maintaining privacy. This decentralized approach not only protects sensitive information but also enhances model generalization across varied clinical contexts. Various FL frameworks have been developed for healthcare applications, with NVFlare offering high-efficiency simulation tools for research and robust production capabilities for enterprise users.

Purpose of the thesis and research questions[edit | edit source]

This work aims to develop a Federated Learning system for automatic detection of gastrointestinal polyps, by integrating deep-learning inference tasks with privacy-preserving training. This research addresses the challenges of AI-assisted colonoscopy, with a particular focus on diagnostic accuracy, computational efficiency, and protection of patient data.

The proposed system will employ an IoT-based smart endoscope device prototyped with an embedded device running a Debian GNU/Linux distribution and integrated with a YOLOv5 deep learning model for real-time polyp detection. The FL training infrastructure will be implemented using NVFlare, with two separate nodes to jointly train the global model without sharing raw medical data. This approach aims to enhance model generalization across various clinical environments while ensuring data privacy regulations. The key points of this research will be as follows:

  1. Development of a deep-Learning model for polyp-detection: a YOLOv5 Deep Learning model will be selected to be used as the training model in the federated learning framework and as the inference model to perform detection tasks on smart endoscopic devices.
  2. Integration of Nvidia NVFlare Federated Learning paradigm: the developed model will be implemented in the NVFlare framework as the starting point for the federated learning scenario.
  3. Investigation of the impact of decentralized training: a comparison between FL training and centralized training will be assessed to determine whether a federated learning environment improves model robustness across different nodes.
  4. Implementation of FL polyp detection in real practice: the purpose is to investigate the applicability of NVFlare framework in real-world environments and evaluate its ability to maintain model accuracy, reduce reliance on centralized datasets, and ensure compliance with healthcare privacy regulations.
  5. Design a system that integrates federated training and detection: the goal is to design a server-edge system where devices initially perform inference locally, then select only poorly detected images to be used for federated training. This iterative process focuses on enhancing model accuracy by training the model only in challenging cases and avoiding redundant training on well-classified images.
  6. Model Optimization for the designed workflow: finally, we aim to ensure that the model runs on edge devices with adequate computational resources, without compromising detection accuracy.

Based on these objectives, the research aims to answer the following questions:

  1. How can Federated Learning be integrated with Deep Learning YOLOv5 model for polyp detection while preserving patient data privacy?
  2. How does the Federated Learning approach impact model generalization and robustness across different nodes?
  3. What is the trade-off between model accuracy and computational efficiency during YOLOv5 model deployment in real-time polyp detection for an IoT-edge device system?
  4. What are the practical challenges in implementing the FL-based AI models in real-world medical environments and how could they be addressed?

This thesis attempts to find answers to the questions above while demonstrating the applicability of the federated framework in polyp detection.