Overlap-aware segmentation for topological reconstruction of obscured objects

This paper introduces OASIS, a novel segmentation-regression framework that employs a weighted loss function to prioritize overlapping regions during training, significantly improving the intensity and topological reconstruction of faint, obscured electron tracks in the challenging context of the MIGDAL experiment.

The Gist

The Big Problem: The "Flashlight in a Storm"

Imagine you are trying to take a photo of a tiny, faint firefly (an electron) flying through a massive, blinding thunderstorm (a nuclear recoil). In the world of particle physics, specifically in an experiment called MIGDAL, scientists are trying to do exactly this.

They are looking for a rare event where a nucleus gets hit by a particle and, as a result, kicks out a tiny electron. The problem is that the "hit" (the nucleus) creates a huge, bright, messy trail of light, while the "kick" (the electron) is a tiny, faint trail that often gets completely swallowed up by the brightness of the storm.

In standard computer vision, if you ask an AI to look at this photo and separate the firefly from the storm, the AI usually gets confused. It sees the bright storm and assumes everything belongs to it, or it tries to split the image evenly, missing the faint firefly entirely.

The Solution: OASIS (The "Smart Spotlight")

The authors of this paper created a new AI framework called OASIS (Overlap-Aware Segmentation of ImageS).

Think of training a normal AI like teaching a student to grade a test where every question is worth the same number of points. If the student gets the easy questions right but misses the hard, tricky ones, they still get a decent grade.

OASIS changes the rules of the test. It tells the AI: "Hey, the part of the image where the bright storm and the faint firefly overlap is the most important part. If you get that wrong, you get a huge penalty. If you get the easy parts wrong, it's less of a problem."

By giving extra "points" (or penalties) to the messy, overlapping areas during its training, the AI learns to pay special attention to the difficult spots where the two signals mix.

How It Works (The Recipe)

1. The Network: They used a standard AI architecture called U-Net (think of it as a very skilled artist who can look at a messy painting and try to separate the colors).
2. The Special Sauce: They added a custom "loss function." In AI terms, a "loss function" is how the computer measures how wrong it is. OASIS's loss function has a special knob that turns up the volume on errors made in the overlap zones.
3. The Training: They showed the AI thousands of images. Some had real "storms" (nuclear tracks) with fake "fireflies" (electron tracks) added in. Others had just storms. The AI learned to separate the two, but because of the special penalty system, it became an expert at finding the faint firefly even when it was buried under the storm.

The Results: Finding the Invisible

The team tested this on the MIGDAL experiment data. Here is what they found:

• Before OASIS: When the AI tried to guess the energy of the faint electron, it was often off by about 41%. It was basically guessing in the dark.
• After OASIS: By using the "overlap-aware" training, the error dropped to just 13%.
• The "Firefly" Test: In cases where the electron was very faint and almost entirely hidden by the bright nuclear track, OASIS could still see it. It successfully separated the two signals, allowing scientists to measure the electron's energy and direction much more accurately.
• No False Alarms: The AI didn't start seeing fireflies where there were none. When shown a picture with only a storm (no electron), it correctly said, "I don't see a firefly here," most of the time.

Why This Matters

The paper claims that this method is a game-changer for the MIGDAL experiment. Because the probability of this rare event happening increases when the electron is very low energy (and therefore even fainter and harder to see), being able to reconstruct these faint signals is crucial.

Without OASIS, scientists might miss the most interesting part of the data. With OASIS, they can finally "see" the faint electron tracks that were previously buried in the noise, allowing them to test theories about dark matter and how particles interact.

Summary in One Sentence

The paper introduces OASIS, a smart AI training method that forces computers to focus extra hard on the messy, overlapping parts of an image, allowing them to successfully separate a tiny, faint signal from a massive, bright background that would normally hide it completely.

Technical Summary

Technical Summary: OASIS for Overlap-Aware Segmentation in Scientific Imaging

Problem Statement
Scientific imaging frequently encounters the challenge of decoupling overlapping objects where pixel intensities represent physically meaningful quantities. In domains ranging from astronomy (blended galaxies) to particle physics (overlapping particle tracks), standard deep learning segmentation-regression algorithms typically treat all image regions equally. This approach fails to prioritize regions of overlap, where the attribution of intensity to specific objects is most ambiguous. This is particularly critical in the MIGDAL experiment, which seeks to image the Migdal effect—a rare process where nuclear scattering induces electron emission. The target signal is a faint electron recoil (ER) track that is often heavily obscured by an order-of-magnitude brighter nuclear recoil (NR) track. Standard methods struggle to reconstruct the energy and topology of these faint, buried ER tracks, limiting the ability to test theoretical models of dark matter interactions.

Methodology: The OASIS Framework
The authors introduce Overlap-Aware Segmentation of ImageS (OASIS), a segmentation-regression framework designed to disentangle overlapping objects by explicitly weighting overlap regions during training.

• Network Architecture: OASIS utilizes a U-Net backbone with a four-stage encoder, a bottleneck, and a decoder path. Skip connections link corresponding encoder and decoder stages. The network processes an input intensity map and outputs $k$ distinct intensity maps, one for each object class. In the MIGDAL application, $k=2$ (ER and NR channels).
• Loss Function: The core innovation lies in a custom weighted loss function, $L$, comprising two terms:
  1. Channel- and Region-Weighted Regression Loss ($L_{seg-reg}$): This term minimizes the mean absolute error between predicted and ground-truth intensities. Crucially, it incorporates programmable weights:
     • $w_c$: Channel-specific weights to prioritize specific object classes (e.g., the faint ER).
     • $W_{region}(x, y)$: Region-specific weights that can be tuned to assign higher penalties to errors in specific spatial regions, such as areas where object masks intersect (overlap regions).
  2. Smoothness Regularization ($L_{smooth}$): A standard term penalizing excessive variation between neighboring pixel predictions to ensure physical continuity.
• Training Strategy: The framework was applied to the MIGDAL experiment using a hybrid dataset. Real NR tracks (from experimental data) were combined with simulated ER tracks to create ground-truth labels. This approach avoids Sim2Real gaps for the complex NR background while providing noise-free ER labels. The training set included both hybrid signal events (ER embedded in NR) and NR-only events to prevent the model from hallucinating ER signals where none exist.

Key Contributions

1. Novel Loss Function Design: The paper introduces the concept of "region-targeted weights" in segmentation regression. Unlike previous instance-segmentation frameworks that handle occlusion but do not regress physical intensities, OASIS explicitly tunes the loss function to prioritize reconstruction accuracy in pixel-overlap regions.
2. Systematic Ablation Study: The authors conducted 36 distinct training campaigns, varying the ratios of channel weights ($w_{ER}/w_{NR}$) and region weights (ER-only, NR-only, and overlap regions). This systematic approach isolated the impact of overlap-aware training.
3. Generalizable Framework: While demonstrated on the MIGDAL experiment, the OASIS framework is presented as a general methodology applicable to any scientific imaging problem requiring the decomposition of overlapping sources into constituent intensity maps, including multi-band astronomy and volumetric medical imaging.

Results
The study evaluated OASIS on a test set of 2,086 signal events and 7,159 NR-only events. The results highlight the following:

• Impact of Overlap Weights: The overlap-region weight ratio ($R_o$) was identified as the single most important hyperparameter. Increasing $R_o$ from 1 (unweighted) to 8 significantly improved the reconstruction of low-energy ER tracks (3–6 keV), where overlap is most dominant.
• Quantitative Improvements: Averaging over training campaigns, the addition of overlap-targeted weights improved the median intensity reconstruction error ($\Delta I_{median}$) for low-energy electrons from −41.1% to −13.3%.
• Topological and Angular Performance: The best-performing model (with weights $r_c=2, R_n=0.5, R_o=8$) achieved a median Intersection-over-Union (IoU) of 0.72 for 4–5 keV ERs, compared to 0.59 for unweighted regression. Furthermore, the angular consistency (the difference between the reconstructed and truth track angles) remained within 20° for ER energies down to 4 keV.
• False Positive Control: The model successfully distinguished between signal events and NR-only events, predicting effectively zero ER signal in the vast majority of NR-only inputs, demonstrating robustness against hallucination.

Significance and Claims
The paper claims that OASIS provides a generalizable path toward recovering obscured signals in overlap-dominated regions, a persistent challenge in scientific imaging. Specifically for the MIGDAL experiment, the methodology enables the reconstruction of faint electron recoil tracks buried within bright nuclear recoils with sufficient accuracy to measure energy and emission angles. This capability is essential for validating theoretical predictions of the Migdal effect, particularly in the low-energy regime where the cross-section increases exponentially.

The authors modestly frame the angular results as a "promising first step" toward demonstrating sensitivity in measuring angular distributions, noting that the current angular consistency is a measure of topological similarity between prediction and truth rather than absolute angular resolution relative to the true recoil direction. The work concludes that by shifting the training objective to prioritize ambiguous overlap regions, OASIS significantly enhances the fidelity of pixel-wise intensity attribution and object topology reconstruction, offering a tool that can be adapted to diverse fields including astronomy (galaxy deblending) and medical imaging.