use std::{error::Error, fmt};
use crate::representation::{EventFrame, EventFrameData, RepresentationKind};
impl EventFrame {
pub fn connected_components(&self, connectivity: u8) -> Result<EventFrame, ClusterError> {
if connectivity != 4 && connectivity != 8 {
return Err(ClusterError::InvalidConnectivity(connectivity));
}
let (channels, height, width) = self.shape();
let plane = width
.checked_mul(height)
.ok_or(ClusterError::SizeOverflow)?;
let foreground = foreground_mask(self.data(), plane, channels);
let mut labels = vec![0_u64; plane];
let mut next = 0_u64;
let mut stack = Vec::new();
let neighbours = neighbour_offsets(connectivity);
for start in 0..plane {
if !foreground[start] || labels[start] != 0 {
continue;
}
next += 1;
labels[start] = next;
stack.push(start);
while let Some(index) = stack.pop() {
let (x, y) = (index % width, index / width);
for &(dx, dy) in neighbours {
let (nx, ny) = (x as isize + dx, y as isize + dy);
if nx < 0 || ny < 0 || nx >= width as isize || ny >= height as isize {
continue;
}
let neighbour = ny as usize * width + nx as usize;
if foreground[neighbour] && labels[neighbour] == 0 {
labels[neighbour] = next;
stack.push(neighbour);
}
}
}
}
Ok(EventFrame::from_parts(
EventFrameData::U64(labels),
width,
height,
RepresentationKind::Labels,
vec!["labels".to_owned()],
))
}
}
fn foreground_mask(data: &EventFrameData, plane: usize, channels: usize) -> Vec<bool> {
let mut mask = vec![false; plane];
let mut mark = |nonzero: &dyn Fn(usize) -> bool| {
for c in 0..channels {
let base = c * plane;
for (i, slot) in mask.iter_mut().enumerate() {
if nonzero(base + i) {
*slot = true;
}
}
}
};
match data {
EventFrameData::U8(v) => mark(&|k| v[k] != 0),
EventFrameData::U16(v) => mark(&|k| v[k] != 0),
EventFrameData::U64(v) => mark(&|k| v[k] != 0),
EventFrameData::F32(v) => mark(&|k| v[k] != 0.0),
}
mask
}
fn neighbour_offsets(connectivity: u8) -> &'static [(isize, isize)] {
const FOUR: [(isize, isize); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)];
const EIGHT: [(isize, isize); 8] = [
(-1, -1),
(0, -1),
(1, -1),
(-1, 0),
(1, 0),
(-1, 1),
(0, 1),
(1, 1),
];
if connectivity == 4 {
&FOUR
} else {
&EIGHT
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum ClusterError {
SizeOverflow,
InvalidConnectivity(u8),
}
impl fmt::Display for ClusterError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::SizeOverflow => formatter.write_str("label field dimensions are too large"),
Self::InvalidConnectivity(value) => {
write!(formatter, "connectivity must be 4 or 8, got {value}")
}
}
}
}
impl Error for ClusterError {}
#[cfg(test)]
mod tests {
use super::ClusterError;
use crate::representation::{EventFrame, EventFrameData, RepresentationKind};
fn frame(width: usize, height: usize, pixels: &[u8]) -> EventFrame {
EventFrame::from_parts(
EventFrameData::U8(pixels.to_vec()),
width,
height,
RepresentationKind::Count,
vec!["count".to_owned()],
)
}
fn labels(frame: &EventFrame) -> Vec<u64> {
let EventFrameData::U64(values) = frame.data() else {
panic!("labels are u64");
};
values.clone()
}
#[test]
fn rejects_bad_connectivity() {
let f = frame(2, 2, &[0, 0, 0, 0]);
assert_eq!(
f.connected_components(6).unwrap_err(),
ClusterError::InvalidConnectivity(6)
);
}
#[test]
fn empty_frame_is_all_background() {
let out = frame(3, 3, &[0; 9]).connected_components(8).unwrap();
assert_eq!(out.shape(), (1, 3, 3));
assert!(labels(&out).iter().all(|&l| l == 0));
}
#[test]
fn two_diagonal_blobs_split_under_4_but_merge_under_8() {
let pixels = [1, 0, 0, 1];
let four = frame(2, 2, &pixels).connected_components(4).unwrap();
assert_eq!(labels(&four), vec![1, 0, 0, 2]);
let eight = frame(2, 2, &pixels).connected_components(8).unwrap();
assert_eq!(labels(&eight), vec![1, 0, 0, 1]);
}
#[test]
fn foreground_is_any_nonzero_channel() {
let data = EventFrameData::U8(vec![0, 0, 0, 1]);
let f = EventFrame::from_parts(
data,
2,
1,
RepresentationKind::Polarity,
vec!["a".to_owned(), "b".to_owned()],
);
assert_eq!(labels(&f.connected_components(4).unwrap()), vec![0, 1]);
}
}