Skip to content

Commit 8a67a6a

Browse files
jasperdewintherjasperdewinther
committed
now also has decoding
1 parent 1988d0a commit 8a67a6a

2 files changed

Lines changed: 165 additions & 50 deletions

File tree

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ log = "0.4"
3333
thiserror = "1"
3434
rav1e = { version = "0.8", default-features = false }
3535
ivf = { version = "0.1", default-features = false }
36+
rav1d-safe = "0.3.1"
3637

3738
[dev-dependencies]
3839
bevy = { version = "0.11", default-features = false, features = ["bevy_pbr"] }

‎examples/bevy.rs‎

Lines changed: 164 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use half::f16;
77
use rav1e::{
88
config::SpeedSettings, prelude::v_frame, Config, Context, EncoderConfig, EncoderStatus,
99
};
10+
use rav1d_safe::src::managed::{Decoder, Planes};
1011
use smooth_bevy_cameras::{controllers::unreal::*, LookTransformPlugin};
1112
use vdb_rs::VdbReader;
1213

@@ -52,7 +53,7 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
5253
let mut vdb_reader = VdbReader::new(BufReader::new(f)).unwrap();
5354
let grid_names = vdb_reader.available_grids();
5455

55-
let grid_to_load = std::env::args().nth(2).unwrap_or_else(|| {
56+
let _grid_to_load = std::env::args().nth(2).unwrap_or_else(|| {
5657
println!(
5758
"Grid name not specified, defaulting to first available grid.\nAvailable grids: {:?}",
5859
grid_names
@@ -68,9 +69,9 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
6869
.unwrap();
6970
let extent = max_aabb - min_aabb + 1;
7071

71-
info!("Extent: {}", extent);
72+
info!("min_aabb: {}, max_aabb: {}, extent: {}", min_aabb, max_aabb, extent);
7273

73-
// MY GLORIOUS COMPRESSION
74+
// ENCODE: VDB -> AV1/IVF
7475
{
7576
let grid = vdb_reader
7677
.read_grid::<half::f16>(&grid_names.first().unwrap())
@@ -85,7 +86,8 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
8586
dbg!("Error pos: {}", pos);
8687
}
8788
let actual_index =
88-
(pos + Vec3::new(x as f32, y as f32, z as f32)).as_uvec3();
89+
(pos + Vec3::new(x as f32, y as f32, z as f32)).as_uvec3()
90+
- min_aabb.as_uvec3();
8991
let id = actual_index.x
9092
+ actual_index.y * extent.x as u32
9193
+ actual_index.z * extent.x as u32 * extent.y as u32;
@@ -130,23 +132,17 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
130132
}
131133
Err(EncoderStatus::Encoded) => (),
132134
Err(EncoderStatus::LimitReached) => {
133-
ctx.flush();
134135
break;
135136
}
136137
Err(EncoderStatus::NeedMoreData) => {
137-
info!("Requesting more data");
138-
frame_idx += 1;
139138
if frame_idx >= slices {
140139
ctx.flush();
141-
142-
// we need to continue the loop here to receive more packets
143140
continue;
144141
}
145142
info!("Writing frame {} out of {}", frame_idx, slices);
146143

147144
let mut frame = f.clone();
148145
for (plane_idx, plane) in frame.planes.iter_mut().enumerate() {
149-
info!("Plane idx: {}", plane_idx);
150146
if plane_idx != 0 {
151147
continue;
152148
}
@@ -159,39 +155,13 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
159155
for (col, dst) in dst_row.iter_mut().take(width).enumerate() {
160156
let multiplier = ((1 << 8) - 1) as f32;
161157
let id = col + row_idx * width + frame_idx * width * height;
162-
if voxel_grid[id] != f16::ZERO {
163-
// dbg!(
164-
// col,
165-
// row_idx,
166-
// frame_idx,
167-
// id,
168-
// voxel_grid[id],
169-
// voxel_grid[id].to_f32() * multiplier
170-
// );
171-
}
172-
match plane_idx {
173-
0 => {
174-
*dst = (voxel_grid[id].to_f32() * multiplier * 10.0)
175-
.clamp(0.0, 255.0)
176-
as u16
177-
}
178-
1 => {
179-
*dst = (voxel_grid[id].to_f32() * multiplier * 10.0)
180-
.clamp(0.0, 255.0)
181-
as u16
182-
}
183-
2 => {
184-
*dst = (voxel_grid[id].to_f32() * multiplier * 10.0)
185-
.clamp(0.0, 255.0)
186-
as u16
187-
}
188-
_ => {}
189-
}
158+
*dst = (voxel_grid[id].to_f32() * multiplier * 10.0)
159+
.clamp(0.0, 255.0)
160+
as u16;
190161
}
191162
}
192163
}
193164

194-
info!("Sending frame");
195165
match ctx.send_frame(frame) {
196166
Ok(_) => {}
197167
Err(EncoderStatus::EnoughData) => {
@@ -201,23 +171,123 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
201171
panic!("Unable to send frame {frame_idx}: {e:?}");
202172
}
203173
}
174+
frame_idx += 1;
204175
}
205176
_ => {}
206177
}
207178
}
208179
}
209180

210-
let grid = vdb_reader.read_grid::<half::f16>(&grid_to_load).unwrap();
211-
let instances: Vec<Cuboid> = grid
212-
.iter()
213-
.map(|(pos, voxel, level)| {
214-
Cuboid::new(
215-
pos * 0.1,
216-
(pos + level.scale()) * 0.1,
217-
u32::from_le_bytes(f32::to_le_bytes(voxel.to_f32())),
218-
)
219-
})
220-
.collect();
181+
// DECODE: IVF/AV1 -> voxel grid
182+
let decoded_voxels = {
183+
let mut ivf_file =
184+
BufReader::new(std::fs::File::open("out.ivf").expect("Failed to open out.ivf"));
185+
186+
let header = ivf::read_header(&mut ivf_file).expect("Failed to read IVF header");
187+
let width = header.w as usize;
188+
let height = header.h as usize;
189+
let slices = extent.z as usize;
190+
191+
info!(
192+
"Decoding IVF: {}x{}, expecting {} slices",
193+
width, height, slices
194+
);
195+
196+
let mut decoder = Decoder::new().expect("Failed to create rav1d decoder");
197+
198+
let mut decoded_grid = vec![f16::ZERO; width * height * slices];
199+
let mut frames_decoded = 0usize;
200+
201+
// Read all IVF packets
202+
let mut packets = Vec::new();
203+
loop {
204+
match ivf::read_packet(&mut ivf_file) {
205+
Ok(pkt) => packets.push(pkt),
206+
Err(_) => break,
207+
}
208+
}
209+
info!("Read {} packets from IVF", packets.len());
210+
211+
for pkt in &packets {
212+
match decoder.decode(&pkt.data) {
213+
Ok(Some(frame)) => {
214+
decode_frame_to_grid(&frame, &mut decoded_grid, width, height, slices, frames_decoded);
215+
frames_decoded += 1;
216+
info!("Decoded frame {} / {}", frames_decoded, slices);
217+
}
218+
Ok(None) => {}
219+
Err(e) => panic!("Decode error: {e}"),
220+
}
221+
222+
// Drain any buffered frames
223+
loop {
224+
match decoder.get_frame() {
225+
Ok(Some(frame)) => {
226+
decode_frame_to_grid(&frame, &mut decoded_grid, width, height, slices, frames_decoded);
227+
frames_decoded += 1;
228+
info!("Decoded frame {} / {}", frames_decoded, slices);
229+
}
230+
Ok(None) => break,
231+
Err(e) => panic!("Decode error draining: {e}"),
232+
}
233+
}
234+
}
235+
236+
// Flush remaining frames
237+
match decoder.flush() {
238+
Ok(remaining_frames) => {
239+
for frame in &remaining_frames {
240+
decode_frame_to_grid(frame, &mut decoded_grid, width, height, slices, frames_decoded);
241+
frames_decoded += 1;
242+
info!("Decoded frame {} / {} (flushed)", frames_decoded, slices);
243+
}
244+
}
245+
Err(e) => info!("Flush completed with: {e}"),
246+
}
247+
248+
info!(
249+
"Decoding complete: {} frames decoded, grid size: {}",
250+
frames_decoded,
251+
decoded_grid.len()
252+
);
253+
254+
decoded_grid
255+
};
256+
257+
// Visualize the decoded voxels
258+
let width = extent.x as usize;
259+
let height = extent.y as usize;
260+
let slices = extent.z as usize;
261+
262+
// Use a threshold to filter out AV1 lossy compression artifacts.
263+
// AV1 at speed 9 creates ringing at boundaries; pixel values of ~5-15
264+
// are common artifacts, which decode to 0.002-0.006. Use 0.005 to be safe.
265+
let artifact_threshold = 0.008;
266+
267+
let mut instances: Vec<Cuboid> = Vec::new();
268+
for z in 0..slices {
269+
for y in 0..height {
270+
for x in 0..width {
271+
let id = x + y * width + z * width * height;
272+
let voxel = decoded_voxels[id];
273+
if voxel.to_f32() > artifact_threshold {
274+
let pos = Vec3::new(
275+
x as f32 + min_aabb.x as f32,
276+
y as f32 + min_aabb.y as f32,
277+
z as f32 + min_aabb.z as f32,
278+
);
279+
instances.push(Cuboid::new(
280+
pos * 0.1,
281+
(pos + 1.0) * 0.1,
282+
u32::from_le_bytes(f32::to_le_bytes(voxel.to_f32())),
283+
));
284+
}
285+
}
286+
}
287+
}
288+
289+
info!("Rendering {} decoded voxels", instances.len());
290+
221291
let cuboids = Cuboids::new(instances);
222292
let aabb = cuboids.aabb();
223293
commands
@@ -243,3 +313,47 @@ fn setup(mut commands: Commands, mut color_options_map: ResMut<CuboidMaterialMap
243313
Vec3::Y,
244314
));
245315
}
316+
317+
fn decode_frame_to_grid(
318+
frame: &rav1d_safe::src::managed::Frame,
319+
decoded_grid: &mut [f16],
320+
width: usize,
321+
height: usize,
322+
slices: usize,
323+
frame_idx: usize,
324+
) {
325+
if frame_idx >= slices {
326+
return;
327+
}
328+
329+
let multiplier = ((1 << 8) - 1) as f32;
330+
331+
match frame.planes() {
332+
Planes::Depth8(planes) => {
333+
let y_plane = planes.y();
334+
for row in 0..y_plane.height().min(height) {
335+
let row_data = y_plane.row(row);
336+
for col in 0..row_data.len().min(width) {
337+
let pixel = row_data[col];
338+
// Reverse the encoding: encoded = (f16 * 255.0 * 10.0).clamp(0, 255)
339+
// So: f16 = pixel / 255.0 / 10.0
340+
let value = pixel as f32 / multiplier / 10.0;
341+
let id = col + row * width + frame_idx * width * height;
342+
decoded_grid[id] = f16::from_f32(value);
343+
}
344+
}
345+
}
346+
Planes::Depth16(planes) => {
347+
let y_plane = planes.y();
348+
for row in 0..y_plane.height().min(height) {
349+
let row_data = y_plane.row(row);
350+
for col in 0..row_data.len().min(width) {
351+
let pixel = row_data[col];
352+
let value = pixel as f32 / multiplier / 10.0;
353+
let id = col + row * width + frame_idx * width * height;
354+
decoded_grid[id] = f16::from_f32(value);
355+
}
356+
}
357+
}
358+
}
359+
}

0 commit comments

Comments
 (0)