Overlap-aware segmentation for topological reconstruction of obscured objects

TL;DR

OASIS tackles the hard problem of disentangling overlapping object signals in scientific imaging by introducing a region- and channel-weighted segmentation-regression loss that prioritizes overlap areas. Implemented on a U-Net backbone, OASIS outputs k intensity maps and uses a loss L = α1 L_seg-reg + α2 L_smooth, with explicit formulas for L_seg-reg and L_smooth to emphasize ambiguous regions. Applied to the MIGDAL experiment, OASIS demonstrates substantial gains in low-energy electron-recoil reconstruction and topology (IoU) over unweighted baselines, achieving ER energy reconstructions within a median of −14% and IoU of 0.855, and yielding angular consistency within ~20° for ER energies ≥4 keV. The method is computationally efficient, generalizable to other overlapping-source problems, and released as open source for cross-domain adoption in fields such as astronomy and medical imaging.

Abstract

The separation of overlapping objects presents a significant challenge in scientific imaging. While deep learning segmentation-regression algorithms can predict pixel-wise intensities, they typically treat all regions equally rather than prioritizing overlap regions where attribution is most ambiguous. Recent advances in instance segmentation show that weighting regions of pixel overlap in training can improve segmentation boundary predictions in regions of overlap, but this idea has not yet been extended to segmentation regression. We address this with Overlap-Aware Segmentation of ImageS (OASIS): a new segmentation-regression framework with a weighted loss function designed to prioritize regions of object-overlap during training, enabling extraction of pixel intensities and topological features from heavily obscured objects. We demonstrate OASIS in the context of the MIGDAL experiment, which aims to directly image the Migdal effect--a rare process where electron emission is induced by nuclear scattering--in a low-pressure optical time projection chamber. This setting poses an extreme test case, as the target for reconstruction is a faint electron recoil track which is often heavily-buried within the orders-of-magnitude brighter nuclear recoil track. Compared to unweighted training, OASIS improves median intensity reconstruction errors from -32% to -14% for low-energy electron tracks (4-5 keV) and improves topological intersection-over-union scores from 0.828 to 0.855. These performance gains demonstrate OASIS's ability to recover obscured signals in overlap-dominated regions. The framework provides a generalizable methodology for scientific imaging where pixels represent physical quantities and overlap obscures features of interest. All code is openly available to facilitate cross-domain adoption.

Figures (6)

• Figure 1: Schematic of OASIS. An input intensity-map image of dimension $n_x\times n_y$ containing $k$ distinct object classes ($k=5$ in this example) is passed as input into the network. The backbone network, where here we use U--Net, processes the images which are then passed into a segmentation-regression head, shown as the orange block. This maps the backbone network's output to $k$ intensity maps of dimension $n_x\times n_y$, each corresponding to a distinct object class. During training, OASIS is optimized using a custom loss function that compares the predicted output intensity maps with truth. The loss function incorporates weights that can be tuned to assign higher penalties to reconstruction errors in overlap regions (white patches in the input image) and to channels with fainter objects.
• Figure 2: Illustration showing the hybrid signal simulation construction of two events: one with a 5.8keV ER (top row) and the other with a 9.8keV ER (bottom row). For each event, the hybrid signal (left column) is constructed by stitching a simulated ER track (middle column) with a real, measured NR track (right column). The stitching point is the truth vertex position of the simulated ER and the estimated vertex of the real NR (both shown as white dots).
• Figure 3: Comparison of OASIS's track reconstruction performance between the unweighted and weighted training campaigns. Left: Median (points) and 25th-75th percentile ranges (error bands) of the percent error of predicted ER track intensity compared to truth, $\Delta I$, versus truth ER energy. Right: Median (points) and 25th-75th percentile ranges (error bands) of pixel IoU between OASIS's predicted ER and truth versus truth ER energy. In both cases, training OASIS with weights prioritizing regions of overlap lead to significant improvements at low ER energies, where pixel overlap is more significant in the ER channel.
• Figure 4: Four test set examples comparing OASIS's ER reconstruction performance to truth. In panels (a)-(c) both the truth and predicted ERs have principal curves shown in white with estimated directional axes shown in black. Panel (d) is an NR-only input event. Predicted ER and truth ER energies for each panel are: (a) $\hat{E}_\mathrm{ER}=$ 5.2keV, $E_\mathrm{truth,ER}=$ 5.2keV. (b) $\hat{E}_\mathrm{ER}=$ 5.5keV, $E_\mathrm{truth,ER}=$ 5.2keV. (c) $\hat{E}_\mathrm{ER}=$ 0.89keV, $E_\mathrm{truth,ER}=$ 5.7keV. (d) $\hat{E}_\mathrm{ER}=$ 0.6eV, $E_\mathrm{truth,ER}=$ 0eV, where $\hat{E}_\mathrm{ER}\equiv\frac{\hat{I}_\mathrm{ER}}{I_\mathrm{truth,ER}}E_{\mathrm{truth,ER}}$.
• Figure 5: Left: Histograms of reconstructed and truth ER intensities for the test set signal sample (blue and orange, respectively), and reconstructed ER intensities for the NR-only sample (green). The black dashed line represents the 3keV false positive threshold for ER detection in the NR-only sample. Right: Model's predicted ER track intensity ($\hat{I}_\mathrm{ER}$) versus truth ER track intensity ($I_\mathrm{truth,ER}$) for all test set ERs satisfying $4\leq E_\mathrm{truth,ER}\leq 15$$\,$keV. The red line shows equality between prediction and truth illustrating that OASIS, on average, nearly perfectly reconstructs energy.