swiftUI coreml deeplabv3去除背景

news2024/11/24 10:31:31

现在手机的性能越来越好，好多深度学习的框架都能能够跑在手机上。因此就集成一下一个官方的深度学习model试一下。
其他的框架生成的模型都能通过相应的工具转换成mlmodel用，转换也比较简单。
下面以替换图像去背景为例，不过官方模型这个效果一般，当个例子实践一下吧。
效果如下：
请添加图片描述

首先用到了两个东西需要下载一下：

DeepLabV3.mlmodel 下载地址
github方便处理uiimage到pixbuffer的一个三方库 CoreMLHelpers 地址

接下来把 deeplabv3直接拖到项目中，看一下详情显示如下：
请添加图片描述
可以看出其入参数是 513*513 的图像，结果也是513x513的灰度图。

通过这个可以看出其中是人的标签是15，我们需要的就是这个标签，其他的标签不需要。
把 CoreMLHelpers 中的同名文件夹拖到项目当中，其中主要是一些方便图像到数据转换的工具。

因此我们需要做的就是如下几步：

把图片缩放成 513x513 的图像；
送入模型处理；
通过人物结果标签生成蒙版图；
图像+蒙版显示结果；

页面代码如下：

import SwiftUI

// 图像增加蒙版功能
func maskImage(image: UIImage, mask: UIImage) -> UIImage {
    let imageRef = image.cgImage!
    let maskRef = mask.cgImage!
    let maskedRef = imageRef.masking(maskRef)
    return UIImage.init(cgImage: maskedRef!)
}

struct ContentView: View {
    @State var previewImage: Image = Image("sample")
    
    var body: some View {
        VStack(spacing: 0) {
            previewImage
                .resizable()
                .scaledToFit()
                .frame(width: 260)
            
            VStack {
                
            }
            .frame(height: 20)
            
            Button {
                guard let model = try? DeepLabV3(configuration: .init()) else {
                    return
                }
                
                print("model loaded \(model)")
                
                guard let image = UIImage(named: "sample") else {
                    return
                }
                
                let out = try? model.prediction(image: image.pixelBuffer(width: 513, height: 513)!)
                
                print("prediection result \(out)")
                
                print("prediection \(out?.semanticPredictions)")
                
                let mask = out?.semanticPredictions.image(channel: 1)
                
                previewImage = Image(uiImage: maskImage(image: .init(named: "sample")!, mask: mask!))
            } label: {
                Text("去背景")
                    .tint(.white)
            }
            .padding([.top, .bottom], 10)
            .padding([.leading, .trailing], 20)
            .background {
                Color.green
            }
            .cornerRadius(10)
            
            
        }
        .padding()
        
    }
}

其中为了选取图片中标签为 15 的人物标签，修改了一下 MLMultiArray+Image.swift 中 image 方法的实现。

    // Loop through all the pixels and all the channels and copy them over.
    for c in 0..<channels {
      for y in 0..<height {
        for x in 0..<width {
          let value = ptr[c*cStride + y*yStride + x*xStride]
//          let scaled = (value - min) * T(255) / (max - min)
//          let pixel = clamp(scaled, min: T(0), max: T(255)).toUInt8
            let pixel = (value == T(15) ? 255: 0).toUInt8
          pixels[(y*width + x)*bytesPerPixel + c] = pixel
        }
      }
    }