use eframe::egui::mutex::Mutex;
use sophus_autodiff::linalg::SVec;
use sophus_image::{
color_map::BlueWhiteRedBlackColorMap,
prelude::*,
ArcImage4U8,
ArcImageF32,
MutImage4U8,
MutImageF32,
};
use crate::camera::ClippingPlanesF32;
pub struct DepthImage {
pub ndc_z_image: ArcImageF32,
pub clipping_planes: ClippingPlanesF32,
color_mapped_cache: Mutex<Option<ArcImage4U8>>,
metric_depth_cache: Mutex<Option<ArcImageF32>>,
}
pub fn ndc_z_to_color(ndc_z: f32) -> SVec<u8, 4> {
if ndc_z == 1.0 {
return SVec::<u8, 4>::new(0, 0, 0, 255);
}
let z = 0.9 * ndc_z;
let rgb = BlueWhiteRedBlackColorMap::f32_to_rgb(z * z);
SVec::<u8, 4>::new(rgb[0], rgb[1], rgb[2], 255)
}
impl DepthImage {
pub fn new(ndc_z_image: ArcImageF32, clipping_planes: ClippingPlanesF32) -> Self {
DepthImage {
ndc_z_image,
clipping_planes,
color_mapped_cache: Mutex::new(None),
metric_depth_cache: Mutex::new(None),
}
}
pub fn color_mapped(&self) -> ArcImage4U8 {
let mut cached_image = self.color_mapped_cache.lock();
match cached_image.as_mut() {
Some(cached_image) => cached_image.clone(),
None => {
let mut image_rgba = MutImage4U8::from_image_size(self.ndc_z_image.image_size());
for v in 0..image_rgba.image_size().height {
for u in 0..image_rgba.image_size().width {
let z = self.ndc_z_image.pixel(u, v);
*image_rgba.mut_pixel(u, v) = ndc_z_to_color(z);
}
}
let shared = ArcImage4U8::from(image_rgba);
*cached_image = Some(shared.clone());
shared
}
}
}
pub fn metric_depth(&self) -> ArcImageF32 {
let mut cached_image = self.metric_depth_cache.lock();
match cached_image.as_mut() {
Some(cached_image) => cached_image.clone(),
None => {
let mut depth = MutImageF32::from_image_size(self.ndc_z_image.image_size());
for v in 0..depth.image_size().height {
for u in 0..depth.image_size().width {
let z = self.ndc_z_image.pixel(u, v);
*depth.mut_pixel(u, v) = self.clipping_planes.metric_z_from_ndc_z(z);
}
}
let shared = ArcImageF32::from(depth);
*cached_image = Some(shared.clone());
shared
}
}
}
}