Oregami
Repositories/oxedyne/fe2o3

oxedyne/fe2o3/fe2o3_infer/src/face/align.rs

8.2 KiB, 1 run

created by r1870400018:19747, which is this file's identity for as long as the history lasts, whatever it is later renamed to

download · who wrote it · its history

1//! The similarity transform that puts five facial landmarks on a fixed template,
2//! and the bilinear warp that follows it.
3
4use crate::face::Image;
5
6use oxedyne_fe2o3_core::prelude::*;
7
8/// The five-point template an aligned crop is warped onto: right eye, left eye,
9/// nose tip, right corner of the mouth, left corner of the mouth, in the
10/// hundred and twelve pixel square the embedder consumes.
11pub const TEMPLATE: [(f32, f32); 5] = [
12 (38.2946, 51.6963),
13 (73.5318, 51.5014),
14 (56.0252, 71.7366),
15 (41.5493, 92.3655),
16 (70.7299, 92.2041),
17];
18
19/// The template's own centroid, carried as a constant so that the transform
20/// matches the reference implementation digit for digit.
21const TEMPLATE_MEAN: (f32, f32) = (56.0262, 71.9008);
22
23/// Side of the aligned crop, in pixels.
24pub const CROP: usize = 112;
25
26/// An affine map, `[[a, b, tx], [c, d, ty]]`, taking a source point to a
27/// destination point.
28#[derive(Clone, Copy, Debug, PartialEq)]
29pub struct Affine {
30 /// Row-major coefficients.
31 pub m: [[f64; 3]; 2],
32}
33
34impl Affine {
35 /// Applies the map to a point.
36 pub fn apply(&self, x: f64, y: f64) -> (f64, f64) {
37 (
38 self.m[0][0] * x + self.m[0][1] * y + self.m[0][2],
39 self.m[1][0] * x + self.m[1][1] * y + self.m[1][2],
40 )
41 }
42
43 /// Inverts the map, which a warp needs because it walks the destination.
44 pub fn invert(&self) -> Outcome<Self> {
45 let det = self.m[0][0] * self.m[1][1] - self.m[0][1] * self.m[1][0];
46 if det.abs() < 1e-12 {
47 return Err(err!(
48 "An affine map with determinant {} cannot be inverted, which means the five \
49 landmarks are collinear.", det;
50 Invalid, Input, Range));
51 }
52 let (a, b, c, d) = (self.m[0][0], self.m[0][1], self.m[1][0], self.m[1][1]);
53 let (ia, ib, ic, id) = (d / det, -b / det, -c / det, a / det);
54 let tx = -(ia * self.m[0][2] + ib * self.m[1][2]);
55 let ty = -(ic * self.m[0][2] + id * self.m[1][2]);
56 Ok(Self { m: [[ia, ib, tx], [ic, id, ty]] })
57 }
58}
59
60/// Singular value decomposition of a two by two matrix, `a = u · diag(s) · vt`.
61///
62/// Two by two is small enough to solve in closed form, which avoids an
63/// iterative routine and keeps the result reproducible.
64fn svd2(a: [[f64; 2]; 2]) -> ([[f64; 2]; 2], [f64; 2], [[f64; 2]; 2]) {
65 let e = (a[0][0] + a[1][1]) / 2.0;
66 let f = (a[0][0] - a[1][1]) / 2.0;
67 let g = (a[1][0] + a[0][1]) / 2.0;
68 let h = (a[1][0] - a[0][1]) / 2.0;
69 let q = (e * e + h * h).sqrt();
70 let r = (f * f + g * g).sqrt();
71 let mut s0 = q + r;
72 let mut s1 = q - r;
73 let a1 = g.atan2(f);
74 let a2 = h.atan2(e);
75 let theta = (a2 - a1) / 2.0;
76 let phi = (a2 + a1) / 2.0;
77 let (cp, sp) = (phi.cos(), phi.sin());
78 let (ct, st) = (theta.cos(), theta.sin());
79 // With these two angles the decomposition reads `a = rot(phi) · s · rot(theta)`,
80 // so the right factor is already the one a product wants on the right.
81 let mut u = [[cp, -sp], [sp, cp]];
82 let mut vt = [[ct, -st], [st, ct]];
83 if s1 < 0.0 {
84 s1 = -s1;
85 vt[1][0] = -vt[1][0];
86 vt[1][1] = -vt[1][1];
87 }
88 if s1 > s0 {
89 core::mem::swap(&mut s0, &mut s1);
90 u = [[u[0][1], u[0][0]], [u[1][1], u[1][0]]];
91 vt = [[vt[1][0], vt[1][1]], [vt[0][0], vt[0][1]]];
92 }
93 (u, [s0, s1], vt)
94}
95
96/// Multiplies two two by two matrices.
97fn mul2(a: [[f64; 2]; 2], b: [[f64; 2]; 2]) -> [[f64; 2]; 2] {
98 [
99 [a[0][0] * b[0][0] + a[0][1] * b[1][0], a[0][0] * b[0][1] + a[0][1] * b[1][1]],
100 [a[1][0] * b[0][0] + a[1][1] * b[1][0], a[1][0] * b[0][1] + a[1][1] * b[1][1]],
101 ]
102}
103
104/// Determinant of a two by two matrix.
105fn det2(a: [[f64; 2]; 2]) -> f64 {
106 a[0][0] * a[1][1] - a[0][1] * a[1][0]
107}
108
109/// Builds the similarity transform taking five detected landmarks onto the
110/// template, by the Umeyama construction.
111///
112/// This is the least-squares rotation, uniform scale and translation, and it is
113/// what the reference implementation of the embedder's preprocessing computes.
114pub fn similarity(src: &[(f32, f32); 5]) -> Outcome<Affine> {
115 let mut src_mean = (0.0f64, 0.0f64);
116 for p in src {
117 src_mean.0 += p.0 as f64;
118 src_mean.1 += p.1 as f64;
119 }
120 src_mean.0 /= 5.0;
121 src_mean.1 /= 5.0;
122 let dst_mean = (TEMPLATE_MEAN.0 as f64, TEMPLATE_MEAN.1 as f64);
123
124 let mut sd = [[0.0f64; 2]; 5];
125 let mut dd = [[0.0f64; 2]; 5];
126 for i in 0..5 {
127 sd[i][0] = src[i].0 as f64 - src_mean.0;
128 sd[i][1] = src[i].1 as f64 - src_mean.1;
129 dd[i][0] = TEMPLATE[i].0 as f64 - dst_mean.0;
130 dd[i][1] = TEMPLATE[i].1 as f64 - dst_mean.1;
131 }
132
133 let mut a = [[0.0f64; 2]; 2];
134 for i in 0..5 {
135 a[0][0] += dd[i][0] * sd[i][0];
136 a[0][1] += dd[i][0] * sd[i][1];
137 a[1][0] += dd[i][1] * sd[i][0];
138 a[1][1] += dd[i][1] * sd[i][1];
139 }
140 for r in a.iter_mut() {
141 for v in r.iter_mut() {
142 *v /= 5.0;
143 }
144 }
145
146 let mut d = [1.0f64, 1.0];
147 if det2(a) < 0.0 {
148 d[1] = -1.0;
149 }
150 let (u, s, vt) = svd2(a);
151 let smax = s[0].max(s[1]);
152 let tol = smax * 2.0 * (f32::MIN_POSITIVE as f64);
153 let rank = (s[0] > tol) as usize + (s[1] > tol) as usize;
154
155 let rot = if rank == 1 {
156 if det2(u) * det2(vt) > 0.0 {
157 mul2(u, vt)
158 } else {
159 let dm = [[d[0], 0.0], [0.0, -1.0]];
160 mul2(u, mul2(dm, vt))
161 }
162 } else {
163 let dm = [[d[0], 0.0], [0.0, d[1]]];
164 mul2(u, mul2(dm, vt))
165 };
166
167 let mut var = 0.0f64;
168 for i in 0..5 {
169 var += sd[i][0] * sd[i][0];
170 }
171 for i in 0..5 {
172 var += sd[i][1] * sd[i][1];
173 }
174 var /= 5.0;
175 if var <= 0.0 {
176 return Err(err!(
177 "Five landmarks that all coincide give no scale to align by."; Invalid, Input, Range));
178 }
179 let scale = (s[0] * d[0] + s[1] * d[1]) / var;
180
181 let tsx = rot[0][0] * src_mean.0 + rot[0][1] * src_mean.1;
182 let tsy = rot[1][0] * src_mean.0 + rot[1][1] * src_mean.1;
183 Ok(Affine { m: [
184 [rot[0][0] * scale, rot[0][1] * scale, dst_mean.0 - scale * tsx],
185 [rot[1][0] * scale, rot[1][1] * scale, dst_mean.1 - scale * tsy],
186 ] })
187}
188
189/// Warps an image through an affine map into a square crop, sampling bilinearly
190/// and reading zero outside the source.
191pub fn warp(img: &Image<'_>, t: &Affine, side: usize) -> Outcome<Vec<u8>> {
192 let inv = res!(t.invert());
193 let ch = img.channels;
194 let mut out = vec![0u8; side * side * ch];
195 for y in 0..side {
196 for x in 0..side {
197 let (sx, sy) = inv.apply(x as f64 + 0.0, y as f64 + 0.0);
198 let x0 = sx.floor();
199 let y0 = sy.floor();
200 let fx = sx - x0;
201 let fy = sy - y0;
202 let dst = (y * side + x) * ch;
203 for c in 0..ch {
204 let p00 = img.sample(x0, y0, c);
205 let p10 = img.sample(x0 + 1.0, y0, c);
206 let p01 = img.sample(x0, y0 + 1.0, c);
207 let p11 = img.sample(x0 + 1.0, y0 + 1.0, c);
208 let top = p00 + (p10 - p00) * fx;
209 let bot = p01 + (p11 - p01) * fx;
210 let v = top + (bot - top) * fy;
211 out[dst + c] = v.round().clamp(0.0, 255.0) as u8;
212 }
213 }
214 }
215 Ok(out)
216}
217
218/// Warps a face out of an image onto the template, giving the crop the embedder
219/// consumes.
220pub fn align_crop(img: &Image<'_>, landmarks: &[(f32, f32); 5]) -> Outcome<Vec<u8>> {
221 let t = res!(similarity(landmarks));
222 warp(img, &t, CROP)
223}
224
225#[cfg(test)]
226mod tests {
227 use super::*;
228
229 #[test]
230 fn the_template_maps_to_itself() -> Outcome<()> {
231 let t = res!(similarity(&TEMPLATE));
232 for p in TEMPLATE.iter() {
233 let (x, y) = t.apply(p.0 as f64, p.1 as f64);
234 req!(((x - p.0 as f64).abs() < 1e-4), true);
235 req!(((y - p.1 as f64).abs() < 1e-4), true);
236 }
237 Ok(())
238 }
239
240 #[test]
241 fn a_scaled_and_turned_face_comes_back_to_the_template() -> Outcome<()> {
242 // Place the template in a larger frame, rotated by a fifth of a radian
243 // and scaled by three, and check the transform undoes exactly that.
244 let (c, s) = (0.2f64.cos(), 0.2f64.sin());
245 let mut src = [(0.0f32, 0.0f32); 5];
246 for i in 0..5 {
247 let (x, y) = (TEMPLATE[i].0 as f64, TEMPLATE[i].1 as f64);
248 src[i] = (
249 (3.0 * (c * x - s * y) + 100.0) as f32,
250 (3.0 * (s * x + c * y) + 40.0) as f32,
251 );
252 }
253 let t = res!(similarity(&src));
254 for i in 0..5 {
255 let (x, y) = t.apply(src[i].0 as f64, src[i].1 as f64);
256 req!(((x - TEMPLATE[i].0 as f64).abs() < 1e-3), true);
257 req!(((y - TEMPLATE[i].1 as f64).abs() < 1e-3), true);
258 }
259 Ok(())
260 }
261
262 #[test]
263 fn a_decomposition_reproduces_its_matrix() {
264 for a in [
265 [[3.0f64, 1.0], [0.5, 2.0]],
266 [[-1.0f64, 2.0], [3.0, 0.25]],
267 [[0.0f64, 1.0], [1.0, 0.0]],
268 ] {
269 let (u, s, vt) = svd2(a);
270 let m = mul2(u, mul2([[s[0], 0.0], [0.0, s[1]]], vt));
271 for r in 0..2 {
272 for c in 0..2 {
273 assert!((m[r][c] - a[r][c]).abs() < 1e-9, "{:?} against {:?}", m, a);
274 }
275 }
276 assert!(s[0] >= s[1] && s[1] >= 0.0, "singular values {:?}", s);
277 }
278 }
279}