Introduction
Coherent imaging techniques, such as ptychography, offer unparalleled multi-scale views of materials across various scientific and technological fields. Advances in brighter sources and high-rate detectors have increased the potential of ptychography, but this comes with a significant increase in data and computational demands, hindering real-time imaging. Traditional ptychography involves scanning a coherent beam across a sample, recording far-field diffraction patterns, and then computationally inverting these patterns to reconstruct the sample image. This inversion, or phase retrieval, is computationally intensive, particularly with the high data rates from state-of-the-art instruments. Conventional phase retrieval methods require acquiring tens or hundreds of diffraction patterns before producing results, limiting real-time capabilities. Furthermore, the spatial overlap needed for numerical convergence restricts the sample volume imaged and can damage dose-sensitive specimens. This work addresses these limitations by utilizing deep learning at the edge to enable real-time ptychographic imaging. This AI-driven approach drastically accelerates the image reconstruction process and reduces the need for oversampling, making it suitable for low-dose imaging of various samples, including those sensitive to radiation damage.
Literature Review
Ptychography is a widely used high-resolution imaging technique in X-ray, optical, and electron microscopy. X-ray ptychography excels in non-destructive nanoscale imaging of large objects with minimal sample preparation. Optical ptychography enables 3D imaging of large samples, and advancements in algorithms have facilitated gigapixel-scale imaging. Electron ptychography has achieved sub-angstrom resolution. However, the computational demands of ptychography, particularly with increased data rates from modern instruments, pose a significant challenge. Deep learning has emerged as a promising solution to accelerate data analysis in coherent imaging, with deep convolutional neural networks showing promise in outperforming conventional iterative algorithms in terms of speed and, increasingly, reconstruction quality. However, real-time coherent imaging using deep learning has not been previously demonstrated.
Methodology
The proposed workflow consists of three concurrent components: measurement, online training, and live inference. Diffraction patterns are captured using a detector while scanning an X-ray beam across the sample in a spiral pattern. The diffraction data are sent to high-performance computing (HPC) resources for phase retrieval using conventional iterative algorithms. These results, paired with the corresponding diffraction patterns, create labeled data for training a neural network. The trained neural network is periodically sent to an edge device for low-latency inference. Diffraction patterns are streamed concurrently to the edge device via a network. The edge device, updated with the latest model, processes individual diffraction patterns and sends reconstructed data back to the beamline computer for real-time sample imaging. The workflow is automated, enabling real-time experimental steering and error detection. A modified version of PtychoNN (PtychoNN 2.0) is used for inference, chosen for its speed and accuracy. Continual learning is employed to maintain high accuracy and adapt to new sample features as the experiment progresses. The neural network is trained using iterative phase retrieval results as ground truth, minimizing the mean absolute error between the target and AI inference. The training is performed online, constantly updating the model during the experiment. Inference times are optimized using TensorRT on an NVIDIA A40 Xavier edge device. The accuracy of AI inference is compared to conventional iterative phase retrieval across various experimental conditions, including different overlap ratios and low photon counts. Strategies to handle low photon count data include scaling up experimental intensity or scaling down training data intensity and retraining the model.
Key Findings
The AI-enabled workflow successfully performs real-time ptychographic imaging at a frame rate up to 2 kHz. The accuracy of AI inference closely matches that of conventional iterative phase retrieval under similar experimental conditions (high overlap). Importantly, the workflow maintains high accuracy even with sparse sampling (low overlap ratios), significantly reducing the photon dose required for imaging. For example, the workflow achieves over 90% accuracy even without overlap, reducing the dose by a factor of 6.25 compared to conventional methods. Further dose reduction is achieved by upscaling experimental intensity or scaling down training data and retraining the model. Upscaling experimental data intensity by a factor of 10, resulted in an accuracy of 86%. Scaling down training data intensity by a factor of 10,000, resulted in over 80% accuracy. Continual learning ensures that the model adapts to new sample features, maintaining high accuracy throughout the experiment. The inference speed is ultimately limited by the network connection, with 100 Hz achieved for 512x512 pixels and 2 kHz for 128x128 pixels. Using a powerful GPU, inference time was reduced to 70 μs per image (14 kHz). The workflow is shown to be robust to variations in count rate, maintaining reliable phase inference even with count rate changes by a factor of 16. Beyond a factor of 40 the reconstruction is no longer usable.
Discussion
This work demonstrates a significant advancement in ptychographic imaging, overcoming the computational bottlenecks that limit real-time capabilities. The combination of deep learning at the edge and HPC resources provides a practical solution for high-throughput ptychographic experiments, particularly at next-generation light sources and advanced electron microscopes. The ability to perform low-dose imaging is particularly impactful for sensitive samples, enabling high-resolution imaging without causing significant damage. The automated workflow and online training strategy streamline the experimental process, improving efficiency and reducing the reliance on extensive post-processing. The high accuracy of AI inference, coupled with its speed, opens up new possibilities for dynamic imaging and real-time sample manipulation during experiments.
Conclusion
This study successfully demonstrates real-time streaming ptychographic imaging using an AI-enabled workflow. The use of deep learning at the edge dramatically reduces computational time and data requirements. The methodology achieves high-accuracy reconstructions, even under low-dose conditions, making it suitable for a broad range of applications involving beam-sensitive samples. Future research could focus on further improving the model's robustness to variations in experimental parameters and exploring its applicability to other coherent imaging modalities.
Limitations
The accuracy of the AI inference is dependent on the quality of the training data generated by iterative phase retrieval. The model is trained for a specific range of refractive indices and a fixed illumination probe; samples outside this range may require retraining. While the model demonstrates resilience to some fluctuations, significant variations in count rate may affect reconstruction fidelity. The current implementation's speed is partially limited by network bandwidth.
Related Publications
Explore these studies to deepen your understanding of the subject.