1use crate::kernels::geometry::{js_ceil, js_floor, js_round, Rect};
10use crate::kernels::pool::SyncPtr;
11
12pub struct Vote<'a> {
14 pub x0: i32,
15 pub y0: i32,
16 pub w: i32,
17 pub h: i32,
18 pub bits: &'a [u8],
20 pub clean: &'a [u8],
22}
23
24impl Vote<'_> {
25 #[inline]
27 pub fn any_negative(&self, ax0: i32, ax1: i32, ay: i32) -> bool {
28 (ax0..=ax1).any(|ax| self.verdict(ax, ay) < 0)
29 }
30
31 #[inline]
33 pub fn verdict(&self, ax: i32, ay: i32) -> i32 {
34 let lx = ax - self.x0;
35 let ly = ay - self.y0;
36 if lx < 0 || ly < 0 || lx >= self.w || ly >= self.h {
37 return 0;
38 }
39 let i = (ly * self.w + lx) as usize;
40 let bit = 1u8 << (i & 7);
41 if self.bits.get(i >> 3).is_some_and(|b| b & bit != 0) {
42 -1
43 } else if self.clean.get(i >> 3).is_some_and(|b| b & bit != 0) {
44 1
45 } else {
46 0
47 }
48 }
49}
50
51pub struct Neighbour<'a> {
52 pub rgba: &'a [u8],
53 pub pose_x: f64,
54 pub pose_y: f64,
55 pub occlusions: &'a [Rect],
56 pub vote: Option<Vote<'a>>,
57}
58
59pub struct MaskInput<'a> {
60 pub rgba: &'a [u8],
61 pub labels: &'a [u8],
62 pub width: usize,
63 pub height: usize,
64 pub region: Rect,
65 pub code: u8,
66 pub pose_x: f64,
67 pub pose_y: f64,
68 pub prev: Option<Neighbour<'a>>,
69 pub next: Option<Neighbour<'a>>,
70 pub vote: Option<Vote<'a>>,
71 pub factor: i32,
72 pub noise: f64,
73}
74
75struct Resolved<'n, 'a> {
76 rgba: &'a [u8],
77 dx: i64,
78 dy: i64,
79 occlusions: &'a [Rect],
80 vote: Option<&'n Vote<'a>>,
81}
82
83fn resolve<'n, 'a>(n: &'n Neighbour<'a>, cx: i32, cy: i32) -> Resolved<'n, 'a> {
84 Resolved {
85 rgba: n.rgba,
86 dx: (cx - js_round(n.pose_x)) as i64,
87 dy: (cy - js_round(n.pose_y)) as i64,
88 occlusions: n.occlusions,
89 vote: n.vote.as_ref(),
90 }
91}
92
93#[inline]
95fn span(x: f64, width: f64) -> (i64, i64) {
96 (x.ceil() as i64, (x + width).ceil() as i64)
97}
98
99pub fn sum_limit(noise: f64) -> u32 {
102 (0..=765u32)
103 .take_while(|&s| s as f64 / 3.0 <= noise)
104 .last()
105 .unwrap_or(0)
106}
107
108struct Ctx<'n, 'a> {
110 input: &'n MaskInput<'a>,
111 neighbours: [Option<Resolved<'n, 'a>>; 2],
112 w: i64,
113 h: i64,
114 rx0: i64,
115 rx1: i64,
116 factor: i64,
117 limit: u32,
118}
119
120struct NeighbourRow<'r, 'n, 'a> {
122 n: &'r Resolved<'n, 'a>,
123 labels: &'a [u8],
124 rgba: &'a [u8],
125 lo: i64,
127 hi: i64,
128 ay: i32,
130 fy: f64,
132 occluded: bool,
134}
135
136impl NeighbourRow<'_, '_, '_> {
137 #[inline]
138 fn occludes(&self, sx: i64) -> bool {
139 self.occluded
140 && self.n.occlusions.iter().any(|o| {
141 if !(self.fy >= o.y && self.fy < o.y + o.height) {
142 return false;
143 }
144 let (a, b) = span(o.x, o.width);
145 sx >= a - self.n.dx && sx < b - self.n.dx
146 })
147 }
148}
149
150pub fn consistency_mask(input: &MaskInput<'_>, out: &mut [u8]) {
154 let (w, h) = (input.width as i64, input.height as i64);
155 let (cx, cy) = (js_round(input.pose_x), js_round(input.pose_y));
156 let ctx = Ctx {
157 input,
158 neighbours: [
159 input.prev.as_ref().map(|n| resolve(n, cx, cy)),
160 input.next.as_ref().map(|n| resolve(n, cx, cy)),
161 ],
162 w,
163 h,
164 rx0: js_floor(input.region.x).clamp(0, w as i32) as i64,
165 rx1: js_ceil(input.region.x + input.region.width).clamp(0, w as i32) as i64,
166 factor: input.factor.max(1) as i64,
167 limit: sum_limit(input.noise),
168 };
169 let ry0 = js_floor(input.region.y).clamp(0, h as i32) as usize;
170 let ry1 = js_ceil(input.region.y + input.region.height).clamp(0, h as i32) as usize;
171 let width = input.width;
172 let out = &mut out[..input.width * input.height];
173 out[..ry0 * width].fill(1);
175 out[ry1.max(ry0) * width..].fill(1);
176 let rows = ry1.saturating_sub(ry0);
177 let chunks = crate::kernels::pool::chunks_for(rows * width, 64 * 1024);
178 let base = SyncPtr(out.as_mut_ptr());
179 crate::kernels::pool::par_for(chunks, |c| {
180 let (a, b) = (
181 ry0 + crate::kernels::pool::split(rows, chunks, c),
182 ry0 + crate::kernels::pool::split(rows, chunks, c + 1),
183 );
184 for sy in a..b {
185 let row = unsafe { std::slice::from_raw_parts_mut(base.get().add(sy * width), width) };
187 row.fill(1);
188 mask_row(&ctx, sy as i64, row);
189 }
190 });
191}
192
193fn mask_row(ctx: &Ctx<'_, '_>, sy: i64, out_row: &mut [u8]) {
194 let input = ctx.input;
195 let (w, h, factor) = (ctx.w, ctx.h, ctx.factor);
196 let width = input.width;
197 let labels = input.labels;
198 let ay = sy.div_euclid(factor) as i32;
199 let mut rows: [Option<NeighbourRow<'_, '_, '_>>; 2] = [None, None];
200 let mut count = 0usize;
201 for n in ctx.neighbours.iter().flatten() {
202 let iy = sy + n.dy;
203 if iy < 0 || iy >= h {
204 continue;
205 }
206 let base = (iy * w) as usize;
207 let fy = iy as f64;
208 rows[count] = Some(NeighbourRow {
209 n,
210 labels: &labels[base..base + width],
211 rgba: &n.rgba[base * 4..(base + width) * 4],
212 lo: (-n.dx).max(ctx.rx0),
213 hi: (w - n.dx).min(ctx.rx1),
214 ay: iy.div_euclid(factor) as i32,
215 fy,
216 occluded: n
217 .occlusions
218 .iter()
219 .any(|o| fy >= o.y && fy < o.y + o.height),
220 });
221 count += 1;
222 }
223 let rows = &rows;
224 let row_vote = input
225 .vote
226 .as_ref()
227 .filter(|v| ay >= v.y0 && ay < v.y0 + v.h);
228 let base = (sy * w) as usize;
229 let lab_row = &labels[base..base + width];
230 let rgb_row = &input.rgba[base * 4..(base + width) * 4];
231 let mut sx = ctx.rx0;
232 let (fast_lo, fast_hi) = rows
236 .iter()
237 .flatten()
238 .fold((ctx.rx0, ctx.rx1), |(lo, hi), r| {
239 if r.occluded {
240 (lo, lo)
241 } else {
242 (lo.max(r.lo), hi.min(r.hi))
243 }
244 });
245 while sx < ctx.rx1 {
246 if sx >= fast_lo && sx + 4 <= fast_hi {
247 let votes_clear = row_vote.is_none_or(|v| {
248 !v.any_negative(
249 sx.div_euclid(factor) as i32,
250 (sx + 3).div_euclid(factor) as i32,
251 ay,
252 )
253 });
254 if votes_clear && block_agrees(ctx, rows, lab_row, rgb_row, sx as usize) {
255 sx += 4;
256 continue;
257 }
258 for x in sx..sx + 4 {
259 mask_pixel(ctx, rows, row_vote, ay, lab_row, rgb_row, out_row, x);
260 }
261 sx += 4;
262 continue;
263 }
264 mask_pixel(ctx, rows, row_vote, ay, lab_row, rgb_row, out_row, sx);
265 sx += 1;
266 }
267}
268
269#[inline]
273fn block_agrees(
274 ctx: &Ctx<'_, '_>,
275 rows: &[Option<NeighbourRow<'_, '_, '_>>; 2],
276 lab_row: &[u8],
277 rgb_row: &[u8],
278 x: usize,
279) -> bool {
280 let code = ctx.input.code;
281 #[cfg(target_feature = "simd128")]
282 {
283 use core::arch::wasm32::*;
284 let _ = code;
285 unsafe {
288 let codes = u32x4_splat(ctx.input.code as u32);
289 let labels4 = |row: &[u8], at: usize| {
290 let v = v128_load32_zero(row.as_ptr().add(at) as *const u32);
291 u32x4_extend_low_u16x8(u16x8_extend_low_u8x16(v))
292 };
293 let inside = u32x4_eq(labels4(lab_row, x), codes);
294 if !v128_any_true(inside) {
295 return true;
296 }
297 let a = v128_load(rgb_row.as_ptr().add(x * 4) as *const v128);
298 let limit = u32x4_splat(ctx.limit);
299 let rgb = u32x4_splat(0x00ff_ffff);
300 let mut bad = u32x4_splat(0);
301 for r in rows.iter().flatten() {
302 let ix = (x as i64 + r.n.dx) as usize;
303 let b = v128_load(r.rgba.as_ptr().add(ix * 4) as *const v128);
304 let d = v128_and(v128_or(u8x16_sub_sat(a, b), u8x16_sub_sat(b, a)), rgb);
305 let sum = u32x4_extadd_pairwise_u16x8(u16x8_extadd_pairwise_u8x16(d));
306 let differs = u32x4_gt(sum, limit);
307 let counted = u32x4_eq(labels4(r.labels, ix), codes);
308 bad = v128_or(bad, v128_and(differs, counted));
309 }
310 !v128_any_true(v128_and(bad, inside))
311 }
312 }
313 #[cfg(not(target_feature = "simd128"))]
314 {
315 (x..x + 4).all(|x| {
316 lab_row[x] != code
317 || rows.iter().flatten().all(|r| {
318 let ix = (x as i64 + r.n.dx) as usize;
319 if r.labels[ix] != code {
320 return true;
321 }
322 let (i, j) = (x * 4, ix * 4);
323 let sum = (rgb_row[i] as i32 - r.rgba[j] as i32).unsigned_abs()
324 + (rgb_row[i + 1] as i32 - r.rgba[j + 1] as i32).unsigned_abs()
325 + (rgb_row[i + 2] as i32 - r.rgba[j + 2] as i32).unsigned_abs();
326 sum <= ctx.limit
327 })
328 })
329 }
330}
331
332#[allow(clippy::too_many_arguments)]
334#[inline]
335fn mask_pixel(
336 ctx: &Ctx<'_, '_>,
337 rows: &[Option<NeighbourRow<'_, '_, '_>>; 2],
338 row_vote: Option<&Vote<'_>>,
339 ay: i32,
340 lab_row: &[u8],
341 rgb_row: &[u8],
342 out_row: &mut [u8],
343 sx: i64,
344) {
345 let code = ctx.input.code;
346 let factor = ctx.factor;
347 let x = sx as usize;
348 if lab_row[x] != code {
349 return;
350 }
351 if row_vote.is_some_and(|v| v.verdict(sx.div_euclid(factor) as i32, ay) < 0) {
352 out_row[x] = 0;
353 return;
354 }
355 let k = x * 4;
356 let (r, g, b) = (
357 rgb_row[k] as i32,
358 rgb_row[k + 1] as i32,
359 rgb_row[k + 2] as i32,
360 );
361 let (mut checked, mut condemned, mut excused) = (0u32, false, 0u32);
362 for row in rows.iter().flatten() {
363 if sx < row.lo || sx >= row.hi {
366 continue;
367 }
368 let ix = (sx + row.n.dx) as usize;
369 if row.labels[ix] != code || row.occludes(sx) {
370 continue;
371 }
372 checked += 1;
373 let j = ix * 4;
374 let (nr, ng, nb) = (
375 row.rgba[j] as i32,
376 row.rgba[j + 1] as i32,
377 row.rgba[j + 2] as i32,
378 );
379 if r == nr && g == ng && b == nb {
380 continue;
381 }
382 let diff = ((r - nr).abs() + (g - ng).abs() + (b - nb).abs()) as f64 / 3.0;
383 if diff <= ctx.input.noise {
384 continue;
385 }
386 if row
387 .n
388 .vote
389 .is_some_and(|v| v.verdict((ix as i64).div_euclid(factor) as i32, row.ay) < 0)
390 {
391 excused += 1;
392 } else {
393 condemned = true;
394 break;
395 }
396 }
397 if condemned || (excused > 0 && checked >= 2) {
398 out_row[x] = 0;
399 }
400}
401
402pub struct FixedMaskInput<'a> {
405 pub rgba: &'a [u8],
406 pub labels: &'a [u8],
407 pub width: usize,
408 pub height: usize,
409 pub region: Rect,
410 pub code: u8,
411 pub prev: Option<&'a [u8]>,
412 pub next: Option<&'a [u8]>,
413 pub noise: f64,
414}
415
416pub fn fixed_mask(input: &FixedMaskInput<'_>, out: &mut [u8]) {
425 let (w, h) = (input.width, input.height);
426 let out = &mut out[..w * h];
427 out.fill(1);
428 let neighbours: Vec<&[u8]> = [input.prev, input.next].into_iter().flatten().collect();
429 if neighbours.is_empty() {
430 return;
431 }
432 let x0 = js_floor(input.region.x).clamp(0, w as i32) as usize;
433 let x1 = js_ceil(input.region.x + input.region.width).clamp(0, w as i32) as usize;
434 let y0 = js_floor(input.region.y).clamp(0, h as i32) as usize;
435 let y1 = js_ceil(input.region.y + input.region.height).clamp(0, h as i32) as usize;
436 let rows = y1.saturating_sub(y0);
437 let limit = sum_limit(input.noise);
438 let chunks = crate::kernels::pool::chunks_for(rows * w, 64 * 1024);
439 let base = SyncPtr(out.as_mut_ptr());
440 crate::kernels::pool::par_for(chunks, |c| {
441 let (a, b) = (
442 y0 + crate::kernels::pool::split(rows, chunks, c),
443 y0 + crate::kernels::pool::split(rows, chunks, c + 1),
444 );
445 for y in a..b {
446 let row = unsafe { std::slice::from_raw_parts_mut(base.get().add(y * w), w) };
448 for (x, flag) in row.iter_mut().enumerate().take(x1).skip(x0) {
449 let i = y * w + x;
450 if input.labels[i] != input.code {
451 continue;
452 }
453 let k = i * 4;
454 let px = &input.rgba[k..k + 3];
455 let differs = |n: &&[u8]| {
456 let q = &n[k..k + 3];
457 let sum = px[0].abs_diff(q[0]) as u32
458 + px[1].abs_diff(q[1]) as u32
459 + px[2].abs_diff(q[2]) as u32;
460 sum > limit
461 };
462 if neighbours.iter().all(differs) {
463 *flag = 0;
464 }
465 }
466 }
467 });
468}
469
470#[cfg(test)]
471mod tests {
472 use super::*;
473
474 fn frame(w: usize, h: usize, value: u8) -> Vec<u8> {
475 (0..w * h)
476 .flat_map(|_| [value, value, value, 255])
477 .collect()
478 }
479
480 #[test]
481 fn lone_disagreeing_neighbour_condemns_unless_excused_by_its_verdict() {
482 let (w, h) = (4usize, 4usize);
483 let current = frame(w, h, 10);
484 let other = frame(w, h, 200);
485 let labels = vec![1u8; w * h];
486 let region = Rect {
487 x: 0.0,
488 y: 0.0,
489 width: 4.0,
490 height: 4.0,
491 };
492 let mut out = vec![0u8; w * h];
493 let base = |vote| MaskInput {
494 rgba: ¤t,
495 labels: &labels,
496 width: w,
497 height: h,
498 region,
499 code: 1,
500 pose_x: 0.0,
501 pose_y: 0.0,
502 prev: Some(Neighbour {
503 rgba: &other,
504 pose_x: 0.0,
505 pose_y: 0.0,
506 occlusions: &[],
507 vote,
508 }),
509 next: None,
510 vote: None,
511 factor: 1,
512 noise: 10.0,
513 };
514 consistency_mask(&base(None), &mut out);
515 assert!(out.iter().all(|&v| v == 0));
516 let bits = vec![0xffu8; 2];
517 let clean = vec![0u8; 2];
518 let excused = base(Some(Vote {
519 x0: 0,
520 y0: 0,
521 w: 4,
522 h: 4,
523 bits: &bits,
524 clean: &clean,
525 }));
526 consistency_mask(&excused, &mut out);
527 assert!(out.iter().all(|&v| v == 1));
528 }
529
530 fn strip(values: [u8; 4]) -> Vec<u8> {
532 values.iter().flat_map(|&v| [v, v, v, 255]).collect()
533 }
534
535 fn fixed(current: &[u8], prev: Option<&[u8]>, next: Option<&[u8]>, labels: &[u8]) -> Vec<u8> {
536 let mut out = vec![9u8; 4];
537 fixed_mask(
538 &FixedMaskInput {
539 rgba: current,
540 labels,
541 width: 4,
542 height: 1,
543 region: Rect {
544 x: 0.0,
545 y: 0.0,
546 width: 4.0,
547 height: 1.0,
548 },
549 code: 1,
550 prev,
551 next,
552 noise: 10.0,
553 },
554 &mut out,
555 );
556 out
557 }
558
559 #[test]
560 fn a_fixed_pixel_is_condemned_only_when_it_differs_from_every_neighbour() {
561 let labels = [1u8; 4];
562 let current = strip([250, 200, 100, 100]);
565 let prev = strip([100, 120, 100, 95]);
566 let next = strip([100, 200, 180, 108]);
567 assert_eq!(
568 fixed(¤t, Some(&prev), Some(&next), &labels),
569 vec![0, 1, 1, 1]
570 );
571 }
572
573 #[test]
574 fn a_lone_fixed_neighbour_that_disagrees_condemns() {
575 let labels = [1u8; 4];
577 let current = strip([250, 100, 100, 100]);
578 let next = strip([100, 100, 100, 100]);
579 assert_eq!(
580 fixed(¤t, None, Some(&next), &labels),
581 vec![0, 1, 1, 1]
582 );
583 assert_eq!(
584 fixed(¤t, Some(&next), None, &labels),
585 vec![0, 1, 1, 1]
586 );
587 }
588
589 #[test]
590 fn a_fixed_pixel_without_neighbours_or_outside_the_region_is_consistent() {
591 let current = strip([250, 250, 250, 250]);
592 let other = strip([0, 0, 0, 0]);
593 assert_eq!(fixed(¤t, None, None, &[1; 4]), vec![1, 1, 1, 1]);
594 assert_eq!(
595 fixed(¤t, Some(&other), Some(&other), &[1, 2, 1, 2]),
596 vec![0, 1, 0, 1]
597 );
598 }
599}