Lightweight Optimal-Transport Harmonization on Edge Devices

AI summary

Overview

Research area: Computer vision — image color harmonization for augmented reality, combining classical optimal transport theory with lightweight neural network design and on-device deployment.

Technical level: Intermediate. The paper sits at the intersection of optimal transport theory, CNN-based image harmonization, and mobile inference engineering; the math (Monge-Kantorovich linear maps, Lipschitz bounds) is nontrivial, but the method itself is comparatively simple.

Scope: The paper proposes MKL-Harmonizer, an encoder that predicts the 12 parameters of a Monge-Kantorovich Linear optimal transport filter for color harmonization, releases a 327-image ARCore dataset with pixel-perfect masks, and benchmarks accuracy, human-perceived quality, and edge-device speed against existing harmonization baselines.

What This Paper Is About

When a virtual object is pasted into a real camera image, it usually looks out of place because it was captured under different lighting and camera conditions. Image harmonization fixes the pasted region's colors so the composite looks seamless, but existing methods are too heavy for augmented reality, where harmonization must run every frame on a phone or headset. This paper asks whether a very simple linear color transform — specifically the Monge-Kantorovich Linear (MKL) optimal transport filter — can be predicted cheaply enough to run on-device while still looking realistic.

Key Contributions

1. MKL-Harmonizer: A lightweight harmonization method that trains a compact encoder to predict the 12 parameters of a Monge-Kantorovich Linear optimal transport map, rather than regressing pixels or filter coefficients.
2. An AR-specific dataset and toolkit: The authors modified a sample ARCore application into a data-gathering tool and collected 327 pairs of composite-and-mask images captured in the wild, with masks taken directly from the rendering engine (pixel-perfect). The dataset spans indoor and outdoor scenes across different times of day, weather, and lighting.
3. A theoretical justification for a linear filter: The paper bounds the approximation error of the MKL map in terms of the Lipschitz constant of the true transport map and the probability that transformed colors fall outside the valid color gamut, showing linear filters suffice when the true transformation is smooth and colors are not concentrated at gamut boundaries.
4. Speed-quality evidence: The authors benchmark against Harmonizer, PCT-Net, and INR-Harmonization, run a user study with 20 participants (642 total grades), and measure real-time frame rates on Google Pixel 4a and Google Pixel 7 devices.

Main Findings

• An ideal linear filter is already very strong on iHarmony4. A non-predictive "Ideal Linear OT filter" — the exact MKL transform computed per image from ground truth — reaches MSE 7.6 ± 0.2, PSNR 43.6 ± 0.1, and fMSE 45.9 ± 0.9 on the iHarmony4 256x256 test set. The main text describes this figure as approximately 7.0. This supports the claim that real harmonization maps are often close to linear.

• MKL-Harmonizer does not win on MSE-style metrics. On iHarmony4 256x256, the authors report for their L1 variant MSE 65.0 ± 1.6, PSNR 34.1 ± 0.1, fMSE 438 ± 7, and for their L2 variant MSE 66.3 ± 1.7, PSNR 33.9 ± 0.1, fMSE 451 ± 7. PCT-Net scores MSE 29.1 ± 0.9, PSNR 38.0 ± 0.1, fMSE 201 ± 4; Harmonizer scores MSE 40.1 ± 1.2, PSNR 36.6 ± 0.1, fMSE 258 ± 5; INR scores MSE 67.2 ± 1.8, PSNR 35.3 ± 0.1, fMSE 392 ± 7; classical color transfer (CT) scores MSE 284 ± 6.9, PSNR 27.5 ± 0.1, fMSE 1836 ± 23; the unharmonized input scores MSE 182 ± 5, PSNR 31 ± 0.1, fMSE 984 ± 17.

• Human observers rate the method on par with leading baselines on real AR data. In a four-way user study (Harmonizer, PCT-Net, INR-Harmonization, and MKL encoder) on ARCore images, 20 participants each graded around 30 image sets, producing 642 grades. The authors state their method is rated perceptually on par with the leading baselines, and that it achieves both a high perceptual score and the fastest performance.

• MSE-based metrics misrepresent perceptual quality in this setting. The paper points to the INR model, which has notably lower MSE than PCT-Net, yet is rated much higher by human observers, and argues this gap follows from the training-inference mismatch caused by mask imperfections in standard datasets.

• The "exposure bias" explanation. Masks in datasets such as iHarmony4 often include background pixels near the object boundary, letting models learn a shortcut by comparing leaked background inside the mask with visible background outside it. When a 3D engine produces a pixel-perfect mask, that cue disappears and performance can degrade significantly.

• Speed advantage grows with resolution. On an NVIDIA RTX 4060Ti, the method runs at 175.01 iterations per second at 256x256, 166.76 at 512x512, 137.21 at 1024x2048, and 40.85 at 4096x4096. For comparison, PCT-Net reaches 104.57, 98.65, 63.74, and 11.84; Harmonizer 95.01, 89.82, 47.63, and 7.45; INR 6.35, 3.22, 0.81, and 0.12; DoveNet is reported only at 256x256 with 123.39.

• Real-time on-device inference is achieved. Using the LiteRT Next Kotlin API on Google Pixel 4a and Google Pixel 7, harmonization runs at 12 to 15 fps. Because the tested implementation makes two unnecessary passes through CPU buffers, the authors estimate zero-copy routines could double this to a 24–30 fps range.

• Filter-based methods avoid upscaling artifacts. At high resolution, the authors observe coarse stripes in PCT-Net output and pixelated, JPEG-like areas in INR output, while Harmonizer and their method — both filter-based — show no such defects. All models show biases: Harmonizer tends to brighten images, while their predictions tend to be darker.

• Known failure case for the linear assumption. Dark objects, whose color distribution is concentrated near a corner of the color cube, often produce implausible results because the clipping error is no longer negligible.

Methodology in Plain English

The intuition is that matching the color statistics of an inserted object to its surroundings can be treated as an optimal transport problem: find the cheapest way to move one color distribution onto another. When both distributions are approximated as Gaussians, the optimal map has a closed-form linear solution — the Monge-Kantorovich Linear filter — completely defined by the source and target means and covariance matrices. Because the target distribution is unknown for a newly inserted object, the authors train a neural network to predict it.

To create training targets, they first computed the exact MKL transform per image in iHarmony4 that optimally maps the unharmonized colors to the ground-truth colors. These precomputed 12-dimensional vectors serve as the supervisory signal. The network is an EfficientNet-B0, also used by PCT-Net and Harmonizer, with its first convolutional layer modified to accept four input channels (RGB plus mask), trained on 256x256 inputs and outputting a 12-dimensional vector.

The authors test two prediction targets: predicting the target statistics (mean and covariance), or directly predicting the matrix A and shift S. They find direct prediction of A and S more robust and lower-loss. The final loss combines an L1 loss on the filter parameters with an L1 per-pixel content loss on the masked composited output, weighted by alpha = 10. L1 is chosen over L2 because the problem is ill-posed: L2 would push the model toward the arithmetic mean of all valid solutions, while L1 permits any valid solution and yields sharper results. Content loss alone causes the model to collapse toward identity transforms, so the label loss is necessary.

Training used 210 epochs on an NVIDIA RTX 4090 with the Adam optimizer, learning rate starting at 1e-3, reduced at 50 epochs to 1e-4 and at 110 epochs to 5e-5, batch size 64. For evaluation, the authors report MSE, PSNR, and fMSE on the iHarmony test set at 256x256, but place the most weight on a user study using ARCore data, since real AR composites have no ground truth. They also re-assembled iHarmony4 into an artifact-free PNG version ("iHarmony4-clean") after finding high-frequency artifacts in unmasked regions that distorted MSE and PSNR, and evaluate all baselines on this cleaned set. Theoretically, they prove that clipping colors back into the [0,1] cube equals Euclidean projection, then bound expected squared error as the sum of a clipping term and a linearity term.

Why This Matters

Impact on research. The paper challenges the assumption that harmonization needs dense pixel-level or high-capacity models. By showing that an ideal linear optimal transport filter scores far better than baselines on iHarmony4 (MSE 7.6 ± 0.2 versus 29.1 ± 0.9 for PCT-Net), it reframes the problem as one of predicting a simple transform rather than learning a complex one. It also introduces a critical methodological point — mask leakage in standard benchmarks creates an exposure bias that can make MSE-style rankings misleading — and backs it with a new pixel-perfect-mask AR dataset.

Real-world applications:

• Augmented reality apps and XR headsets that insert virtual objects into live camera feeds and need frame-rate harmonization.
• Mobile commerce and virtual try-on, where a product or garment rendered over a camera view must match ambient lighting.
• Social media camera filters and lens effects on phones.
• Any embedded or battery-constrained device requiring photorealistic compositing without server round-trips.

Industry relevance. The authors are affiliated with Magicly AI (Dubai, UAE) and Glam AI (San Francisco, USA), and the paper explicitly targets mainstream AR toolchains — ARKit, ARCore, Meta Spark, Snap Lens Studio — which today rely on light estimation (direction and intensity, environmental cube maps, spherical harmonic lighting, global exposure or white-balance adjustments) rather than advanced color harmonization, primarily due to latency and mobile compute limits. A method that runs at 12–15 fps on Google Pixel 4a and Pixel 7, with a plausible path to 24–30 fps, directly addresses that barrier.

Future Directions

• Fixing the dark-object failure mode. The authors note that when color distributions concentrate near the corner of the gamut — dark objects — clipping error becomes non-negligible and results are often implausible. A targeted remedy is not proposed.
• Better video and temporal consistency. The model is not trained on video, so sequential frame predictions can vary significantly. The current mitigation is an exponential moving average applied to outputs; the authors state video harmonization remains a challenging problem overall.
• Eliminating the CPU copies in the mobile pipeline. The measured implementation makes two unnecessary passes through CPU buffers; removing them via zero-copy routines is projected to roughly double the frame rate, but was not implemented.
• Expanding the AR dataset. The released set has 327 composite-mask pairs, and the authors explicitly release the data-gathering toolkit to support further acquisition by other researchers.

Target Audience

Researchers and engineers working on image harmonization, computational photography, and augmented reality graphics pipelines; practitioners deploying neural networks on mobile and XR hardware under latency constraints; and anyone studying optimal transport applications in vision or evaluating benchmark validity, particularly the mask-quality and exposure-bias issues in iHarmony4-style datasets. Readers interested in a concrete example of trading raw MSE for perceptual quality and inference speed will also find the benchmarking design useful.

Authors’ abstract

Color harmonization adjusts the colors of an inserted object so that it perceptually matches the surrounding image, resulting in a seamless composite. The harmonization problem naturally arises in augmented reality (AR), yet harmonization algorithms are not currently integrated into AR pipelines because real-time solutions are scarce. In this work, we address color harmonization for AR by proposing a lightweight approach that supports on-device inference. For this, we leverage classical optimal transport theory by training a compact encoder to predict the Monge-Kantorovich transport map. We benchmark our MKL-Harmonizer algorithm against state-of-the-art methods and demonstrate that for real composite AR images our method achieves the best aggregated score. We release our dedicated AR dataset of composite images with pixel-accurate masks and data-gathering toolkit to support further data acquisition by researchers.