Pathology Image Restoration via Mixture of Prompts

TL;DR

This work tackles the challenge of producing high-quality all-in-focus pathology images from a single focal plane by introducing a two-stage restoration pipeline guided by a novel mixture of prompts. A coarse-stage P-former uses pathology-oriented prompts with an MoE-based gating to restore semantic content, followed by a fine-stage P-diffusion that refines high-frequency details under semantic and structural constraints from $P_P$ and $P_E$. The prompts are generated from a defocus-aware framework combining a defocus encoder and pathology foundation models, enabling robust priors even on degraded inputs. Across in-house and 3DHistech datasets, the approach achieves strong distortion and perceptual metrics and improves downstream diagnostic performance, illustrating the practical impact of integrating domain-specific priors into a two-stage restoration framework. A remaining challenge is diffusion-time, which motivates future work to accelerate inference without compromising fidelity.

Abstract

In digital pathology, acquiring all-in-focus images is essential to high-quality imaging and high-efficient clinical workflow. Traditional scanners achieve this by scanning at multiple focal planes of varying depths and then merging them, which is relatively slow and often struggles with complex tissue defocus. Recent prevailing image restoration technique provides a means to restore high-quality pathology images from scans of single focal planes. However, existing image restoration methods are inadequate, due to intricate defocus patterns in pathology images and their domain-specific semantic complexities. In this work, we devise a two-stage restoration solution cascading a transformer and a diffusion model, to benefit from their powers in preserving image fidelity and perceptual quality, respectively. We particularly propose a novel mixture of prompts for the two-stage solution. Given initial prompt that models defocus in microscopic imaging, we design two prompts that describe the high-level image semantics from pathology foundation model and the fine-grained tissue structures via edge extraction. We demonstrate that, by feeding the prompt mixture to our method, we can restore high-quality pathology images from single-focal-plane scans, implying high potentials of the mixture of prompts to clinical usage. Code will be publicly available at https://github.com/caijd2000/MoP.

Figures (7)

• Figure 1: Our method addresses key issues in pathology image restoration by extracting a mixture of prompts. These prompts are utilized to guide a two-stage restoration process, consisting of a coarse-stage P-former and a fine-stage P-diffusion.
• Figure 2: The framework of our method. (a) Mixture of prompts. The defocus prompt $P_D$ is utilized to restore the pathology prompt $P_P$ and weight the edge prompt $P_E$ from the low-quality image. (b) P-former uses the pathology prompt $P_{P}$ to enhance Restormer zamir2022restormer blocks within MoE. (c) P-diffusion is constrained structurally and semantically by the edge prompt $P_{E}$ and $P_{P}$.
• Figure 3: Visual comparison between different methods on the in-house dataset. The distinguished regions within the yellow bounding boxes are zoomed in at the bottom right.
• Figure 4: Within the same field of view, the original image and the heatmap generated by the defocus estimator exhibit variations with changes in the focal distance. The red and orange boxes highlight localized magnified views for closer inspection.
• Figure 5: The visualization obtained using t-SNE illustrates the effect of defocus-aware prompt restoration, showcasing the feature distributions of $P_{LP}$ (green), $P_{P}$ (blue) and $P_{HP}$ (orange).