CGImage+FaceCrop.swift
2.66 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
//
// CGImage+FaceCrop.swift
//
// Created by Kirill Kudaev on 9/9/19.
// Copyright © 2019 Ancestry. All rights reserved.
//
import Foundation
import Vision
public extension CGImage {
@available(iOS 11.0, *)
func faceCrop(margin: CGFloat = 200, completion: @escaping (FaceCropResult) -> Void) {
let req = VNDetectFaceRectanglesRequest { request, error in
if let error = error {
completion(.failure(error))
return
}
guard let results = request.results, !results.isEmpty else {
completion(.notFound)
return
}
var faces: [VNFaceObservation] = []
for result in results {
guard let face = result as? VNFaceObservation else { continue }
faces.append(face)
}
let croppingRect = self.getCroppingRect(for: faces, margin: margin)
let faceImage = self.cropping(to: croppingRect)
guard let result = faceImage else {
completion(.notFound)
return
}
completion(.success(result))
}
do {
try VNImageRequestHandler(cgImage: self, options: [:]).perform([req])
} catch let error {
completion(.failure(error))
}
}
@available(iOS 11.0, *)
private func getCroppingRect(for faces: [VNFaceObservation], margin: CGFloat) -> CGRect {
var totalX = CGFloat(0)
var totalY = CGFloat(0)
var totalW = CGFloat(0)
var totalH = CGFloat(0)
var minX = CGFloat.greatestFiniteMagnitude
var minY = CGFloat.greatestFiniteMagnitude
let numFaces = CGFloat(faces.count)
for face in faces {
let w = face.boundingBox.width * CGFloat(width)
let h = face.boundingBox.height * CGFloat(height)
let x = face.boundingBox.origin.x * CGFloat(width)
let y = (1 - face.boundingBox.origin.y) * CGFloat(height) - h
totalX += x
totalY += y
totalW += w
totalH += h
minX = .minimum(minX, x)
minY = .minimum(minY, y)
}
let avgX = totalX / numFaces
let avgY = totalY / numFaces
let avgW = totalW / numFaces
let avgH = totalH / numFaces
let offset = margin + avgX - minX
return CGRect(x: avgX - offset, y: avgY - offset, width: avgW + (offset * 2), height: avgH + (offset * 2))
}
}
public enum FaceCropResult {
case success(CGImage)
case notFound
case failure(Error)
}