Skip to content

Commit

Permalink
improved npz reader
Browse files Browse the repository at this point in the history
  • Loading branch information
@KeKsBoTer
KeKsBoTer committed Mar 22, 2024
1 parent b668d45 commit e4cca96
Show file tree
Hide file tree
Showing 2 changed files with 138 additions and 108 deletions.
2 changes: 1 addition & 1 deletion public/demo.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
<div class="container is-max-desktop is-mobile">
<div class="column">
<h1 class="title is-3">Web-Splat - Demo Scenes</h1>
<p class="subtitle is-5">Scenes compressed with <a href="https://keksboter.github.io/3dgs">Compressed 3D Gaussian Splatting for Accelerated Novel View Synthesis</a></p>
<p class="subtitle is-5">Scenes compressed with <a href="https://keksboter.github.io/c3dgs">Compressed 3D Gaussian Splatting for Accelerated Novel View Synthesis</a></p>
</div>
<div class="column" style="margin-bottom: 1rem;">
<a href="https://jonbarron.info/mipnerf360/"><h2 class="title is-3 is-link" style="color:var(--link)">Mip-NeRF360</h2></a>
Expand Down
244 changes: 137 additions & 107 deletions src/npz.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,20 +47,11 @@ impl<'a, R: Read + Seek> NpzReader<'a, R> {
(b"\x50\x4B\x03\x04").as_bytes()
}
}
fn get_npz_const<'a, T: npyz::Deserialize + Copy, R: Read + Seek>(
npz_file: &mut NpzArchive<&'a mut R>,
field_name: &str,
) -> Result<T, anyhow::Error> {
Ok(npz_file
.by_name(field_name)?
.ok_or(anyhow::anyhow!("field not present in npz file"))?
.into_vec::<T>()?[0]
.clone())
}

impl<'a, R: Read + Seek> PointCloudReader for NpzReader<'a, R> {
fn read(
&mut self,
sh_deg: u32,
_sh_deg: u32,
) -> Result<
(
Vec<Gaussian>,
Expand All @@ -70,72 +61,70 @@ impl<'a, R: Read + Seek> PointCloudReader for NpzReader<'a, R> {
),
anyhow::Error,
> {
let opacity_scale: f32 = get_npz_const(&mut self.npz_file, "opacity_scale").unwrap_or(1.0);
let now = Instant::now();
let opacity_scale: f32 = get_npz_value(&mut self.npz_file, "opacity_scale")?.unwrap_or(1.0);
let opacity_zero_point: i32 =
get_npz_const(&mut self.npz_file, "opacity_zero_point").unwrap_or(0);
get_npz_value(&mut self.npz_file, "opacity_zero_point")?.unwrap_or(0);

let scaling_scale: f32 = get_npz_const(&mut self.npz_file, "scaling_scale").unwrap_or(1.0);
let scaling_scale: f32 = get_npz_value(&mut self.npz_file, "scaling_scale")?.unwrap_or(1.0);
let scaling_zero_point: f32 =
get_npz_const::<i32, _>(&mut self.npz_file, "scaling_zero_point").unwrap_or(0) as f32;
get_npz_value::<i32>(&mut self.npz_file, "scaling_zero_point")?.unwrap_or(0) as f32;

let rotation_scale: f32 =
get_npz_const(&mut self.npz_file, "rotation_scale").unwrap_or(1.0);
get_npz_value(&mut self.npz_file, "rotation_scale")?.unwrap_or(1.0);
let rotation_zero_point: f32 =
get_npz_const::<i32, _>(&mut self.npz_file, "rotation_zero_point").unwrap_or(0) as f32;
get_npz_value::<i32>(&mut self.npz_file, "rotation_zero_point")?.unwrap_or(0) as f32;

let features_dc_scale: f32 =
get_npz_const(&mut self.npz_file, "features_dc_scale").unwrap_or(1.0);
get_npz_value(&mut self.npz_file, "features_dc_scale")?.unwrap_or(1.0);
let features_dc_zero_point: i32 =
get_npz_const(&mut self.npz_file, "features_dc_zero_point").unwrap_or(0);
get_npz_value(&mut self.npz_file, "features_dc_zero_point")?.unwrap_or(0);

let features_rest_scale: f32 =
get_npz_const(&mut self.npz_file, "features_rest_scale").unwrap_or(1.0);
get_npz_value(&mut self.npz_file, "features_rest_scale")?.unwrap_or(1.0);
let features_rest_zero_point: i32 =
get_npz_const(&mut self.npz_file, "features_rest_zero_point").unwrap_or(0);
get_npz_value(&mut self.npz_file, "features_rest_zero_point")?.unwrap_or(0);

let scaling_factor_scale: f32 =
get_npz_const(&mut self.npz_file, "scaling_factor_scale").unwrap_or(1.);
let scaling_factor_zero_point: i32 =
get_npz_const(&mut self.npz_file, "scaling_factor_zero_point").unwrap_or(0);
let mut scaling_factor: Option<Vec<i8>> = None;
let mut scaling_factor_zero_point: i32 = 0;
let mut scaling_factor_scale: f32 = 1.0;
if self.npz_file.by_name("scaling_factor_scale")?.is_some() {
scaling_factor_scale =
get_npz_value(&mut self.npz_file, "scaling_factor_scale")?.unwrap_or(1.);
scaling_factor_zero_point =
get_npz_value(&mut self.npz_file, "scaling_factor_zero_point")?.unwrap_or(0);

let now = Instant::now();
let xyz: Vec<Point3<f16>> = self
.npz_file
.by_name("xyz")
.unwrap()
.unwrap()
.into_vec()?
scaling_factor = Some(try_get_npz_array(&mut self.npz_file, "scaling_factor")?);
}

let xyz: Vec<Point3<f16>> = try_get_npz_array::<f16>(&mut self.npz_file, "xyz")?
.as_slice()
.chunks_exact(3)
.map(|c: &[f16]| Point3::new(c[0], c[1], c[2]).cast().unwrap())
.collect();

let scaling: Vec<Vector3<f32>> = self
.npz_file
.by_name("scaling")
.unwrap()
.unwrap()
.into_vec()?
.as_slice()
.iter()
.map(|c: &i8| ((*c as f32 - scaling_zero_point) * scaling_scale).max(0.))
.collect::<Vec<f32>>()
.chunks_exact(3)
.map(|c: &[f32]| Vector3::new(c[0], c[1], c[2]).normalize())
.collect();
let scaling: Vec<Vector3<f32>> = if scaling_factor.is_none() {
// if no scaling factor is present, we assume the scaling is not normalized
try_get_npz_array::<i8>(&mut self.npz_file, "scaling")?
.as_slice()
.iter()
.map(|c: &i8| ((*c as f32 - scaling_zero_point) * scaling_scale).exp())
.collect::<Vec<f32>>()
.chunks_exact(3)
.map(|c: &[f32]| Vector3::new(c[0], c[1], c[2]))
.collect()
} else {
try_get_npz_array::<i8>(&mut self.npz_file, "scaling")?
.as_slice()
.iter()
.map(|c: &i8| ((*c as f32 - scaling_zero_point) * scaling_scale).max(0.))
.collect::<Vec<f32>>()
.chunks_exact(3)
.map(|c: &[f32]| Vector3::new(c[0], c[1], c[2]).normalize())
.collect()
};

let scaling_factor: Vec<i8> = self
.npz_file
.by_name("scaling_factor")?
.unwrap()
.into_vec()?;

let rotation: Vec<Quaternion<f32>> = self
.npz_file
.by_name("rotation")
.unwrap()
.unwrap()
.into_vec()?
let rotation: Vec<Quaternion<f32>> = try_get_npz_array(&mut self.npz_file, "rotation")?
.as_slice()
.iter()
.map(|c: &i8| ((*c as f32 - rotation_zero_point) * rotation_scale))
Expand All @@ -144,74 +133,72 @@ impl<'a, R: Read + Seek> PointCloudReader for NpzReader<'a, R> {
.map(|c| Quaternion::new(c[0], c[1], c[2], c[3]).normalize())
.collect();

let opacity = self
.npz_file
.by_name("opacity")?
.unwrap()
.into_vec::<i8>()?;

let features_indices: Vec<u32> = self
.npz_file
.by_name("feature_indices")?
.unwrap()
.into_vec()?
.as_slice()
.iter()
.map(|c: &i32| *c as u32)
.collect::<Vec<u32>>();

let gaussian_indices = self
.npz_file
.by_name("gaussian_indices")?
.unwrap()
.into_vec()?
.as_slice()
.iter()
.map(|c: &i32| *c as u32)
.collect::<Vec<u32>>();

let features_dc: Vec<i8> = self
.npz_file
.by_name("features_dc")
.unwrap()
.unwrap()
.into_vec()?;

let features_rest: Vec<i8> = self
.npz_file
.by_name("features_rest")
.unwrap()
.unwrap()
.into_vec()?;
let opacity: Vec<i8> = try_get_npz_array(&mut self.npz_file, "opacity")?;

let mut feature_indices: Option<Vec<u32>> = None;
if self.npz_file.by_name("feature_indices")?.is_some() {
feature_indices = Some(
try_get_npz_array::<i32>(&mut self.npz_file, "feature_indices")?
.as_slice()
.iter()
.map(|c: &i32| *c as u32)
.collect::<Vec<u32>>(),
);
}

let mut gaussian_indices: Option<Vec<u32>> = None;
if self.npz_file.by_name("gaussian_indices")?.is_some() {
gaussian_indices = Some(
try_get_npz_array::<i32>(&mut self.npz_file, "gaussian_indices")?
.as_slice()
.iter()
.map(|c: &i32| *c as u32)
.collect::<Vec<u32>>(),
);
}

let features_dc: Vec<i8> = try_get_npz_array(&mut self.npz_file, "features_dc")?;

let features_rest: Vec<i8> = try_get_npz_array(&mut self.npz_file, "features_rest")?;

let num_points: usize = xyz.len();
let sh_deg = self.sh_deg;
let num_sh_coeffs = sh_num_coefficients(sh_deg);

let vertices: Vec<Gaussian> = (0..num_points)
let gaussians: Vec<Gaussian> = (0..num_points)
.map(|i| Gaussian {
xyz: xyz[i],
opacity: opacity[i],
scale_factor: scaling_factor[i],
geometry_idx: gaussian_indices[i],
sh_idx: features_indices[i],
scale_factor: match scaling_factor {
Some(ref sf) => sf[i],
None => 0,
},
geometry_idx: match gaussian_indices {
Some(ref gi) => gi[i],
None => i as u32,
},
sh_idx: match feature_indices {
Some(ref fi) => fi[i],
None => i as u32,
},
})
.collect();

let mut sh_buffer = Vec::new();
let mut sh_coefs = Vec::new();

let sh_coeffs_length = num_sh_coeffs as usize * 3;
let rest_num_coefs = sh_coeffs_length - 3;
for i in 0..(features_dc.len() / 3) {
sh_buffer.write_i8(features_dc[i * 3 + 0]).unwrap();
sh_buffer.write_i8(features_dc[i * 3 + 1]).unwrap();
sh_buffer.write_i8(features_dc[i * 3 + 2]).unwrap();
sh_coefs.write_i8(features_dc[i * 3 + 0]).unwrap();
sh_coefs.write_i8(features_dc[i * 3 + 1]).unwrap();
sh_coefs.write_i8(features_dc[i * 3 + 2]).unwrap();
for j in 0..rest_num_coefs {
sh_buffer
sh_coefs
.write_i8(features_rest[i * rest_num_coefs + j])
.unwrap();
}
}
let covar_buffer = (0..rotation.len())
let covars = (0..rotation.len())
.map(|i| {
let cov = build_cov(rotation[i], scaling[i]);
GeometricInfo {
Expand All @@ -231,7 +218,8 @@ impl<'a, R: Read + Seek> PointCloudReader for NpzReader<'a, R> {
scaling_factor: Quantization::new(scaling_factor_zero_point, scaling_factor_scale),
};

return Ok((vertices, sh_buffer, covar_buffer, quantization));

return Ok((gaussians, sh_coefs, covars, quantization));
}

fn file_sh_deg(&self) -> Result<u32, anyhow::Error> {
Expand All @@ -242,3 +230,45 @@ impl<'a, R: Read + Seek> PointCloudReader for NpzReader<'a, R> {
Ok(self.num_points)
}
}


// tries to read an array
// if the array is not present None is returned
fn get_npz_array_optional<T: npyz::Deserialize + Copy>(
reader: &mut NpzArchive<impl Read + Seek>,
field_name: &str,
) -> Result<Option<Vec<T>>, anyhow::Error> {
let a = reader.by_name(field_name)?;
if let Some(a) = a {
let val = a.into_vec::<T>()?;
Ok(Some(val))
} else {
Ok(None)
}
}

// tries to read an array
// if the array is not present it is treated as an error
fn try_get_npz_array<T: npyz::Deserialize + Copy>(
reader: &mut NpzArchive<impl Read + Seek>,
field_name: &str,
) -> Result<Vec<T>, anyhow::Error> {
return Ok(reader
.by_name(field_name)?
.ok_or(anyhow::format_err!("array {field_name} missing"))?
.into_vec::<T>()?);
}

// reads a single optional value
fn get_npz_value<T: npyz::Deserialize + Copy>(
reader: &mut NpzArchive<impl Read + Seek>,
field_name: &str,
) -> Result<Option<T>, anyhow::Error> {
if let Some(arr) = get_npz_array_optional(reader, field_name)? {
arr.get(0)
.ok_or(anyhow::format_err!("array empty"))
.map(|v| Some(*v))
} else {
Ok(None)
}
}

0 comments on commit e4cca96

Please sign in to comment.