use crate::kernels::geometry::{js_ceil, js_round, Rect};
use crate::kernels::pool::SyncPtr;
pub const QUALITY_BLOCK: usize = 16;
pub struct TileBuffers<'a> {
pub size: usize,
pub pixels: &'a mut [u8],
pub coverage: &'a mut [u8],
pub provisional: &'a mut [u8],
pub quality: &'a mut [u8],
pub conflicts: &'a mut [u8],
pub owner: &'a mut [u32],
pub score: &'a mut [f32],
pub frozen: &'a [u8],
}
pub struct Observation<'a> {
pub rgba: &'a [u8],
pub width: usize,
pub height: usize,
pub labels: Option<(&'a [u8], u8)>,
pub occlusions: &'a [Rect],
pub consistent: Option<&'a [u8]>,
pub confidence: f64,
pub uncertain: bool,
pub frame: u32,
pub latest: bool,
}
#[derive(Default, Debug, PartialEq)]
pub struct TileStats {
pub added: u32,
pub conflicts: u32,
pub uncertain: u32,
pub provisional_delta: i32,
pub changed: bool,
pub settled_blocks: u32,
pub conflict_blocks: Vec<(u32, u32)>,
}
#[inline]
fn bit(bits: &[u8], p: usize) -> bool {
bits[p >> 3] & (1 << (p & 7)) != 0
}
#[inline]
fn set_bit(bits: &mut [u8], p: usize) {
bits[p >> 3] |= 1 << (p & 7);
}
#[inline]
fn clear_bit(bits: &mut [u8], p: usize) {
bits[p >> 3] &= !(1 << (p & 7));
}
#[inline]
fn occluded(occlusions: &[Rect], sx: i64, sy: i64) -> bool {
occlusions.iter().any(|o| o.contains(sx as f64, sy as f64))
}
#[inline]
fn px32(rgba: &[u8], i: usize) -> u32 {
u32::from_le_bytes([
rgba[i * 4],
rgba[i * 4 + 1],
rgba[i * 4 + 2],
rgba[i * 4 + 3],
])
}
struct Geometry<'o, 'a> {
obs: &'o Observation<'a>,
world: Rect,
size: i64,
blocks: i64,
dx: i64,
dy: i64,
w: i64,
h: i64,
tile_x: f64,
tile_y: f64,
}
#[derive(Default, Clone, Copy)]
struct Part {
added: u32,
conflicts: u32,
uncertain: u32,
provisional_delta: i32,
changed: bool,
settled: u32,
}
struct Rows<'t> {
pixels: &'t mut [u8],
coverage: &'t mut [u8],
provisional: &'t mut [u8],
quality: &'t mut [u8],
conflicts: &'t mut [u8],
owner: &'t mut [u32],
score: &'t mut [f32],
frozen: &'t [u8],
flagged: &'t mut [u8],
by0: i64,
}
#[allow(clippy::too_many_arguments)]
pub fn composite_tile(
tile: &mut TileBuffers<'_>,
obs: &Observation<'_>,
world: Rect,
ox: i64,
oy: i64,
tx: i64,
ty: i64,
) -> TileStats {
let size = tile.size as i64;
let b = QUALITY_BLOCK as i64;
let blocks = size / b;
let mut stats = TileStats::default();
let tile_rect = Rect {
x: (tx * size) as f64,
y: (ty * size) as f64,
width: size as f64,
height: size as f64,
};
let px = world.x.max(tile_rect.x);
let py = world.y.max(tile_rect.y);
let pw = ((world.x + world.width).min(tile_rect.x + tile_rect.width) - px).max(0.0);
let ph = ((world.y + world.height).min(tile_rect.y + tile_rect.height) - py).max(0.0);
let bx0 = ((px - tile_rect.x) / b as f64).floor() as i64;
let by0 = ((py - tile_rect.y) / b as f64).floor() as i64;
let bx1 = js_ceil((px + pw - tile_rect.x) / b as f64) as i64;
let by1 = js_ceil((py + ph - tile_rect.y) / b as f64) as i64;
if by1 <= by0 || bx1 <= bx0 {
return stats;
}
let geo = Geometry {
obs,
world,
size,
blocks,
dx: tx * size - ox,
dy: ty * size - oy,
w: obs.width as i64,
h: obs.height as i64,
tile_x: tile_rect.x,
tile_y: tile_rect.y,
};
let rows = (by1 - by0) as usize;
let chunks = crate::kernels::pool::chunks_for(rows * ((bx1 - bx0) * b * b) as usize, 16 * 1024);
let mut parts = vec![Part::default(); chunks];
let mut flagged = vec![0u8; rows * blocks as usize];
let (row_px, row_bits, row_blocks) = (
(b * size) as usize,
(b * size / 8) as usize,
blocks as usize,
);
let base = |r: usize| by0 as usize + r;
let ptrs = (
SyncPtr(tile.pixels.as_mut_ptr()),
SyncPtr(tile.coverage.as_mut_ptr()),
SyncPtr(tile.provisional.as_mut_ptr()),
SyncPtr(tile.quality.as_mut_ptr()),
SyncPtr(tile.conflicts.as_mut_ptr()),
SyncPtr(tile.owner.as_mut_ptr()),
SyncPtr(tile.score.as_mut_ptr()),
SyncPtr(flagged.as_mut_ptr()),
SyncPtr(parts.as_mut_ptr()),
);
let frozen = tile.frozen;
crate::kernels::pool::par_for(chunks, |c| {
let (ra, rb) = (
crate::kernels::pool::split(rows, chunks, c),
crate::kernels::pool::split(rows, chunks, c + 1),
);
if ra == rb {
return;
}
let (a, n) = (base(ra), rb - ra);
let mut part = unsafe {
Rows {
pixels: std::slice::from_raw_parts_mut(
ptrs.0.get().add(a * row_px * 4),
n * row_px * 4,
),
coverage: std::slice::from_raw_parts_mut(
ptrs.1.get().add(a * row_bits),
n * row_bits,
),
provisional: std::slice::from_raw_parts_mut(
ptrs.2.get().add(a * row_bits),
n * row_bits,
),
quality: std::slice::from_raw_parts_mut(
ptrs.3.get().add(a * row_blocks),
n * row_blocks,
),
conflicts: std::slice::from_raw_parts_mut(
ptrs.4.get().add(a * row_blocks),
n * row_blocks,
),
owner: std::slice::from_raw_parts_mut(
ptrs.5.get().add(a * row_blocks),
n * row_blocks,
),
score: std::slice::from_raw_parts_mut(
ptrs.6.get().add(a * row_blocks),
n * row_blocks,
),
frozen: &frozen[a * row_blocks..(a + n) * row_blocks],
flagged: std::slice::from_raw_parts_mut(
ptrs.7.get().add(ra * row_blocks),
n * row_blocks,
),
by0: a as i64,
}
};
let mut sum = Part::default();
for by in a as i64..(a + n) as i64 {
for bx in bx0..bx1 {
composite_block(&geo, &mut part, bx, by, &mut sum);
}
}
unsafe { *ptrs.8.get().add(c) = sum };
});
for p in &parts {
stats.added += p.added;
stats.conflicts += p.conflicts;
stats.uncertain += p.uncertain;
stats.provisional_delta += p.provisional_delta;
stats.changed |= p.changed;
stats.settled_blocks += p.settled;
}
for r in 0..rows {
for bx in bx0..bx1 {
if flagged[r * blocks as usize + bx as usize] != 0 {
stats
.conflict_blocks
.push((bx as u32, (by0 + r as i64) as u32));
}
}
}
stats
}
#[inline]
fn columns(geo: &Geometry<'_, '_>, sx0: i64, sx1: i64) -> (i64, i64) {
((-geo.dx).max(sx0), (geo.w - geo.dx).min(sx1))
}
#[inline]
fn row_occluded(occlusions: &[Rect], sy: i64) -> bool {
let fy = sy as f64;
occlusions.iter().any(|o| fy >= o.y && fy < o.y + o.height)
}
#[cfg(target_feature = "simd128")]
#[inline]
fn row_stats_simd(
obs: &Observation<'_>,
pixels: &[u8],
coverage: &[u8],
dst0: usize,
src0: usize,
) -> (u32, u32, u32, u32) {
use core::arch::wasm32::*;
let inside: u32 = match obs.labels {
None => 0xffff,
Some((labels, code)) => unsafe {
let l = v128_load(labels.as_ptr().add(src0) as *const v128);
u8x16_bitmask(u8x16_eq(l, u8x16_splat(code))) as u32
},
};
let cov = coverage[dst0 >> 3] as u32 | (coverage[(dst0 >> 3) + 1] as u32) << 8;
let (mut eq, mut mism) = (0u32, 0u32);
let rgb = u32x4_splat(0x00ff_ffff);
let limit = u32x4_splat(75);
for g in 0..4 {
let (t, o) = unsafe {
(
v128_load(pixels.as_ptr().add((dst0 + 4 * g) * 4) as *const v128),
v128_load(obs.rgba.as_ptr().add((src0 + 4 * g) * 4) as *const v128),
)
};
eq |= (i32x4_bitmask(i32x4_eq(t, o)) as u32) << (4 * g);
let d = v128_and(v128_or(u8x16_sub_sat(t, o), u8x16_sub_sat(o, t)), rgb);
let sum = u32x4_extadd_pairwise_u16x8(u16x8_extadd_pairwise_u8x16(d));
mism |= (i32x4_bitmask(u32x4_gt(sum, limit)) as u32) << (4 * g);
}
let ov = cov & inside;
(
inside.count_ones(),
ov.count_ones(),
(ov & eq).count_ones(),
(ov & mism).count_ones(),
)
}
#[cfg(not(target_feature = "simd128"))]
#[inline]
fn row_stats_scalar(
obs: &Observation<'_>,
pixels: &[u8],
coverage: &[u8],
dst0: usize,
src0: usize,
) -> (u32, u32, u32, u32) {
const N: usize = QUALITY_BLOCK;
let (t, o) = (
&pixels[dst0 * 4..(dst0 + N) * 4],
&obs.rgba[src0 * 4..(src0 + N) * 4],
);
let inside: u32 = match obs.labels {
None => (1 << N) - 1,
Some((labels, code)) => labels[src0..src0 + N]
.iter()
.enumerate()
.fold(0, |m, (i, &l)| m | ((l == code) as u32) << i),
};
let cov = coverage[dst0 >> 3] as u32 | (coverage[(dst0 >> 3) + 1] as u32) << 8;
let (mut eq, mut mism) = (0u32, 0u32);
for (i, (a, b)) in t.chunks_exact(4).zip(o.chunks_exact(4)).enumerate() {
eq |= ((a == b) as u32) << i;
let sum =
a[0].abs_diff(b[0]) as u32 + a[1].abs_diff(b[1]) as u32 + a[2].abs_diff(b[2]) as u32;
mism |= ((sum > 75) as u32) << i;
}
let ov = cov & inside;
(
inside.count_ones(),
ov.count_ones(),
(ov & eq).count_ones(),
(ov & mism).count_ones(),
)
}
fn composite_block(geo: &Geometry<'_, '_>, t: &mut Rows<'_>, bx: i64, by: i64, sum: &mut Part) {
let obs = geo.obs;
let world = geo.world;
let (size, b) = (geo.size, QUALITY_BLOCK as i64);
let (w, h) = (geo.w, geo.h);
let (tile_x, tile_y) = (geo.tile_x, geo.tile_y);
let ly = by - t.by0;
let q = (ly * geo.blocks + bx) as usize;
let inside = |src: usize| obs.labels.is_none_or(|(labels, code)| labels[src] == code);
let (mut mismatch, mut overlap, mut count, mut identical) = (0u32, 0u32, 0u32, 0u32);
let sy0 = js_ceil(((by * b) as f64).max(world.y - tile_y)) as i64;
let sy1 = js_ceil((((by + 1) * b) as f64).min(world.y + world.height - tile_y)) as i64;
let sx0 = js_ceil(((bx * b) as f64).max(world.x - tile_x)) as i64;
let sx1 = js_ceil((((bx + 1) * b) as f64).min(world.x + world.width - tile_x)) as i64;
let (cx0, cx1) = columns(geo, sx0, sx1);
let full = cx0 == bx * b && cx1 == (bx + 1) * b;
let dst_of = |x: i64, y: i64| ((y - t.by0 * b) * size + x) as usize;
for y in sy0..sy1 {
let sy = y + geo.dy;
if sy < 0 || sy >= h || cx1 <= cx0 {
continue;
}
let occluded_row = row_occluded(obs.occlusions, sy);
if full && !occluded_row {
#[cfg(target_feature = "simd128")]
let row_stats = row_stats_simd;
#[cfg(not(target_feature = "simd128"))]
let row_stats = row_stats_scalar;
let (c, o, i, m) = row_stats(
obs,
t.pixels,
t.coverage,
dst_of(cx0, y),
(sy * w + cx0 + geo.dx) as usize,
);
count += c;
overlap += o;
identical += i;
mismatch += m;
continue;
}
for x in cx0..cx1 {
let sx = x + geo.dx;
if occluded_row && occluded(obs.occlusions, sx, sy) {
continue;
}
let src = (sy * w + sx) as usize;
if !inside(src) {
continue;
}
let dst = dst_of(x, y);
count += 1;
if bit(t.coverage, dst) {
overlap += 1;
if px32(t.pixels, dst) == px32(obs.rgba, src) {
identical += 1;
continue;
}
let (i, j) = (dst * 4, src * 4);
let diff = ((t.pixels[i] as i32 - obs.rgba[j] as i32).abs()
+ (t.pixels[i + 1] as i32 - obs.rgba[j + 1] as i32).abs()
+ (t.pixels[i + 2] as i32 - obs.rgba[j + 2] as i32).abs())
as f64
/ 3.0;
if diff > 25.0 {
mismatch += 1;
}
}
}
}
if count == 0 {
return;
}
let (mut covered_in_block, mut provisional_in_block) = (0u32, 0i32);
for row in 0..b {
let i = dst_of(bx * b, by * b + row) >> 3;
covered_in_block += t.coverage[i].count_ones() + t.coverage[i + 1].count_ones();
provisional_in_block +=
(t.provisional[i].count_ones() + t.provisional[i + 1].count_ones()) as i32;
}
if identical == count && provisional_in_block == 0 {
return;
}
let conflict = overlap >= 12 && mismatch as f64 / overlap as f64 > 0.16;
let complete = overlap == covered_in_block;
let edge = ((geo.tile_x as i64 + bx * b) as f64 - world.x)
.min((geo.tile_y as i64 + by * b) as f64 - world.y)
.min(world.x + world.width - (geo.tile_x as i64 + (bx + 1) * b) as f64)
.min(world.y + world.height - (geo.tile_y as i64 + (by + 1) * b) as f64);
let mut sharpness = 0u64;
if (!conflict && t.frozen[q] == 0 && complete) || t.owner[q] == 0 {
let (gx0, gx1) = ((1 - geo.dx).max(cx0), (w - 1 - geo.dx).min(cx1));
for y in sy0..sy1 {
let sy = y + geo.dy;
if sy < 0 || sy >= h {
continue;
}
let occluded_row = row_occluded(obs.occlusions, sy);
for x in gx0..gx1 {
let sx = x + geo.dx;
if occluded_row && occluded(obs.occlusions, sx, sy) {
continue;
}
let src = (sy * w + sx) as usize;
if inside(src) {
let j = src * 4;
sharpness +=
(obs.rgba[j - 4] as i32 - obs.rgba[j + 4] as i32).unsigned_abs() as u64;
}
}
}
}
let had_provisional = provisional_in_block > 0;
let score = obs.confidence * 100.0
+ (12.0f64).min(sharpness as f64 / count as f64 * 0.15)
+ (6.0f64).min(edge.max(0.0) / 40.0);
let replace = complete
&& t.frozen[q] == 0
&& (obs.latest || (!conflict && score > t.score[q] as f64 + 4.0));
if conflict && obs.latest {
if let Some(consistent) = obs.consistent {
let settled = (sy0..sy1).all(|y| {
let sy = y + geo.dy;
sy < 0
|| sy >= h
|| (cx0..cx1).all(|x| {
let src = (sy * w + x + geo.dx) as usize;
!inside(src)
|| occluded(obs.occlusions, x + geo.dx, sy)
|| consistent[src] != 0
})
});
if settled {
sum.settled += 1;
}
}
}
if conflict {
t.conflicts[q] = 1;
t.flagged[q] = 1;
sum.conflicts += mismatch;
sum.changed = true;
}
if replace || overlap < count || provisional_in_block > 0 {
for y in sy0..sy1 {
let sy = y + geo.dy;
if sy < 0 || sy >= h {
continue;
}
let occluded_row = row_occluded(obs.occlusions, sy);
for x in cx0..cx1 {
let sx = x + geo.dx;
if occluded_row && occluded(obs.occlusions, sx, sy) {
continue;
}
let src = (sy * w + sx) as usize;
if !inside(src) {
continue;
}
let dst = dst_of(x, y);
let fresh = !bit(t.coverage, dst);
let bad = obs.consistent.is_some_and(|c| c[src] == 0);
let was_provisional = bit(t.provisional, dst);
let write = if fresh {
true
} else if bad {
false
} else {
was_provisional || replace
};
if !write {
if bad && !was_provisional && px32(t.pixels, dst) == px32(obs.rgba, src) {
set_bit(t.provisional, dst);
sum.provisional_delta += 1;
provisional_in_block += 1;
sum.changed = true;
}
continue;
}
t.pixels[dst * 4..dst * 4 + 4].copy_from_slice(&obs.rgba[src * 4..src * 4 + 4]);
if fresh {
set_bit(t.coverage, dst);
sum.added += 1;
if obs.uncertain {
sum.uncertain += 1;
}
}
if bad {
if !was_provisional {
set_bit(t.provisional, dst);
sum.provisional_delta += 1;
provisional_in_block += 1;
}
} else if was_provisional {
clear_bit(t.provisional, dst);
sum.provisional_delta -= 1;
provisional_in_block -= 1;
}
sum.changed = true;
}
}
}
let before = (t.quality[q], t.owner[q], t.score[q].to_bits());
let quality = js_round(obs.confidence * 255.0) as u8;
if replace || t.owner[q] == 0 {
t.quality[q] = quality;
t.owner[q] = obs.frame + 1;
t.score[q] = score as f32;
} else if obs.uncertain {
t.quality[q] = if t.owner[q] != 0 {
t.quality[q].min(quality)
} else {
quality
};
}
if had_provisional && provisional_in_block == 0 && !obs.uncertain {
t.quality[q] = t.quality[q].max(quality);
}
if provisional_in_block > 0 {
t.quality[q] = t.quality[q].min(64);
}
if (t.quality[q], t.owner[q], t.score[q].to_bits()) != before {
sum.changed = true;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(not(target_feature = "simd128"))]
#[test]
fn block_row_stats_match_the_per_pixel_rule() {
let mut seed = 0x1234_5678u32;
let mut next = move || {
seed ^= seed << 13;
seed ^= seed >> 17;
seed ^= seed << 5;
seed
};
let n = QUALITY_BLOCK;
for round in 0..2000 {
let (dst0, src0) = (n * (round % 3), 5 + round % 7);
let mut pixels = vec![0u8; (dst0 + n) * 4];
let mut rgba = vec![0u8; (src0 + n) * 4];
for v in pixels.iter_mut().chain(rgba.iter_mut()) {
*v = next() as u8;
}
for i in 0..n {
let (a, b) = ((dst0 + i) * 4, (src0 + i) * 4);
match next() % 5 {
0 => pixels[a..a + 4].copy_from_slice(&rgba[b..b + 4]),
1 => {
pixels[a..a + 4].copy_from_slice(&rgba[b..b + 4]);
pixels[a + 3] ^= 1;
}
k @ (2 | 3) => {
let base = [100u8, 120, 140, 255];
rgba[b..b + 4].copy_from_slice(&base);
let d = if k == 2 { 75 } else { 76 };
pixels[a..a + 4].copy_from_slice(&[
100 + 25,
120 - 25,
140 + (d - 50) as u8,
255,
]);
}
_ => {}
}
}
let coverage: Vec<u8> = (0..(dst0 + n) / 8).map(|_| next() as u8).collect();
let labels: Vec<u8> = (0..src0 + n).map(|_| (next() % 3) as u8).collect();
for with_labels in [false, true] {
let obs = Observation {
rgba: &rgba,
width: src0 + n,
height: 1,
labels: with_labels.then_some((&labels[..], 1)),
occlusions: &[],
consistent: None,
confidence: 1.0,
uncertain: false,
frame: 0,
latest: false,
};
let mut want = (0, 0, 0, 0);
for i in 0..n {
let (dst, src) = (dst0 + i, src0 + i);
if with_labels && labels[src] != 1 {
continue;
}
want.0 += 1;
if bit(&coverage, dst) {
want.1 += 1;
if px32(&pixels, dst) == px32(&rgba, src) {
want.2 += 1;
continue;
}
let (i, j) = (dst * 4, src * 4);
let diff = ((pixels[i] as i32 - rgba[j] as i32).abs()
+ (pixels[i + 1] as i32 - rgba[j + 1] as i32).abs()
+ (pixels[i + 2] as i32 - rgba[j + 2] as i32).abs())
as f64
/ 3.0;
if diff > 25.0 {
want.3 += 1;
}
}
}
assert_eq!(
row_stats_scalar(&obs, &pixels, &coverage, dst0, src0),
want,
"round {round} labels {with_labels}"
);
}
}
}
#[test]
fn fresh_pixels_are_written_and_counted() {
let size = 32usize;
let mut pixels = vec![0u8; size * size * 4];
let mut coverage = vec![0u8; size * size / 8];
let mut provisional = vec![0u8; size * size / 8];
let blocks = (size / QUALITY_BLOCK).pow(2);
let (mut quality, mut conflicts, mut owner, mut score, frozen) = (
vec![0u8; blocks],
vec![0u8; blocks],
vec![0u32; blocks],
vec![0f32; blocks],
vec![0u8; blocks],
);
let mut tile = TileBuffers {
size,
pixels: &mut pixels,
coverage: &mut coverage,
provisional: &mut provisional,
quality: &mut quality,
conflicts: &mut conflicts,
owner: &mut owner,
score: &mut score,
frozen: &frozen,
};
let rgba: Vec<u8> = (0..20 * 10).flat_map(|i| [i as u8, 7, 9, 255]).collect();
let obs = Observation {
rgba: &rgba,
width: 20,
height: 10,
labels: None,
occlusions: &[],
consistent: None,
confidence: 0.9,
uncertain: false,
frame: 3,
latest: false,
};
let world = Rect {
x: 5.0,
y: 6.0,
width: 20.0,
height: 10.0,
};
let stats = composite_tile(&mut tile, &obs, world, 5, 6, 0, 0);
assert_eq!(stats.added, 200);
assert!(stats.changed && stats.conflict_blocks.is_empty());
assert_eq!(tile.owner[0], 4);
assert_eq!(tile.pixels[(6 * size + 5) * 4], 0);
assert_eq!(tile.pixels[(6 * size + 6) * 4], 1);
}
}