From 53bc5f250659c1298ea67aaae0f465c988b01f70 Mon Sep 17 00:00:00 2001 From: Enis Gegic Date: Fri, 27 Sep 2024 16:21:25 +0200 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Improve=20error=20handling=20-Added?= =?UTF-8?q?=20check=20to=20avoid=20producing=20images=20with=20large=20siz?= =?UTF-8?q?e=20-Use=20CIContext=20for=20proper=20image=20conversion=20proc?= =?UTF-8?q?ess=20to=20avoid=20crashes=20on=20some=20devices?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ios/ReactNativeBackgroundRemover.swift | 55 ++++++++++++++----- 1 file changed, 40 insertions(+), 15 deletions(-) diff --git a/packages/react-native-background-remover/ios/ReactNativeBackgroundRemover.swift b/packages/react-native-background-remover/ios/ReactNativeBackgroundRemover.swift index 8a88f7d..79005bf 100644 --- a/packages/react-native-background-remover/ios/ReactNativeBackgroundRemover.swift +++ b/packages/react-native-background-remover/ios/ReactNativeBackgroundRemover.swift @@ -4,49 +4,74 @@ import CoreImage public class BackgroundRemoverSwift: NSObject { @objc - public func removeBackground(_ imageURI: String, resolve: @escaping RCTPromiseResolveBlock, reject: @escaping RCTPromiseRejectBlock)->Void { + public func removeBackground(_ imageURI: String, resolve: @escaping RCTPromiseResolveBlock, reject: @escaping RCTPromiseRejectBlock) -> Void { #if targetEnvironment(simulator) reject("BackgroundRemover", "SimulatorError", NSError(domain: "BackgroundRemover", code: 2)) - return + return #endif if #available(iOS 15.0, *) { - let url = URL(string: imageURI)! - let originalImage = CIImage(contentsOf: url, options: [.applyOrientationProperty: true])! + guard let url = URL(string: imageURI) else { + reject("BackgroundRemover", "Invalid URL", NSError(domain: "BackgroundRemover", code: 3)) + return + } + + guard let originalImage = CIImage(contentsOf: url, options: [.applyOrientationProperty: true]) else { + reject("BackgroundRemover", "Unable to load image", NSError(domain: "BackgroundRemover", code: 4)) + return + } + let imageRequestHandler = VNImageRequestHandler(ciImage: originalImage) var segmentationRequest = VNGeneratePersonSegmentationRequest() - segmentationRequest = VNGeneratePersonSegmentationRequest() segmentationRequest.qualityLevel = .accurate segmentationRequest.outputPixelFormat = kCVPixelFormatType_OneComponent8 DispatchQueue.global(qos: .userInitiated).async { do { try imageRequestHandler.perform([segmentationRequest]) - let pixelBuffer = segmentationRequest.results?.first!.pixelBuffer + guard let pixelBuffer = segmentationRequest.results?.first?.pixelBuffer else { + reject("BackgroundRemover", "No segmentation results", NSError(domain: "BackgroundRemover", code: 5)) + return + } - var maskImage = CIImage(cvPixelBuffer: pixelBuffer!) + var maskImage = CIImage(cvPixelBuffer: pixelBuffer) + // Adjust mask scaling let scaleX = originalImage.extent.width / maskImage.extent.width let scaleY = originalImage.extent.height / maskImage.extent.height - maskImage = maskImage.transformed(by: .init(scaleX: scaleX, y: scaleY)) + + // Avoid up-scaling if mask dimensions are too small + if scaleX > 1 || scaleY > 1 { + maskImage = maskImage.transformed(by: CGAffineTransform(scaleX: scaleX, y: scaleY)) + } let maskedImage = originalImage.applyingFilter("CIBlendWithMask", parameters: [kCIInputMaskImageKey: maskImage]) - // Save the masked image to a temporary file - let tempURL = URL(fileURLWithPath: NSTemporaryDirectory()).appendingPathComponent(url.lastPathComponent) - let uiImage = UIImage(ciImage: maskedImage) - if let data = uiImage.pngData() { + // Convert to UIImage via CGImage for better control + let context = CIContext() + guard let cgMaskedImage = context.createCGImage(maskedImage, from: maskedImage.extent) else { + reject("BackgroundRemover", "Error creating CGImage", NSError(domain: "BackgroundRemover", code: 6)) + return + } + + let uiImage = UIImage(cgImage: cgMaskedImage) + + // Save the image as PNG to preserve transparency + let tempURL = URL(fileURLWithPath: NSTemporaryDirectory()).appendingPathComponent(url.lastPathComponent).appendingPathExtension("png") + if let data = uiImage.pngData() { // Use PNG to preserve transparency try data.write(to: tempURL) resolve(tempURL.absoluteString) + } else { + reject("BackgroundRemover", "Error saving image", NSError(domain: "BackgroundRemover", code: 7)) } + } catch { - reject("BackgroundRemover", "[ReactNativeBackgroundRemover]: Error removing background", error) + reject("BackgroundRemover", "Error removing background", error) } } } else { - reject("BackgroundRemover", "[ReactNativeBackgroundRemover]: You need to have a device with at least iOS 15", NSError(domain: "BackgroundRemover", code: 1)) - return + reject("BackgroundRemover", "You need a device with iOS 15 or later", NSError(domain: "BackgroundRemover", code: 1)) } } }