Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

✨ Improve Error Handling and Image Conversion in removeBackground Method #9

Merged
merged 1 commit into from
Sep 27, 2024
Merged

✨ Improve Error Handling and Image Conversion in removeBackground Method #9

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 40 additions & 15 deletions packages/react-native-background-remover/ios/ReactNativeBackgroundRemover.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}
}
}
Loading