Pre-resize images to exact model dimensions, implement feathered blending to eliminate seam artifacts, cache model constraints to unify coordinate space

DepthPredictor.swift

@@ -138,9 +138,17 @@ final class DepthPredictor {
     private var visionModel: VNCoreMLModel?
     private var _outputHeight = 512
     private var _outputWidth = 1024
+    private var _modelInputWidth: Int = 0
+    private var _modelInputHeight: Int = 0
 
     var outputHeight: Int { _outputHeight }
     var outputWidth: Int { _outputWidth }
+    /// Model's expected input dimensions, read from the CoreML model's image
+    /// constraints at load time. Used to manually resize source images so that
+    /// Vision's `.scaleFit` becomes a no-op (no letterboxing, no implicit
+    /// bilinear downscale). Zero if the model isn't loaded.
+    var modelInputWidth: Int { _modelInputWidth }
+    var modelInputHeight: Int { _modelInputHeight }
     var isLoaded: Bool { visionModel != nil }
 
     /// Load model dynamically from a .mlpackage or .mlmodelc URL.
@@ -174,6 +182,15 @@ final class DepthPredictor {
     }
 
     /// Run a single pass of depth inference on a CGImage.
+    ///
+    /// The image is resized to the model's expected input dimensions using
+    /// high-quality interpolation *before* being handed to Vision. This makes
+    /// `imageCropAndScaleOption = .scaleFit` effectively a no-op and avoids
+    /// two failure modes of letting Vision do the resize:
+    ///   1. Letterboxing on inputs whose aspect ratio doesn't exactly match
+    ///      the model (Vision pads with black, polluting depth predictions).
+    ///   2. Implicit bilinear downscale, which loses high-frequency detail
+    ///      compared to PIL's Lanczos resize used in the Python export script.
     private func runSingleInference(on cgImage: CGImage, completion: @escaping (DepthResult?) -> Void) {
         guard let visionModel else {
             print("[DepthPredictor] Model not loaded")
@@ -181,6 +198,18 @@ final class DepthPredictor {
             return
         }
 
+        // Pre-resize to exact model input dims (matches Python's PIL resize).
+        let prepared: CGImage
+        if _modelInputWidth > 0 && _modelInputHeight > 0,
+           let resized = DepthPredictor.resizeImage(cgImage,
+                                                    toWidth: _modelInputWidth,
+                                                    height: _modelInputHeight) {
+            prepared = resized
+        } else {
+            // Fallback: model dims unknown — let Vision handle scaling.
+            prepared = cgImage
+        }
+
         let request = VNCoreMLRequest(model: visionModel) { [weak self] request, error in
             if let error {
                 print("[DepthPredictor] Inference error: \(error)")
@@ -205,7 +234,7 @@ final class DepthPredictor {
 
         request.imageCropAndScaleOption = .scaleFit
 
-        let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])
+        let handler = VNImageRequestHandler(cgImage: prepared, options: [:])
         do {
             try handler.perform([request])
         } catch {
@@ -237,11 +266,25 @@ final class DepthPredictor {
         patchHalfWidth: Int = 25,
         completion: @escaping (DepthResult?) -> Void
     ) {
-        let imageWidth = cgImage.width
+        // Resize source to model input dims *once*, so both inference passes
+        // and the horizontal shift all happen in the same coordinate space.
+        // This avoids resampling twice and keeps the shift offset exact in
+        // the same pixel grid as the depth output.
+        let prepared: CGImage
+        if _modelInputWidth > 0 && _modelInputHeight > 0,
+           let resized = DepthPredictor.resizeImage(cgImage,
+                                                    toWidth: _modelInputWidth,
+                                                    height: _modelInputHeight) {
+            prepared = resized
+        } else {
+            prepared = cgImage
+        }
+
+        let imageWidth = prepared.width
         let half = imageWidth / 2
 
         // Shift the source image left by half — the seam moves to the center
-        guard let shiftedImage = DepthPredictor.shiftImageHorizontally(cgImage, by: half) else {
+        guard let shiftedImage = DepthPredictor.shiftImageHorizontally(prepared, by: half) else {
             print("[DepthPredictor] Failed to shift image for seam fix")
             completion(nil)
             return
@@ -255,8 +298,8 @@ final class DepthPredictor {
             )
         }
 
-        // 1. Infer depth on the original image
-        runSingleInference(on: cgImage) { [weak self] originalDepth in
+        // 1. Infer depth on the (resized) original image
+        runSingleInference(on: prepared) { [weak self] originalDepth in
             guard let self, let originalDepth else {
                 completion(nil)
                 return
@@ -314,21 +357,36 @@ final class DepthPredictor {
         }
     }
 
-    /// Stitch the seam region using a single output buffer — no intermediate copies.
+    /// Stitch the seam region using a single output buffer with **feathered**
+    /// blending at the patch boundaries — no intermediate copies.
+    ///
+    /// The two inference passes (original and half-shifted) produce slightly
+    /// different absolute depth values even where they agree on geometry,
+    /// because they're independent forward passes through a non-linear model.
+    /// A hard cutover at the patch boundary therefore leaves a visible step.
+    /// To avoid this, we linearly blend from original→shifted as the column
+    /// enters the patch zone and from shifted→original as it leaves, using a
+    /// transition band of `featherWidth` pixels on each side.
     ///
-    /// For each output pixel (row, col), we determine whether it falls in the
-    /// patch zone (the strip around the original seam at column `depthHalf`).
-    /// If so, we read from the shifted depth at the corresponding shifted
-    /// column; otherwise we read from the original depth at `col` directly.
+    /// Layout in *shifted* coordinate space (centered at width/2):
     ///
-    /// This replaces the prior 3-step roll→patch→unroll with 4 temporary arrays.
+    ///     [ original ][ feather ][ shifted ][ feather ][ original ]
+    ///                ^          ^          ^          ^
+    ///         patchLeft   coreLeft    coreRight   patchRight
+    ///
+    /// - Outside `[patchLeft, patchRight)`: pure original.
+    /// - Inside `[coreLeft, coreRight)`: pure shifted.
+    /// - In the two feather bands: linear blend, weight 0→1 across the band.
+    ///
+    /// `featherWidth` is clamped so the feather bands never overlap the core.
     private func stitchSeamFromShiftedDepth(
         original: MLMultiArray,
         shifted: MLMultiArray,
         width: Int,
         height: Int,
         depthHalf: Int,
-        patchHalfWidth: Int
+        patchHalfWidth: Int,
+        featherWidth: Int = 12
     ) -> MLMultiArray? {
         let origView = DepthArrayView(original)
         let shiftView = DepthArrayView(shifted)
@@ -339,6 +397,13 @@ final class DepthPredictor {
         let patchLeft = centerX - dx
         let patchRight = centerX + dx  // exclusive
 
+        // Clamp feather so the two bands don't overlap (each band must fit
+        // within half the patch width, leaving at least one pure-shifted col).
+        let maxFeather = max(0, dx - 1)
+        let feather = min(max(0, featherWidth), maxFeather)
+        let coreLeft = patchLeft + feather
+        let coreRight = patchRight - feather  // exclusive
+
         // Create output MLMultiArray
         let output: MLMultiArray
         do {
@@ -351,6 +416,9 @@ final class DepthPredictor {
         let outStride = output.strides[2].intValue
         let outPtr = output.dataPointer.bindMemory(to: Float32.self, capacity: width * height)
 
+        // Precompute reciprocal once (avoid div-by-zero when feather == 0).
+        let invFeather: Float32 = feather > 0 ? 1.0 / Float32(feather) : 0.0
+
         for row in 0..<height {
             let outBase = row * outStride
             for col in 0..<width {
@@ -358,12 +426,27 @@ final class DepthPredictor {
                 // shifting left by depthHalf means shiftedCol = (col + depthHalf) % width
                 let shiftedCol = (col + depthHalf) % width
 
-                if shiftedCol >= patchLeft && shiftedCol < patchRight {
-                    // This pixel is in the patch zone — use shifted depth
+                if shiftedCol < patchLeft || shiftedCol >= patchRight {
+                    // Outside patch zone — pure original (identity mapping).
+                    outPtr[outBase + col] = origView.value(row: row, col: col)
+                } else if shiftedCol >= coreLeft && shiftedCol < coreRight {
+                    // Core patch zone — pure shifted.
                     outPtr[outBase + col] = shiftView.value(row: row, col: shiftedCol)
                 } else {
-                    // Outside patch — use original depth (identity mapping)
-                    outPtr[outBase + col] = origView.value(row: row, col: col)
+                    // Feather band — linear blend.
+                    // Weight w: 0 at the outer patch edge, 1 at the core edge.
+                    let w: Float32
+                    if shiftedCol < coreLeft {
+                        // Left feather: ramp up as we move right toward coreLeft.
+                        w = Float32(shiftedCol - patchLeft) * invFeather
+                    } else {
+                        // Right feather: ramp down as we move right toward patchRight.
+                        w = Float32(patchRight - 1 - shiftedCol) * invFeather
+                    }
+                    let wClamped = max(0.0 as Float32, min(1.0 as Float32, w))
+                    let origVal = origView.value(row: row, col: col)
+                    let shiftVal = shiftView.value(row: row, col: shiftedCol)
+                    outPtr[outBase + col] = origVal + (shiftVal - origVal) * wClamped
                 }
             }
         }
@@ -512,6 +595,20 @@ final class DepthPredictor {
 
             let compiledURL = try compileModelIfNeeded(at: modelURL)
             let model = try MLModel(contentsOf: compiledURL, configuration: config)
+
+            // Capture the model's expected input dimensions so we can resize
+            // source images ourselves (avoiding Vision's letterboxing + implicit
+            // bilinear downscale). DAP exports use a single ImageType input.
+            if let imageInput = model.modelDescription.inputDescriptionsByName.values
+                .first(where: { $0.imageConstraint != nil }),
+               let constraint = imageInput.imageConstraint {
+                _modelInputWidth = constraint.pixelsWide
+                _modelInputHeight = constraint.pixelsHigh
+                print("[DepthPredictor] Model input: \(_modelInputWidth)x\(_modelInputHeight)")
+            } else {
+                print("[DepthPredictor] Warning: could not read model input image constraint; manual resize disabled")
+            }
+
             visionModel = try VNCoreMLModel(for: model)
 
             print("[DepthPredictor] Model loaded from \(modelURL.path)")
@@ -661,6 +758,42 @@ extension DepthPredictor {
         return cgImage
     }
 
+    /// Resize a CGImage to exact `(width, height)` using high-quality
+    /// interpolation (Lanczos-equivalent on macOS). Returns nil if context
+    /// creation fails.
+    ///
+    /// This is used to pre-resize the source image to the model's expected
+    /// input dimensions *before* handing off to Vision. Doing so makes
+    /// `imageCropAndScaleOption = .scaleFit` a no-op — no letterboxing on
+    /// non-matching aspect ratios, and no implicit bilinear downscale.
+    static func resizeImage(_ cgImage: CGImage, toWidth width: Int, height: Int) -> CGImage? {
+        guard width > 0, height > 0 else { return nil }
+        if cgImage.width == width && cgImage.height == height {
+            return cgImage
+        }
+
+        let colorSpace = CGColorSpaceCreateDeviceRGB()
+        let bitmapInfo = CGBitmapInfo.byteOrder32Little.rawValue
+            | CGImageAlphaInfo.noneSkipLast.rawValue
+
+        guard let ctx = CGContext(
+            data: nil,
+            width: width,
+            height: height,
+            bitsPerComponent: 8,
+            bytesPerRow: 0,
+            space: colorSpace,
+            bitmapInfo: bitmapInfo
+        ) else {
+            print("[DepthPredictor] resizeImage: CGContext creation failed")
+            return nil
+        }
+
+        ctx.interpolationQuality = .high
+        ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
+        return ctx.makeImage()
+    }
+
     /// Save a CIImage as a PNG file (renders via CIContext first).
     static func saveImage(_ ciImage: CIImage, to path: URL) throws {
         let context = CIContext()