Swift - 使用 CoreML 删除图像背景答案

【问题标题】：Swift - Remove image background with CoreMLSwift - 使用 CoreML 删除图像背景
【发布时间】：2022-06-30 18:09:32
【问题描述】：

我正在使用 CoreML 和 DeepLabV3 模型从图像中移除背景：https://developer.apple.com/machine-learning/models/

这对于从照片中去除人/狗/汽车的背景效果很好，但对于其他情况，例如天际线和桌子上的某些物体（请参阅示例图像），它无法检测到背景中的对象。

我应该为此使用不同的方法吗？

谢谢

var imageSegmentationModel = DeepLabV3()
var request :  VNCoreMLRequest?

func setUpModel() {
        if let visionModel = try? VNCoreMLModel(for: imageSegmentationModel.model) {
                request = VNCoreMLRequest(model: visionModel, completionHandler: visionRequestDidComplete)
                request?.imageCropAndScaleOption = .scaleFill
        }
        else {
                fatalError()
        }
}

func predict() {
        DispatchQueue.global(qos: .userInitiated).async {
            guard let request = self.request else { fatalError() }
            let handler = VNImageRequestHandler(cgImage: (self.originalImage?.cgImage)!, options: [:])
            do {
                try handler.perform([request])
            }catch {
                print(error)
            }
        }
   }

func visionRequestDidComplete(request: VNRequest, error: Error?) {
        DispatchQueue.main.async {
            if let observations = request.results as? [VNCoreMLFeatureValueObservation],
                let segmentationmap = observations.first?.featureValue.multiArrayValue {
                
                self.maskImage = segmentationmap.image(min: 0, max: 255)
                print(self.maskImage!.size)
                
                self.maskImage = self.maskImage?.resizedImage(for: self.originalImage!.size)
                if let image:UIImage = self.maskOriginalImage(){
                    print("Success")
                    self.outputImageView.image = image
                }
            }
        }
            
    }

func maskOriginalImage() -> UIImage? {
        if(self.maskImage != nil && self.originalImage != nil){
            let maskReference = self.maskImage?.cgImage!
            let imageMask = CGImage(maskWidth: maskReference!.width,
                                    height: maskReference!.height,
                                    bitsPerComponent: maskReference!.bitsPerComponent,
                                    bitsPerPixel: maskReference!.bitsPerPixel,
                                    bytesPerRow: maskReference!.bytesPerRow,
                                    provider: maskReference!.dataProvider!, decode: nil, shouldInterpolate: true)
            
            let maskedReference = self.originalImage?.cgImage!.masking(imageMask!)
            return UIImage(cgImage: maskedReference!)
            
        }
        return nil
    }

【问题讨论】：

标签： ios swift coreml apple-vision deeplab

【解决方案1】：

如果您查看 DeepLabV3 模型支持的标签，它们包括 ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningTable", "dog", "horse", "motorbike", "person", "pottedPlant", "sheep", "sofa", "train", "tvOrMonitor"]。该模型未经过训练以分割 macbook。

您可以训练自己的模型以包含可取的对象。或者看看saliency 的使用，它在检测图像中的主要对象方面做得不错。

【讨论】：