Diffusion-Shock PDEs for Deep Learning on Position-Orientation Space

We extend regularised diffusion-shock (RDS) filtering from Euclidean space $R^2$ (Schaefer and Weickert in J Math Imaging Vis 66:447-463, 2024. 10.1007/s10851-024-01175-0) to position-orientation space $M_2\cong R^2\times S^1$ . This has numerous advantages, e.g. making it possible to enhance and inpaint crossing structures, since they become disentangled when lifted to $M_2$ . We create a version of the algorithm using gauge frames to mitigate issues caused by lifting to a finite number of orientations. This leads us to study generalisations of diffusion, since the gauge frame diffusion is not generated by the Laplace-Beltrami operator. RDS filtering compares favourably to existing techniques such as total roto-translational variation (TR-TV) flow (Smets et al. in J Math Imaging Vis 63:237-262, 2021. 10.1007/s10851-020-00991-4; Chambolle and Pock in Numer Math 142:611-666, 2019. 10.1007/s00211-019-01026-w), NLM (Buades et al. in Image Process On Line 1:208-212, 2011. 10.5201/ipol.2011.bcm_nlm), and BM3D (Dabov et al. in Trans Image Process 16:2080-2095, 2007. 10.1109/TIP.2007.901238) when denoising images with crossing structures, particularly if they are segmented. Furthermore, we see that $M_2$ RDS inpainting is indeed able to restore crossing structures, unlike $R^2$ RDS inpainting. In addition to the contributions of our SSVM submission (Sherry et al. in: Bubba, Gaburro, Gazzola, Papafitsoros, Pereyra, Schönlieb (eds) 10th International Conference on Scale Space and Variational Methods in Computer Vision II (SSVM), vol. 15668, pp. 205-217. Springer, Cham, 2025. 10.1007/978-3-031-92369-2_16), in this extended work we provide new theorical results and automate RDS filtering by integrating it into a geometric deep learning framework. Regarding our theoretical contributions, we prove that our generalised diffusions are still well posed, smoothing, and analytic. We developed an RDS filtering PDE layer for the PDE-CNN and PDE-G-CNN deep learning frameworks, using a novel gating mechanism. We show that these new RDS PDE layers can be beneficial in various impainting and denoising tasks.