版本记录

版本号	时间
V1.0	2022.03.20 星期日

前言

iOS 11+和macOS 10.13+ 新出了Vision框架，提供了人脸识别、物体检测、物体跟踪等技术，它是基于Core ML的。可以说是人工智能的一部分，接下来几篇我们就详细的解析一下Vision框架。感兴趣的看下面几篇文章。
1. Vision框架详细解析（一） —— 基本概览（一）
2. Vision框架详细解析（二） —— 基于Vision的人脸识别（一）
3. Vision框架详细解析（三） —— 基于Vision的人脸识别（二）
4. Vision框架详细解析（四） —— 在iOS中使用Vision和Metal进行照片堆叠（一）
5. Vision框架详细解析（五） —— 在iOS中使用Vision和Metal进行照片堆叠（二）
6. Vision框架详细解析（六） —— 基于Vision的显著性分析（一）
7. Vision框架详细解析（七） —— 基于Vision的显著性分析（二）
8. Vision框架详细解析（八） —— 基于Vision的QR扫描（一）
9. Vision框架详细解析（九） —— 基于Vision的QR扫描（二）
10. Vision框架详细解析（十） —— 基于Vision的Body Detect和Hand Pose（一）
11. Vision框架详细解析（十一） —— 基于Vision的Body Detect和Hand Pose（二)
12. Vision框架详细解析（十二） —— 基于Vision的Face Detection新特性（一)
13. Vision框架详细解析（十三） —— 基于Vision的Face Detection新特性（二)
14. Vision框架详细解析（十四） —— 基于Vision的人员分割（一)

源码

首先看下工程组织结构

下面就是源码了

1. GreetingProcessor.swift

import UIKit
import Combine
import Vision
import CoreImage.CIFilterBuiltins

class GreetingProcessor {
  static let shared = GreetingProcessor()
  @Published var photoOutput = UIImage()
  let context = CIContext()
  let request = VNGeneratePersonSegmentationRequest()

  func generatePhotoGreeting(greeting: Greeting) {
    guard
      let backgroundImage = greeting.backgroundImage.cgImage,
      let foregroundImage = greeting.foregroundImage.cgImage else {
      print("Missing required images")
      return
    }

    // Create request
    request.qualityLevel = .accurate
    request.revision = VNGeneratePersonSegmentationRequestRevision1
    request.outputPixelFormat = kCVPixelFormatType_OneComponent8

    // Create request handler
    let requestHandler = VNImageRequestHandler(
      cgImage: foregroundImage,
      options: [:])

    do {
      // Process request
      try requestHandler.perform([request])
      guard let mask = request.results?.first else {
        print("Error generating person segmentation mask")
        return
      }

      let foreground = CIImage(cgImage: foregroundImage)
      let maskImage = CIImage(cvPixelBuffer: mask.pixelBuffer)
      let background = CIImage(cgImage: backgroundImage)
      
      guard let output = blendImages(
        background: background,
        foreground: foreground,
        mask: maskImage) else {
          print("Error blending images")
          return
      }
      // Update photoOutput
      if let photoResult = renderAsUIImage(output) {
        self.photoOutput = photoResult
      }
    } catch {
      print("Error processing person segmentation request")
    }
  }

  func blendImages(
    background: CIImage,
    foreground: CIImage,
    mask: CIImage,
    isRedMask: Bool = false
  ) -> CIImage? {
    // scale mask
    let maskScaleX = foreground.extent.width / mask.extent.width
    let maskScaleY = foreground.extent.height / mask.extent.height
    let maskScaled = mask.transformed(by: __CGAffineTransformMake(maskScaleX, 0, 0, maskScaleY, 0, 0))

    // scale background
    let backgroundScaleX = (foreground.extent.width / background.extent.width)
    let backgroundScaleY = (foreground.extent.height / background.extent.height)
    let backgroundScaled = background.transformed(
      by: __CGAffineTransformMake(backgroundScaleX, 0, 0, backgroundScaleY, 0, 0))

    let blendFilter = isRedMask ? CIFilter.blendWithRedMask() : CIFilter.blendWithMask()
    blendFilter.inputImage = foreground
    blendFilter.backgroundImage = backgroundScaled
    blendFilter.maskImage = maskScaled

    return blendFilter.outputImage
  }
  
  private func renderAsUIImage(_ image: CIImage) -> UIImage? {
    guard let cgImage = context.createCGImage(image, from: image.extent) else {
      return nil
    }
    return UIImage(cgImage: cgImage)
  }
  
  func processVideoFrame(
    foreground: CVPixelBuffer,
    background: CGImage
  ) -> CIImage? {
    // Create request handler
    let ciForeground = CIImage(cvPixelBuffer: foreground)
    let personSegmentFilter = CIFilter.personSegmentation()
    personSegmentFilter.inputImage = ciForeground
    if let mask = personSegmentFilter.outputImage {
      guard let output = blendImages(
        background: CIImage(cgImage: background),
        foreground: ciForeground,
        mask: mask,
        isRedMask: true) else {
          print("Error blending images")
          return nil
        }
      return output
    }
    return nil
  }
}

2. CameraViewController.swift

import UIKit
import AVFoundation
import MetalKit

class CameraViewController: UIViewController {
  // AV Foundation
  let captureSession = AVCaptureSession()
  let cameraView = MTKView()
  var background: UIImage?
  var cameraFrameAdded = false

  // swiftlint:disable implicitly_unwrapped_optional
  public var metalDevice: MTLDevice!
  public var metalCommandQueue: MTLCommandQueue!
  public var ciContext: CIContext!
  // swiftlint:enable implicitly_unwrapped_optional

  public var currentCIImage: CIImage? {
    didSet {
      cameraView.draw()
    }
  }

  override func viewDidLoad() {
    super.viewDidLoad()
    setupMetal()
    setupCoreImage()
    setupSession()
  }

  override func viewWillAppear(_ animated: Bool) {
    super.viewWillAppear(animated)
    if !cameraFrameAdded {
      cameraView.frame = view.frame
      view.addSubview(cameraView)
      cameraFrameAdded.toggle()
    }
    startSession()
  }

  override func viewWillDisappear(_ animated: Bool) {
    stopSession()
    super.viewWillDisappear(animated)
  }

  func setupSession() {
    captureSession.beginConfiguration()
    guard let camera = AVCaptureDevice.default(.builtInWideAngleCamera, for: .video, position: .front) else {
      return
    }
    do {
      let videoInput = try AVCaptureDeviceInput(device: camera)
      if captureSession.canAddInput(videoInput) {
        captureSession.addInput(videoInput)
      }
    } catch {
      print("Error setting device input: \(error)")
      return
    }
    let output = AVCaptureVideoDataOutput()
    output.alwaysDiscardsLateVideoFrames = true
    output.setSampleBufferDelegate(self, queue: .main)
    captureSession.addOutput(output)
    output.connections.first?.videoOrientation = .portrait
    captureSession.commitConfiguration()
  }

  func setupMetal() {
    metalDevice = MTLCreateSystemDefaultDevice()
    metalCommandQueue = metalDevice.makeCommandQueue()
    cameraView.device = metalDevice
    cameraView.isPaused = true
    cameraView.enableSetNeedsDisplay = false
    cameraView.delegate = self
    cameraView.framebufferOnly = false
  }

  func setupCoreImage() {
    ciContext = CIContext(mtlDevice: metalDevice)
  }

  func startSession() {
    if !captureSession.isRunning {
      DispatchQueue.global(qos: .default).async { [weak self] in
        self?.captureSession.startRunning()
      }
    }
  }

  func stopSession() {
    if captureSession.isRunning {
      DispatchQueue.global(qos: .default).async { [weak self] in
        self?.captureSession.stopRunning()
      }
    }
  }
}

// MARK: - AVCaptureVideoDataOutputSampleBufferDelegate
extension CameraViewController: AVCaptureVideoDataOutputSampleBufferDelegate {
  func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    // Grab the pixelbuffer frame from the camera output
    guard
      let pixelBuffer = sampleBuffer.imageBuffer,
      let backgroundImage = self.background?.cgImage else {
      return
    }
    DispatchQueue.global().async {
      if let output = GreetingProcessor.shared.processVideoFrame(
        foreground: pixelBuffer,
        background: backgroundImage) {
        DispatchQueue.main.async {
          self.currentCIImage = output
        }
      }
    }
  }
}

// MARK: - MTKViewDelegate
extension CameraViewController: MTKViewDelegate {
  func draw(in view: MTKView) {
    guard
      let commandBuffer = metalCommandQueue.makeCommandBuffer(),
      let ciImage = currentCIImage,
      let currentDrawable = view.currentDrawable else {
      return
    }

    let drawSize = cameraView.drawableSize
    let scaleX = drawSize.width / ciImage.extent.width
    let scaleY = drawSize.height / ciImage.extent.height

    let newImage = ciImage.transformed(by: .init(scaleX: scaleX, y: scaleY))
    // render into the metal texture
    self.ciContext.render(
      newImage,
      to: currentDrawable.texture,
      commandBuffer: commandBuffer,
      bounds: newImage.extent,
      colorSpace: CGColorSpaceCreateDeviceRGB())

    // register drawable to command buffer
    commandBuffer.present(currentDrawable)
    commandBuffer.commit()
  }

  func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
    // No-op
  }
}

3. CameraView.swift

import Foundation
import SwiftUI

struct CameraView: UIViewControllerRepresentable {
  typealias UIViewControllerType = CameraViewController
  private let cameraViewController: CameraViewController

  init(videoBackgroundImage: UIImage?) {
    cameraViewController = CameraViewController()
    cameraViewController.background = videoBackgroundImage
  }


  func makeUIViewController(context: Context) -> CameraViewController {
    cameraViewController
  }

  func updateUIViewController(_ uiViewController: CameraViewController, context: Context) {
  }
}

4. ContentView.swift

import SwiftUI

struct ContentView: View {
  var body: some View {
    TabView {
      PhotoGreetingView()
        .tabItem {
          Image(systemName: "photo")
          Text("Photo Greeting")
        }
        .tag(0)
      VideoGreetingView()
        .tabItem {
          Image(systemName: "video")
          Text("Video Greeting")
        }
        .tag(1)
    }
  }
}

struct ContentView_Previews: PreviewProvider {
  static var previews: some View {
    ContentView()
  }
}

5. PhotoGreetingView.swift

import SwiftUI

struct PhotoGreetingView: View {
  @State private var displayPhotoGreeting = false
  var greeting: Greeting {
    Greeting(backgroundImage: GreetingContent.photoBackground, foregroundImage: GreetingContent.photoForeground)
  }

  var body: some View {
    ScrollView {
      VStack {
        Image(uiImage: greeting.backgroundImage)
          .resizable()
          .aspectRatio(contentMode: .fit)
          .overlay(OverlayContent(title: "Background"), alignment: .bottom)
        Image(uiImage: greeting.foregroundImage)
          .resizable()
          .aspectRatio(contentMode: .fit)
          .overlay(OverlayContent(title: "Family Pic"), alignment: .bottom)
        Button {
          GreetingProcessor.shared.generatePhotoGreeting(greeting: greeting)
          displayPhotoGreeting.toggle()
        } label: {
          Text("Generate Photo Greeting")
            .font(.headline)
            .foregroundColor(Color.blue)
            .padding(/*@START_MENU_TOKEN@*/.all, 10.0/*@END_MENU_TOKEN@*/)
        }
      }
    }
    .sheet(isPresented: $displayPhotoGreeting) {
      PhotoGreetingOutputView()
    }
    .padding()
  }
}

struct PhotoGreetingView_Previews: PreviewProvider {
  static var previews: some View {
    PhotoGreetingView()
  }
}

struct OverlayContent: View {
  let title: String

  var body: some View {
    Text(title)
      .font(.title)
      .fontWeight(.regular)
      .frame(maxWidth: .infinity)
      .foregroundColor(Color.black)
      .background(Color.white)
      .opacity(0.7)
  }
}

struct PhotoGreetingOutputView: View {
  var greetingProcessor = GreetingProcessor.shared

  var body: some View {
    ScrollView {
      VStack {
        Image(uiImage: greetingProcessor.photoOutput)
          .resizable()
          .aspectRatio(contentMode: .fit)
          .overlay(OverlayContent(title: "Output"), alignment: .bottom)
      }
    }
    .padding()
  }
}

6. VideoGreetingView.swift

import SwiftUI

struct VideoGreetingView: View {
  var cameraView = CameraView(videoBackgroundImage: GreetingContent.videoBackground)
  var body: some View {
    cameraView
  }
}

struct VideoGreetingView_Previews: PreviewProvider {
  static var previews: some View {
    VideoGreetingView()
  }
}

7. Greeting.swift

import UIKit

struct Greeting {
  var backgroundImage: UIImage
  var foregroundImage: UIImage
}

8. GreetingContent.swift

import UIKit

enum GreetingContent {
  // swiftlint:disable implicitly_unwrapped_optional
  static let photoBackground: UIImage! = UIImage(named: "Christmas")
  static let photoForeground: UIImage! = UIImage(named: "Family")
  static let videoBackground: UIImage! = UIImage(named: "Thanksgiving")
  // swiftlint:enable implicitly_unwrapped_optional
}

9. AppMain.swift

import SwiftUI

@main
struct AppMain: App {
  var body: some Scene {
    WindowGroup {
      ContentView()
    }
  }
}

后记

本篇主要讲述了基于Vision的人员分割，感兴趣的给个赞或者关注~~~