1pub mod plane;
2pub mod frame;
3pub mod enc;
4pub mod dec;
5
6mod dct;
7mod common;
8mod huffman;
9mod rle;
10
11#[cfg(test)]
12mod tests {
13 use std::{path::Path, fs::{File, self}, io::{Cursor, Seek, Read}, time::Instant, hint::black_box};
14
15 use bitstream_io::{BitWriter, BitWrite, BitReader, BitRead};
16 use byteorder::{ReadBytesExt, LittleEndian};
17 use image::{io::Reader as ImageReader, RgbImage};
18
19 use crate::{dct::*, frame::VideoFrame, plane::VideoPlane, enc::Encoder, dec::Decoder, rle};
20
21 const DCT_B2_NORMALIZER: [i32;8] = [
22 91, 105, 95, 75, 91, 75, 95, 105
23 ];
24
25 pub fn dct_b2_scale(vector: &mut [i32;8]) {
26 vector[0] = (vector[0] * DCT_B2_NORMALIZER[0]) / 256;
27 vector[1] = (vector[1] * DCT_B2_NORMALIZER[1]) / 256;
28 vector[2] = (vector[2] * DCT_B2_NORMALIZER[2]) / 256;
29 vector[3] = (vector[3] * DCT_B2_NORMALIZER[3]) / 256;
30 vector[4] = (vector[4] * DCT_B2_NORMALIZER[4]) / 256;
31 vector[5] = (vector[5] * DCT_B2_NORMALIZER[5]) / 256;
32 vector[6] = (vector[6] * DCT_B2_NORMALIZER[6]) / 256;
33 vector[7] = (vector[7] * DCT_B2_NORMALIZER[7]) / 256;
34 }
35
36 #[test]
37 fn test_dct_2_fp() {
38 let data: [i32;8] = [0 << 8, 10 << 8, 20 << 8, 30 << 8, 40 << 8, 50 << 8, 60 << 8, 70 << 8];
39
40 let mut dct = data;
41 DctMatrix8x8::fdct(&mut dct);
42 dct_b2_scale(&mut dct);
43
44 println!("DCT: {:?}", dct);
45
46 let mut out = dct;
47 dct_b2_scale(&mut out);
48 DctMatrix8x8::idct(&mut out);
49
50 for i in 0..8 {
51 out[i] >>= 8;
52 }
53
54 println!("Output: {:?}", out);
55 }
56
57 #[test]
58 fn test_dct_encode() {
59 let qtable = [5, 10, 11, 13, 16, 16, 18, 21, 10, 10, 13, 15, 16, 18, 21, 23, 11, 13, 16, 16, 18, 21, 21, 23, 13, 13, 16, 16, 18, 21, 23, 25, 13, 16, 16, 18, 20, 21, 25, 30,
6216, 16, 18, 20, 21, 25, 30, 36, 16, 16, 18, 21, 23, 28, 35, 43, 16, 18, 21, 23, 28, 35, 43, 51];
63
64 let mut dct = DctMatrix8x8::new();
65
66 dct.m = [44, 42, 43, 43, 46, 49, 42, 33, 36, 49, 56, 47, 42, 41, 36, 28, 36, 48, 57, 52, 42, 35, 29, 23, 36, 35, 41, 48, 45, 32, 25, 24, 32, 27, 30, 39, 41, 32, 25, 26, 26, 27, 29, 30, 31, 31, 27, 23, 29, 27, 27, 27, 30, 31, 26, 20, 35, 23, 19, 27, 34, 30, 22, 16];
67
68 println!("Input: {:?}", dct);
69
70 for i in 0..64 {
71 dct.m[i] = (dct.m[i] - 128) << 8;
72 }
73
74 dct.dct_transform_rows();
75 dct.dct_transform_columns();
76
77 println!("Before quant: {:?}", dct.m);
78
79 let qdct = dct.encode(&qtable);
80 println!("Quantized: {:?}", qdct);
81
82 let mut dct2 = DctMatrix8x8::decode(&qdct, &qtable);
83
84 println!("After quant: {:?}", dct2.m);
85
86 dct2.dct_inverse_transform_columns();
87 dct2.dct_inverse_transform_rows();
88
89 for i in 0..64 {
90 dct2.m[i] = (dct2.m[i] >> 8) + 128;
91 }
92
93 println!("Output: {:?}", dct2);
94 }
95
96 #[test]
97 fn test_entropy() {
98 let test_data = [10, 0, 0, 5, 3, 0, 0, 0, 0, -10];
99 let mut rle_sequence = Vec::new();
100 rle::rle_encode(&mut rle_sequence, &test_data);
101
102 let mut table = [0;16];
103 rle::update_table(&mut table, &rle_sequence);
104
105 let tree = rle::rle_create_huffman(&table);
106 let mut tmp_buf = Cursor::new(Vec::new());
107 let mut bitwriter = BitWriter::endian(&mut tmp_buf, bitstream_io::LittleEndian);
108
109 for sq in &rle_sequence {
110 let num_zeroes = tree.get_code(sq.num_zeroes);
111 let num_bits = tree.get_code(sq.coeff_size);
112
113 assert!(num_zeroes.len > 0 && num_bits.len > 0);
114
115 bitwriter.write(num_zeroes.len, num_zeroes.val).unwrap();
116 bitwriter.write(num_bits.len, num_bits.val).unwrap();
117
118 if sq.coeff_size > 0 {
119 bitwriter.write_signed(sq.coeff_size as u32, sq.coeff).unwrap();
120 }
121 }
122
123 bitwriter.byte_align().unwrap();
124
125 let rle_coded = tmp_buf.into_inner();
126
127 println!("Test data encoded to {} bytes", rle_coded.len());
128
129 let mut rle_reader = Cursor::new(rle_coded);
130 let mut bitreader = BitReader::endian(&mut rle_reader, bitstream_io::LittleEndian);
131
132 let total_bits = bitreader.seek_bits(std::io::SeekFrom::End(0)).unwrap();
133 bitreader.seek_bits(std::io::SeekFrom::Start(0)).unwrap();
134
135 let mut out_data = [0;10];
136
137 let mut out_idx = 0;
138 while out_idx < out_data.len() {
139 let num_zeroes = tree.read(&mut bitreader, total_bits).unwrap() as usize;
140 out_idx += num_zeroes;
141
142 let num_bits = tree.read(&mut bitreader, total_bits).unwrap();
143
144 if num_bits > 0 {
146 let coeff = bitreader.read_signed::<i16>(num_bits as u32).unwrap();
147 out_data[out_idx] = coeff;
148
149 out_idx += 1;
150 }
151 }
152
153 println!("RLE decoded to: {:?}", out_data);
154
155 test_data.iter().zip(out_data).for_each(|(a, b)| {
156 assert!(*a == b);
157 });
158 }
159
160 #[test]
161 fn test_entropy_2() {
162 let mut infile = File::open("test_coeff.bin").unwrap();
163 let infile_len = infile.seek(std::io::SeekFrom::End(0)).unwrap() as usize;
164 infile.seek(std::io::SeekFrom::Start(0)).unwrap();
165
166 let mut test_data = vec![0;infile_len / 2];
167
168 for i in 0..test_data.len() {
169 test_data[i] = infile.read_i16::<LittleEndian>().unwrap();
170 }
171
172 let mut rle_sequence = Vec::new();
173 rle::rle_encode(&mut rle_sequence, &test_data);
174
175 let mut table = [0;16];
176 rle::update_table(&mut table, &rle_sequence);
177
178 let tree = rle::rle_create_huffman(&table);
179 let mut tmp_buf = Cursor::new(Vec::new());
180 let mut bitwriter = BitWriter::endian(&mut tmp_buf, bitstream_io::LittleEndian);
181
182 let mut bits_written = 0;
183
184 for sq in &rle_sequence {
185 let num_zeroes = tree.get_code(sq.num_zeroes);
186 let num_bits = tree.get_code(sq.coeff_size);
187
188 assert!(num_zeroes.len > 0 && num_bits.len > 0);
189
190 bitwriter.write(num_zeroes.len, num_zeroes.val).unwrap();
191 bitwriter.write(num_bits.len, num_bits.val).unwrap();
192
193 if sq.coeff_size > 0 {
194 bitwriter.write_signed(sq.coeff_size as u32, sq.coeff).unwrap();
195 }
196
197 bits_written += num_zeroes.len + num_bits.len + sq.coeff_size as u32;
198 }
199
200 bitwriter.byte_align().unwrap();
201
202 let rle_coded = tmp_buf.into_inner();
203
204 println!("Test data encoded ({} bytes -> {} bytes, {} bits)", infile_len, rle_coded.len(), bits_written);
205
206 let mut rle_reader = Cursor::new(rle_coded);
207 let mut bitreader = BitReader::endian(&mut rle_reader, bitstream_io::LittleEndian);
208
209 let total_bits = bitreader.seek_bits(std::io::SeekFrom::End(0)).unwrap();
210 bitreader.seek_bits(std::io::SeekFrom::Start(0)).unwrap();
211
212 let mut out_data = vec![0;test_data.len()];
213
214 let mut out_idx = 0;
215 let mut run_idx = 0;
216 while out_idx < out_data.len() {
217 let num_zeroes = tree.read(&mut bitreader, total_bits).unwrap() as usize;
218 out_idx += num_zeroes;
219
220 let num_bits = tree.read(&mut bitreader, total_bits).unwrap();
221
222 let run = &rle_sequence[run_idx];
223 assert!(run.num_zeroes == num_zeroes as u8);
224 assert!(run.coeff_size == num_bits);
225
226 if num_bits > 0 {
228 let coeff = bitreader.read_signed::<i16>(num_bits as u32).unwrap();
229 out_data[out_idx] = coeff;
230 out_idx += 1;
231 }
232
233 run_idx += 1;
234 }
235
236 test_data.iter().zip(out_data).for_each(|(a, b)| {
237 assert!(*a == b);
238 });
239 }
240
241 #[test]
242 fn test_encode_1() {
243 let test_frame = load_frame("test1.png");
244 let outfile = File::create("test.pfv").unwrap();
245 let mut encoder = Encoder::new(outfile, test_frame.width, test_frame.height, 30, 5, 6).unwrap();
246
247 encoder.encode_iframe(&test_frame).unwrap();
248 encoder.encode_pframe(&test_frame).unwrap();
249 encoder.finish().unwrap();
250
251 println!("File written");
252 }
253
254 #[test]
255 fn test_decode_1() {
256 let infile = File::open("test.pfv").unwrap();
257 let mut decoder = Decoder::new(infile, 6).unwrap();
258
259 let mut outframe = 0;
260
261 while decoder.advance_frame(&mut |frame| {
262 let frame_out_path = format!("test_frames_out/{:0>3}.png", outframe);
264 save_frame(frame_out_path, frame);
265 outframe += 1;
266 }).unwrap() {}
267
268 println!("Decoded {} frames", outframe);
269 }
270
271 #[test]
272 fn test_encode_2() {
273 let outfile = File::create("test2.pfv").unwrap();
274 let mut encoder = Encoder::new(outfile, 512, 384, 30, 2, 6).unwrap();
275
276 for frame_id in 1..162 {
277 let frame_path = format!("test_frames/{:0>3}.png", frame_id);
278 let frame = load_frame(frame_path);
279
280 if (frame_id - 1) % 60 == 0 {
281 encoder.encode_iframe(&frame).unwrap();
282 } else {
283 encoder.encode_pframe(&frame).unwrap();
284 }
285
286 println!("Encoded: {} / {}", frame_id, 162);
287 }
288
289 encoder.finish().unwrap();
290
291 println!("File written");
292 }
293
294 #[test]
295 fn test_decode_2() {
296 let infile = File::open("test2.pfv").unwrap();
297 let mut decoder = Decoder::new(infile, 6).unwrap();
298
299 let mut outframe = 0;
300
301 while decoder.advance_frame(&mut |frame| {
302 let frame_out_path = format!("test_frames_out_2/{:0>3}.png", outframe);
304 save_frame(frame_out_path, frame);
305 outframe += 1;
306 println!("Decoded {}", outframe);
307 }).unwrap() {}
308 }
309
310 #[test]
311 fn test_decode_speed_2() {
312 for run in 0..50 {
313 println!("RUN {}", run);
314
315 let mut infile = File::open("test2.pfv").unwrap();
316 let mut filebuf = Vec::new();
317 infile.read_to_end(&mut filebuf).unwrap();
318
319 let infile = Cursor::new(filebuf);
320
321 let mut decoder = Decoder::new(infile, 6).unwrap();
322
323 let mut outframe = 0;
324
325 let start = Instant::now();
326
327 while decoder.advance_frame(&mut |frame| {
328 outframe += 1;
329 black_box(frame);
330 }).unwrap() {}
331
332 let duration = start.elapsed().as_millis();
333 println!("Decoded {} frames in {} ms", outframe, duration);
334 }
335 }
336
337 fn load_frame<Q: AsRef<Path>>(path: Q) -> VideoFrame {
338 let src_img = ImageReader::open(path).unwrap().decode().unwrap().into_rgb8();
339
340 let yuv_pixels: Vec<[u8;3]> = src_img.pixels().map(|rgb| {
341 let y = (0.299 * rgb.0[0] as f32) + (0.587 * rgb.0[1] as f32) + (0.114 * rgb.0[2] as f32);
343 let u = 128.0 - (0.168736 * rgb.0[0] as f32) - (0.331264 * rgb.0[1] as f32) + (0.5 * rgb.0[2] as f32);
344 let v = 128.0 + (0.5 * rgb.0[0] as f32) - (0.418688 * rgb.0[1] as f32) - (0.081312 * rgb.0[2] as f32);
345 [y as u8, u as u8, v as u8]
346 }).collect();
347
348 let y_buffer: Vec<_> = yuv_pixels.iter().map(|x| x[0]).collect();
350 let u_buffer: Vec<_> = yuv_pixels.iter().map(|x| x[1]).collect();
351 let v_buffer: Vec<_> = yuv_pixels.iter().map(|x| x[2]).collect();
352
353 let y_plane = VideoPlane::from_slice(src_img.width() as usize, src_img.height() as usize, &y_buffer);
354 let u_plane = VideoPlane::from_slice(src_img.width() as usize, src_img.height() as usize, &u_buffer);
355 let v_plane = VideoPlane::from_slice(src_img.width() as usize, src_img.height() as usize, &v_buffer);
356
357 VideoFrame::from_planes(src_img.width() as usize, src_img.height() as usize, y_plane, u_plane, v_plane)
358 }
359
360 fn save_frame<Q: AsRef<Path>>(path: Q, frame: &VideoFrame) {
361 if let Some(parent) = path.as_ref().parent() {
362 fs::create_dir_all(parent).unwrap();
363 }
364
365 let plane_u = frame.plane_u.double();
366 let plane_v = frame.plane_v.double();
367
368 let yuv_pixels: Vec<[u8;3]> = frame.plane_y.pixels.iter().enumerate().map(|(idx, y)| {
369 let y = *y;
370 let u = plane_u.pixels[idx];
371 let v = plane_v.pixels[idx];
372
373 [y, u, v]
374 }).collect();
375
376 let mut rgb_buf: Vec<u8> = Vec::new();
377
378 for yuv in yuv_pixels.iter() {
379 let y = yuv[0] as f32;
380 let u = yuv[1] as f32 - 128.0;
381 let v = yuv[2] as f32 - 128.0;
382
383 let r = y + (1.402 * v);
385 let g = y - (0.344136 * u) - (0.714136 * v);
386 let b = y + (1.772 * u);
387
388 rgb_buf.push(r as u8);
389 rgb_buf.push(g as u8);
390 rgb_buf.push(b as u8);
391 }
392
393 let img_buf = RgbImage::from_vec(frame.width as u32, frame.height as u32, rgb_buf).unwrap();
394 img_buf.save(path).unwrap();
395 }
396}