wifi-densepose/v2/crates/wifi-densepose-pointcloud/src/pointcloud.rs

//! Point cloud types + PLY export + Gaussian splat conversion.
#![allow(dead_code)]

use serde::{Deserialize, Serialize};
use std::io::Write;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Point3D {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ColorPoint {
    pub x: f32,
    pub y: f32,
    pub z: f32,
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub intensity: f32,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PointCloud {
    pub points: Vec<ColorPoint>,
    pub timestamp_ms: i64,
    pub source: String,
}

impl PointCloud {
    pub fn new(source: &str) -> Self {
        Self {
            points: Vec::new(),
            timestamp_ms: chrono::Utc::now().timestamp_millis(),
            source: source.to_string(),
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub fn add(&mut self, x: f32, y: f32, z: f32, r: u8, g: u8, b: u8, intensity: f32) {
        self.points.push(ColorPoint {
            x,
            y,
            z,
            r,
            g,
            b,
            intensity,
        });
    }

    pub fn bounds(&self) -> ([f32; 3], [f32; 3]) {
        if self.points.is_empty() {
            return ([0.0; 3], [0.0; 3]);
        }
        let mut min = [f32::MAX; 3];
        let mut max = [f32::MIN; 3];
        for p in &self.points {
            min[0] = min[0].min(p.x);
            min[1] = min[1].min(p.y);
            min[2] = min[2].min(p.z);
            max[0] = max[0].max(p.x);
            max[1] = max[1].max(p.y);
            max[2] = max[2].max(p.z);
        }
        (min, max)
    }
}

/// Write point cloud to PLY format (ASCII).
pub fn write_ply(cloud: &PointCloud, path: &str) -> anyhow::Result<()> {
    let mut f = std::fs::File::create(path)?;
    writeln!(f, "ply")?;
    writeln!(f, "format ascii 1.0")?;
    writeln!(f, "comment Generated by RuView Dense Point Cloud")?;
    writeln!(f, "comment Source: {}", cloud.source)?;
    writeln!(f, "comment Timestamp: {}", cloud.timestamp_ms)?;
    writeln!(f, "element vertex {}", cloud.points.len())?;
    writeln!(f, "property float x")?;
    writeln!(f, "property float y")?;
    writeln!(f, "property float z")?;
    writeln!(f, "property uchar red")?;
    writeln!(f, "property uchar green")?;
    writeln!(f, "property uchar blue")?;
    writeln!(f, "property float intensity")?;
    writeln!(f, "end_header")?;
    for p in &cloud.points {
        writeln!(
            f,
            "{:.4} {:.4} {:.4} {} {} {} {:.4}",
            p.x, p.y, p.z, p.r, p.g, p.b, p.intensity
        )?;
    }
    Ok(())
}

/// Convert point cloud to Gaussian splats for 3D rendering.
#[derive(Serialize, Deserialize)]
pub struct GaussianSplat {
    pub center: [f32; 3],
    pub color: [f32; 3],
    pub opacity: f32,
    pub scale: [f32; 3],
}

pub fn to_gaussian_splats(cloud: &PointCloud) -> Vec<GaussianSplat> {
    // Cluster points into voxels and create one Gaussian per cluster
    let voxel_size = 0.08; // smaller voxels = more detail = visible movement
    let mut cells: std::collections::HashMap<(i32, i32, i32), Vec<&ColorPoint>> =
        std::collections::HashMap::new();

    for p in &cloud.points {
        let key = (
            (p.x / voxel_size).floor() as i32,
            (p.y / voxel_size).floor() as i32,
            (p.z / voxel_size).floor() as i32,
        );
        cells.entry(key).or_default().push(p);
    }

    cells
        .values()
        .map(|pts| {
            let n = pts.len() as f32;

            // Pass 1 — single fused accumulation of all six sums (position +
            // colour). Replaces six separate `.iter().sum()` passes; identical
            // f32 accumulation order, so the result is bit-for-bit unchanged.
            let (mut sum_x, mut sum_y, mut sum_z) = (0.0f32, 0.0f32, 0.0f32);
            let (mut sum_r, mut sum_g, mut sum_b) = (0.0f32, 0.0f32, 0.0f32);
            for p in pts {
                sum_x += p.x;
                sum_y += p.y;
                sum_z += p.z;
                sum_r += p.r as f32;
                sum_g += p.g as f32;
                sum_b += p.b as f32;
            }
            let cx = sum_x / n;
            let cy = sum_y / n;
            let cz = sum_z / n;
            let cr = sum_r / n / 255.0;
            let cg = sum_g / n / 255.0;
            let cb = sum_b / n / 255.0;

            // Pass 2 — spread (mean absolute deviation) needs the centroid, so
            // it is a second fused pass instead of three separate ones.
            let (mut dev_x, mut dev_y, mut dev_z) = (0.0f32, 0.0f32, 0.0f32);
            for p in pts {
                dev_x += (p.x - cx).abs();
                dev_y += (p.y - cy).abs();
                dev_z += (p.z - cz).abs();
            }
            let sx = dev_x / n + 0.01;
            let sy = dev_y / n + 0.01;
            let sz = dev_z / n + 0.01;

            GaussianSplat {
                center: [cx, cy, cz],
                color: [cr, cg, cb],
                opacity: (n / 10.0).min(1.0),
                scale: [sx, sy, sz],
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty_cloud_has_no_splats() {
        let cloud = PointCloud::new("test");
        assert!(to_gaussian_splats(&cloud).is_empty());
    }

    #[test]
    fn single_voxel_centroid_and_scale_are_correct() {
        // Two points inside the same 0.08 m voxel: (0.01,0.01,0.01) and
        // (0.03,0.03,0.03). Centroid = 0.02 each axis; mean-abs-dev = 0.01;
        // scale = 0.01 + 0.01 = 0.02. Colours: r=0 and r=255 → mean 127.5/255.
        let mut cloud = PointCloud::new("test");
        cloud.add(0.01, 0.01, 0.01, 0, 0, 0, 1.0);
        cloud.add(0.03, 0.03, 0.03, 255, 255, 255, 1.0);

        let splats = to_gaussian_splats(&cloud);
        assert_eq!(splats.len(), 1, "both points fall in one voxel");
        let s = &splats[0];
        for axis in 0..3 {
            assert!((s.center[axis] - 0.02).abs() < 1e-5, "center[{axis}]={}", s.center[axis]);
            assert!((s.scale[axis] - 0.02).abs() < 1e-5, "scale[{axis}]={}", s.scale[axis]);
            assert!((s.color[axis] - 127.5 / 255.0).abs() < 1e-5, "color[{axis}]");
        }
        // opacity = n/10 = 0.2
        assert!((s.opacity - 0.2).abs() < 1e-6);
    }

    #[test]
    fn distinct_voxels_yield_distinct_splats() {
        // Two points far apart → two separate voxels → two splats.
        let mut cloud = PointCloud::new("test");
        cloud.add(0.0, 0.0, 0.0, 10, 20, 30, 1.0);
        cloud.add(1.0, 1.0, 1.0, 40, 50, 60, 1.0);
        assert_eq!(to_gaussian_splats(&cloud).len(), 2);
    }
}