gspx 0.1.2

Sparse graph signal processing and spectral graph wavelets in Rust
Documentation
use std::collections::BTreeSet;
use std::fs::File;
use std::io::{BufRead, BufReader, Read};
use std::path::Path;

use ndarray::Array2;
use sprs::TriMat;

use crate::{Laplacian, error::GspError};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PlyFormat {
    Ascii,
    BinaryLittleEndian,
}

type PlyVertices = Vec<[f64; 3]>;
type PlyFaces = Vec<Vec<usize>>;
type PlyParseResult = (PlyVertices, PlyFaces, usize);

fn parse_ply(path: &Path) -> Result<PlyParseResult, GspError> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);
    let mut format = PlyFormat::Ascii;
    let mut vertex_count = 0usize;
    let mut face_count = 0usize;
    let mut vertex_props = Vec::<(String, String)>::new();
    let mut current_element = String::new();

    loop {
        let mut line = String::new();
        let n = reader.read_line(&mut line)?;
        if n == 0 {
            return Err(GspError::Parse(
                "PLY header terminated unexpectedly".to_string(),
            ));
        }
        let line = line.trim();
        if line == "end_header" {
            break;
        }
        if line.is_empty() {
            continue;
        }
        let parts = line.split_whitespace().collect::<Vec<_>>();
        if parts.is_empty() {
            continue;
        }
        match parts[0] {
            "format" if parts.len() >= 2 => {
                format = match parts[1] {
                    "ascii" => PlyFormat::Ascii,
                    "binary_little_endian" => PlyFormat::BinaryLittleEndian,
                    other => {
                        return Err(GspError::UnsupportedFormat(format!(
                            "unsupported PLY format `{other}`"
                        )));
                    }
                };
            }
            "element" if parts.len() >= 3 => {
                current_element = parts[1].to_string();
                if current_element == "vertex" {
                    vertex_count = parts[2]
                        .parse::<usize>()
                        .map_err(|e| GspError::Parse(e.to_string()))?;
                } else if current_element == "face" {
                    face_count = parts[2]
                        .parse::<usize>()
                        .map_err(|e| GspError::Parse(e.to_string()))?;
                }
            }
            "property" if current_element == "vertex" && parts.len() >= 3 => {
                vertex_props.push((parts[2].to_string(), parts[1].to_string()));
            }
            _ => {}
        }
    }

    match format {
        PlyFormat::Ascii => parse_ply_ascii(reader, vertex_count, face_count),
        PlyFormat::BinaryLittleEndian => {
            parse_ply_binary_le(reader, vertex_count, face_count, &vertex_props)
        }
    }
}

fn parse_ply_ascii(
    mut reader: BufReader<File>,
    vertex_count: usize,
    face_count: usize,
) -> Result<PlyParseResult, GspError> {
    let mut vertices = Vec::<[f64; 3]>::with_capacity(vertex_count);
    let mut faces = Vec::<Vec<usize>>::with_capacity(face_count);
    for _ in 0..vertex_count {
        let mut line = String::new();
        reader.read_line(&mut line)?;
        let parts = line.split_whitespace().collect::<Vec<_>>();
        if parts.len() < 3 {
            return Err(GspError::Parse(
                "PLY vertex line has fewer than 3 components".to_string(),
            ));
        }
        vertices.push([
            parts[0]
                .parse::<f64>()
                .map_err(|e| GspError::Parse(e.to_string()))?,
            parts[1]
                .parse::<f64>()
                .map_err(|e| GspError::Parse(e.to_string()))?,
            parts[2]
                .parse::<f64>()
                .map_err(|e| GspError::Parse(e.to_string()))?,
        ]);
    }
    for _ in 0..face_count {
        let mut line = String::new();
        reader.read_line(&mut line)?;
        let parts = line.split_whitespace().collect::<Vec<_>>();
        if parts.is_empty() {
            continue;
        }
        let n = parts[0]
            .parse::<usize>()
            .map_err(|e| GspError::Parse(e.to_string()))?;
        if parts.len() < n + 1 {
            return Err(GspError::Parse(
                "PLY face line has fewer indices than expected".to_string(),
            ));
        }
        let mut face = Vec::with_capacity(n);
        for i in 0..n {
            face.push(
                parts[i + 1]
                    .parse::<usize>()
                    .map_err(|e| GspError::Parse(e.to_string()))?,
            );
        }
        faces.push(face);
    }
    Ok((vertices, faces, vertex_count))
}

fn parse_ply_binary_le(
    mut reader: BufReader<File>,
    vertex_count: usize,
    face_count: usize,
    vertex_props: &[(String, String)],
) -> Result<PlyParseResult, GspError> {
    let mut vertices = Vec::<[f64; 3]>::with_capacity(vertex_count);
    let mut faces = Vec::<Vec<usize>>::with_capacity(face_count);
    let prop_names = vertex_props
        .iter()
        .map(|(n, _)| n.as_str())
        .collect::<Vec<_>>();
    let xyz = if let (Some(ix), Some(iy), Some(iz)) = (
        prop_names.iter().position(|n| *n == "x"),
        prop_names.iter().position(|n| *n == "y"),
        prop_names.iter().position(|n| *n == "z"),
    ) {
        [ix, iy, iz]
    } else {
        [0, 1, 2]
    };

    for _ in 0..vertex_count {
        let mut vals = Vec::<f64>::with_capacity(vertex_props.len());
        for (_, t) in vertex_props {
            vals.push(read_ply_scalar(&mut reader, t)?);
        }
        if vals.len() < 3 {
            return Err(GspError::Parse(
                "binary PLY vertex has fewer than 3 properties".to_string(),
            ));
        }
        vertices.push([vals[xyz[0]], vals[xyz[1]], vals[xyz[2]]]);
    }
    for _ in 0..face_count {
        let mut nbuf = [0u8; 1];
        reader.read_exact(&mut nbuf)?;
        let n = nbuf[0] as usize;
        let mut face = Vec::with_capacity(n);
        for _ in 0..n {
            let mut ibuf = [0u8; 4];
            reader.read_exact(&mut ibuf)?;
            face.push(i32::from_le_bytes(ibuf) as usize);
        }
        faces.push(face);
    }
    Ok((vertices, faces, vertex_count))
}

fn read_ply_scalar<R: Read>(reader: &mut R, t: &str) -> Result<f64, GspError> {
    match t {
        "char" => {
            let mut b = [0u8; 1];
            reader.read_exact(&mut b)?;
            Ok((b[0] as i8) as f64)
        }
        "uchar" => {
            let mut b = [0u8; 1];
            reader.read_exact(&mut b)?;
            Ok(b[0] as f64)
        }
        "short" => {
            let mut b = [0u8; 2];
            reader.read_exact(&mut b)?;
            Ok(i16::from_le_bytes(b) as f64)
        }
        "ushort" => {
            let mut b = [0u8; 2];
            reader.read_exact(&mut b)?;
            Ok(u16::from_le_bytes(b) as f64)
        }
        "int" => {
            let mut b = [0u8; 4];
            reader.read_exact(&mut b)?;
            Ok(i32::from_le_bytes(b) as f64)
        }
        "uint" => {
            let mut b = [0u8; 4];
            reader.read_exact(&mut b)?;
            Ok(u32::from_le_bytes(b) as f64)
        }
        "float" => {
            let mut b = [0u8; 4];
            reader.read_exact(&mut b)?;
            Ok(f32::from_le_bytes(b) as f64)
        }
        "double" => {
            let mut b = [0u8; 8];
            reader.read_exact(&mut b)?;
            Ok(f64::from_le_bytes(b))
        }
        _ => Err(GspError::UnsupportedFormat(format!(
            "unsupported binary PLY property type `{t}`"
        ))),
    }
}

/// Loads a mesh Laplacian from a PLY file.
///
/// # Errors
/// Returns an error when parsing fails.
pub fn load_ply_laplacian(path: &Path) -> Result<Laplacian, GspError> {
    let (_, faces, n_vertices) = parse_ply(path)?;
    let mut edges = BTreeSet::<(usize, usize)>::new();
    for face in faces {
        if face.len() < 2 {
            continue;
        }
        for i in 0..face.len() {
            let u = face[i];
            let v = face[(i + 1) % face.len()];
            let (a, b) = if u <= v { (u, v) } else { (v, u) };
            edges.insert((a, b));
        }
    }
    let mut tri = TriMat::<f64>::new((n_vertices, edges.len()));
    for (idx, (u, v)) in edges.iter().copied().enumerate() {
        tri.add_triplet(u, idx, 1.0);
        tri.add_triplet(v, idx, -1.0);
    }
    let b = tri.to_csc();
    let l = &b * &b.transpose_view();
    Ok(l.to_csc())
}

/// Loads `(x, y, z)` vertex coordinates from a PLY file.
///
/// # Errors
/// Returns an error when parsing fails.
pub fn load_ply_xyz(path: &Path) -> Result<Array2<f64>, GspError> {
    let (verts, _, _) = parse_ply(path)?;
    let mut out = Array2::<f64>::zeros((verts.len(), 3));
    for (i, xyz) in verts.into_iter().enumerate() {
        out[[i, 0]] = xyz[0];
        out[[i, 1]] = xyz[1];
        out[[i, 2]] = xyz[2];
    }
    Ok(out)
}