Optimizing Image Inpainting Pipelines: The Importance of Metrics and Blending in AI Systems
This article from the ClearPix engineering series details a bug in their AI-powered watermark removal pipeline, specifically within the final blending step. It highlights the critical importance of designing precise, region-specific metrics for generative AI systems to accurately identify and resolve artifacts, rather than relying on averaged metrics that can hide critical issues. The solution involved refining the image blending logic to prevent model-generated content from polluting clean image areas, leading to significant quality improvements.
Read original on Dev.to #architectureThe Challenge: 'Processed Look' in Watermark Removal
The ClearPix team encountered a persistent user complaint: images, after watermark removal, had a 'processed look' characterized by a soft, smeared halo. Initially, the generative AI model for inpainting was suspected. However, the root cause was discovered to be in the final *blending step* of the pipeline, where the inpainted region was merged back into the original image. The issue arose because the inpainting model reconstructs the *entire crop*, not just the masked hole. Consequently, the blending logic inadvertently mixed slightly erroneous model output into pixels that were originally clean, especially in a feathered transition band around the watermark.
Flawed Metrics Obscured the Problem
A key architectural learning point from this case study is the danger of overly generalized quality metrics in complex pipelines. ClearPix's original `psnrRegion` metric averaged the quality across both the inpainted hole and its surrounding ring. This averaging hid the degradation occurring in the 'clean' outer ring, making it difficult to pinpoint the blending as the source of the artifacts. The problem was especially pronounced for small watermarks, where the affected transition band was several times larger than the watermark itself.
System Design Lesson: Granular Metrics for Generative AI
When designing generative AI systems, especially those involving image or data manipulation, it is crucial to implement granular, region-specific quality metrics. Splitting metrics by 'inside mask' vs. 'outside mask' or 'restored region' vs. 'clean region' can reveal subtle yet significant issues that averaged metrics would miss. This allows for precise identification of problem areas, whether in the model, post-processing, or blending stages.
Introducing Split Metrics
- PSNR-in / PSNR-out: Separate PSNR calculations for inside the mask and the outer ring.
- ringDelta: Measures the mean absolute difference in the outer ring; ideally 0 for untouched pixels.
- colorShift: Lab ΔE between the filled hole and the surrounding ring, indicating color mismatch.
- SSIM (region): Structural Similarity Index, which often correlates better with human perception than PSNR.
The Architectural Fix: Precise Blending
The solution involved a small but critical change to the blending algorithm. Instead of using a broad Gaussian-blurred feather band, the updated logic ensures that 100% model output is used *only* inside the mask. Outside the mask, the artificial feather band is deleted, and blending is restricted to a minimal, pixel-or-two wide ramp. This ensures that pixels outside the immediate mask boundary remain bit-identical to the original input, preventing pollution from the model's reconstruction of 'clean' areas.
Old blending:
out = original + a * (inpaintedCrop - original)
a = max(binary mask, blur(mask, 6-22px))
New blending (conceptual):
out = original, where mask is 0
out = inpaintedCrop, where mask is 1 (with minimal feathering)This refinement resulted in `ringDelta` dropping to zero for all cases, and `PSNR-out` dramatically increasing, confirming that clean pixels were no longer being touched. This demonstrates how a seemingly minor detail in a pipeline's final stage can have a cascading impact on overall system quality, especially when dealing with generative output.