201 lines
7.0 KiB
Rust
201 lines
7.0 KiB
Rust
pub trait NdRoi<T, D>: seal::Sealed {
|
|
fn roi(&self, rect: bbox::BBox<usize>) -> ndarray::ArrayView<T, D>;
|
|
}
|
|
|
|
pub trait NdRoiMut<T, D>: seal::Sealed {
|
|
fn roi_mut(&mut self, rect: bbox::BBox<usize>) -> ndarray::ArrayViewMut<T, D>;
|
|
}
|
|
|
|
mod seal {
|
|
use ndarray::{Ix2, Ix3};
|
|
pub trait Sealed {}
|
|
impl<T: bytemuck::Pod, S: ndarray::Data<Elem = T>> Sealed for ndarray::ArrayBase<S, Ix2> {}
|
|
impl<T: bytemuck::Pod, S: ndarray::Data<Elem = T>> Sealed for ndarray::ArrayBase<S, Ix3> {}
|
|
}
|
|
|
|
impl<T: bytemuck::Pod, S: ndarray::Data<Elem = T>> NdRoi<T, ndarray::Ix3>
|
|
for ndarray::ArrayBase<S, ndarray::Ix3>
|
|
{
|
|
fn roi(&self, rect: bbox::BBox<usize>) -> ndarray::ArrayView3<T> {
|
|
let y1 = rect.y1();
|
|
let y2 = rect.y2();
|
|
let x1 = rect.x1();
|
|
let x2 = rect.x2();
|
|
self.slice(ndarray::s![y1..y2, x1..x2, ..])
|
|
}
|
|
}
|
|
|
|
impl<T: bytemuck::Pod, S: ndarray::DataMut<Elem = T>> NdRoiMut<T, ndarray::Ix3>
|
|
for ndarray::ArrayBase<S, ndarray::Ix3>
|
|
{
|
|
fn roi_mut(&mut self, rect: bbox::BBox<usize>) -> ndarray::ArrayViewMut3<T> {
|
|
let y1 = rect.y1();
|
|
let y2 = rect.y2();
|
|
let x1 = rect.x1();
|
|
let x2 = rect.x2();
|
|
self.slice_mut(ndarray::s![y1..y2, x1..x2, ..])
|
|
}
|
|
}
|
|
|
|
impl<T: bytemuck::Pod, S: ndarray::Data<Elem = T>> NdRoi<T, ndarray::Ix2>
|
|
for ndarray::ArrayBase<S, ndarray::Ix2>
|
|
{
|
|
fn roi(&self, rect: bbox::BBox<usize>) -> ndarray::ArrayView2<T> {
|
|
let y1 = rect.y1();
|
|
let y2 = rect.y2();
|
|
let x1 = rect.x1();
|
|
let x2 = rect.x2();
|
|
self.slice(ndarray::s![y1..y2, x1..x2])
|
|
}
|
|
}
|
|
|
|
impl<T: bytemuck::Pod, S: ndarray::DataMut<Elem = T>> NdRoiMut<T, ndarray::Ix2>
|
|
for ndarray::ArrayBase<S, ndarray::Ix2>
|
|
{
|
|
fn roi_mut(&mut self, rect: bbox::BBox<usize>) -> ndarray::ArrayViewMut2<T> {
|
|
let y1 = rect.y1();
|
|
let y2 = rect.y2();
|
|
let x1 = rect.x1();
|
|
let x2 = rect.x2();
|
|
self.slice_mut(ndarray::s![y1..y2, x1..x2])
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_roi() {
|
|
let arr = ndarray::Array3::<u8>::zeros((10, 10, 3));
|
|
let rect = bbox::BBox::new(1, 1, 3, 3);
|
|
let roi = arr.roi(rect);
|
|
assert_eq!(roi.shape(), &[3, 3, 3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_roi_2d() {
|
|
let arr = ndarray::Array2::<u8>::zeros((10, 10));
|
|
let rect = bbox::BBox::new(1, 1, 3, 3);
|
|
let roi = arr.roi(rect);
|
|
assert_eq!(roi.shape(), &[3, 3]);
|
|
}
|
|
|
|
/// ```text
|
|
/// ┌──────────────────┐
|
|
/// │ padded │
|
|
/// │ ┌────────┐ │
|
|
/// │ │ │ │
|
|
/// │ │original│ │
|
|
/// │ │ │ │
|
|
/// │ └────────┘ │
|
|
/// │ zeroed │
|
|
/// └──────────────────┘
|
|
/// ```
|
|
///
|
|
/// Returns the padded bounding box and the padded segment
|
|
/// The padded is the padded bounding box
|
|
/// The original is the original bounding box
|
|
/// Returns the padded bounding box as zeros and the original bbox as the original segment
|
|
pub trait NdRoiZeroPadded<T, D: ndarray::Dimension>: seal::Sealed {
|
|
fn roi_zero_padded(
|
|
&self,
|
|
original: bbox::BBox<usize>,
|
|
padded: bbox::BBox<usize>,
|
|
) -> (bbox::BBox<usize>, ndarray::Array<T, D>);
|
|
}
|
|
|
|
impl<T: bytemuck::Pod + num::Zero> NdRoiZeroPadded<T, ndarray::Ix2> for ndarray::Array2<T> {
|
|
fn roi_zero_padded(
|
|
&self,
|
|
original: bbox::BBox<usize>,
|
|
padded: bbox::BBox<usize>,
|
|
) -> (bbox::BBox<usize>, ndarray::Array2<T>) {
|
|
// The co-ordinates of both the original and the padded bounding boxes must be contained in
|
|
// self or it will panic
|
|
|
|
let self_bbox = bbox::BBox::new(0, 0, self.shape()[1], self.shape()[0]);
|
|
if !self_bbox.contains_bbox(original) {
|
|
panic!("original bounding box is not contained in self");
|
|
}
|
|
if !self_bbox.contains_bbox(padded) {
|
|
panic!("padded bounding box is not contained in self");
|
|
}
|
|
|
|
let original_roi_in_padded =
|
|
original.with_top_left(original.top_left().with_origin(padded.top_left()));
|
|
let original_segment = self.roi(original);
|
|
let mut padded_segment = padded.zeros_ndarray_2d::<T>();
|
|
padded_segment
|
|
.roi_mut(original_roi_in_padded)
|
|
.assign(&original_segment);
|
|
|
|
(padded, padded_segment)
|
|
}
|
|
}
|
|
|
|
impl<T: bytemuck::Pod + num::Zero> NdRoiZeroPadded<T, ndarray::Ix3> for ndarray::Array3<T> {
|
|
fn roi_zero_padded(
|
|
&self,
|
|
original: bbox::BBox<usize>,
|
|
padded: bbox::BBox<usize>,
|
|
) -> (bbox::BBox<usize>, ndarray::Array3<T>) {
|
|
let self_bbox = bbox::BBox::new(0, 0, self.shape()[1], self.shape()[0]);
|
|
if !self_bbox.contains_bbox(original) {
|
|
panic!("original bounding box is not contained in self");
|
|
}
|
|
if !self_bbox.contains_bbox(padded) {
|
|
panic!("padded bounding box is not contained in self");
|
|
}
|
|
|
|
let original_roi_in_padded =
|
|
original.with_top_left(original.top_left().with_origin(padded.top_left()));
|
|
let original_segment = self.roi(original);
|
|
let mut padded_segment = padded.zeros_ndarray_3d::<T>(self.len_of(ndarray::Axis(2)));
|
|
padded_segment
|
|
.roi_mut(original_roi_in_padded)
|
|
.assign(&original_segment);
|
|
|
|
(padded, padded_segment)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_roi_zero_padded() {
|
|
let arr = ndarray::Array2::<u8>::ones((10, 10));
|
|
let original = bbox::BBox::new(1, 1, 3, 3);
|
|
let clamp = bbox::BBox::new(0, 0, 10, 10);
|
|
let padded = original.padding(2).clamp_box(clamp);
|
|
let (padded, padded_segment) = arr.roi_zero_padded(original.cast(), padded.cast());
|
|
assert_eq!(padded, bbox::BBox::new(0, 0, 6, 6));
|
|
assert_eq!(padded_segment.shape(), &[6, 6]);
|
|
}
|
|
|
|
#[test]
|
|
pub fn bbox_clamp_failure_preload() {
|
|
let segment_mask = ndarray::Array2::<u8>::zeros((512, 512));
|
|
let og = bbox::BBox::new(475.0, 79.625, 37.0, 282.15);
|
|
let clamp = bbox::BBox::new(0.0, 0.0, 512.0, 512.0);
|
|
let padded = og.scale(1.2).clamp_box(clamp);
|
|
let padded = padded.round();
|
|
let (_bbox, _segment) = segment_mask.roi_zero_padded(og.cast(), padded.cast());
|
|
}
|
|
|
|
#[test]
|
|
pub fn bbox_clamp_failure_preload_2() {
|
|
let segment_mask = ndarray::Array2::<u8>::zeros((512, 512));
|
|
let bbox = bbox::BBox::new(354.25, 98.5, 116.25, 413.5);
|
|
// let padded = bbox::BBox::new(342.625, 57.150000000000006, 139.5, 454.85);
|
|
let clamp = bbox::BBox::new(0.0, 0.0, 512.0, 512.0);
|
|
let padded = bbox.scale(1.2).clamp_box(clamp);
|
|
let padded = padded.round();
|
|
let (_bbox, _segment) = segment_mask.roi_zero_padded(bbox.round().cast(), padded.cast());
|
|
}
|
|
|
|
#[test]
|
|
fn test_roi_zero_padded_3d() {
|
|
let arr = ndarray::Array3::<u8>::ones((10, 10, 3));
|
|
let original = bbox::BBox::new(1, 1, 3, 3);
|
|
let clamp = bbox::BBox::new(0, 0, 10, 10);
|
|
let padded = original.padding(2).clamp_box(clamp);
|
|
let (padded, padded_segment) = arr.roi_zero_padded(original.cast(), padded.cast());
|
|
assert_eq!(padded, bbox::BBox::new(0, 0, 6, 6));
|
|
assert_eq!(padded_segment.shape(), &[6, 6, 3]);
|
|
}
|