pfv_rs/
enc.rs

1use std::io::{Write, Cursor};
2
3use bitstream_io::{BitWriter, BitWrite};
4use byteorder::{WriteBytesExt, LittleEndian};
5
6use crate::common::{EncodedIFrame, PFV_MAGIC, PFV_VERSION, EncodedPFrame};
7use crate::frame::VideoFrame;
8use crate::dct::{Q_TABLE_INTER, Q_TABLE_INTRA};
9use crate::plane::VideoPlane;
10use crate::rle::{rle_encode, rle_create_huffman, update_table};
11
12pub struct Encoder<W: Write> {
13    width: usize,
14    height: usize,
15    framerate: u32,
16    prev_frame: VideoFrame,
17    px_err: f32,
18    qtable_inter_l: [i32;64],
19    qtable_inter_c: [i32;64],
20    qtable_intra_l: [i32;64],
21    qtable_intra_c: [i32;64],
22    writer: W,
23    finished: bool,
24    #[cfg(feature = "multithreading")]
25    threadpool: rayon::ThreadPool
26}
27
28impl<W: Write> Drop for Encoder<W> {
29    fn drop(&mut self) {
30        if !self.finished {
31            self.finish().unwrap();
32        }   
33    }
34}
35
36impl<W: Write> Encoder<W> {
37    pub fn new(writer: W, width: usize, height: usize, framerate: u32, quality: i32, #[cfg(feature = "multithreading")] num_threads: usize) -> Result<Encoder<W>, std::io::Error> {
38        assert!(quality >= 0 && quality <= 10);
39
40        let qscale = quality as f32 * 0.25;
41        let px_err = quality as f32 * 1.5;
42
43        #[cfg(feature = "multithreading")]
44        let mut enc = {
45            Encoder { width: width, height: height, framerate: framerate,
46                prev_frame: VideoFrame::new_padded(width, height),
47                px_err: px_err,
48                qtable_inter_l: Q_TABLE_INTER.map(|x| (x as f32 * qscale * 0.5).max(1.0) as i32),
49                qtable_inter_c: Q_TABLE_INTER.map(|x| (x as f32 * qscale).max(1.0) as i32),
50                qtable_intra_l: Q_TABLE_INTRA.map(|x| (x as f32 * qscale * 0.5).max(1.0) as i32),
51                qtable_intra_c: Q_TABLE_INTRA.map(|x| (x as f32 * qscale).max(1.0) as i32),
52                writer: writer,
53                finished: false,
54                threadpool: rayon::ThreadPoolBuilder::new().num_threads(num_threads).build().unwrap() }
55        };
56
57        #[cfg(not(feature = "multithreading"))]
58        let mut enc = {
59            Encoder { width: width, height: height, framerate: framerate,
60                prev_frame: VideoFrame::new_padded(width, height),
61                px_err: px_err,
62                qtable_inter_l: Q_TABLE_INTER.map(|x| (x as f32 * qscale * 0.5).max(1.0) as i32),
63                qtable_inter_c: Q_TABLE_INTER.map(|x| (x as f32 * qscale).max(1.0) as i32),
64                qtable_intra_l: Q_TABLE_INTRA.map(|x| (x as f32 * qscale * 0.5).max(1.0) as i32),
65                qtable_intra_c: Q_TABLE_INTRA.map(|x| (x as f32 * qscale).max(1.0) as i32),
66                writer: writer,
67                finished: false, }
68        };
69
70        enc.write_header()?;
71
72        Ok(enc)
73    }
74
75    pub fn encode_iframe(self: &mut Encoder<W>, frame: &VideoFrame) -> Result<(), std::io::Error> {
76        assert!(frame.width == self.width && frame.height == self.height);
77        assert!(frame.plane_y.width == frame.width && frame.plane_y.height == frame.height);
78        assert!(frame.plane_u.width == frame.width / 2 && frame.plane_u.height == frame.height / 2);
79        assert!(frame.plane_v.width == frame.width / 2 && frame.plane_v.height == frame.height / 2);
80        assert!(!self.finished);
81
82        #[cfg(feature = "multithreading")]
83        {
84            let enc_y = frame.plane_y.encode_plane(&self.qtable_intra_l, 0, &self.threadpool);
85            let dec_y = VideoPlane::decode_plane(&enc_y, &self.qtable_intra_l, &self.threadpool);
86
87            let enc_u = frame.plane_u.encode_plane(&self.qtable_intra_c, 128, &self.threadpool);
88            let dec_u = VideoPlane::decode_plane(&enc_u, &self.qtable_intra_c, &self.threadpool);
89
90            let enc_v = frame.plane_v.encode_plane(&self.qtable_intra_c, 128, &self.threadpool);
91            let dec_v = VideoPlane::decode_plane(&enc_v, &self.qtable_intra_c, &self.threadpool);
92
93            let enc_frame = EncodedIFrame { y: enc_y, u: enc_u, v: enc_v };
94
95            self.prev_frame.plane_y.blit(&dec_y, 0, 0, 0, 0, dec_y.width, dec_y.height);
96            self.prev_frame.plane_u.blit(&dec_u, 0, 0, 0, 0, dec_u.width, dec_u.height);
97            self.prev_frame.plane_v.blit(&dec_v, 0, 0, 0, 0, dec_v.width, dec_v.height);
98
99            Encoder::<W>::write_iframe_packet(&enc_frame, &mut self.writer)?;
100        }
101
102        #[cfg(not(feature = "multithreading"))]
103        {
104            let enc_y = frame.plane_y.encode_plane(&self.qtable_intra_l, 0);
105            let dec_y = VideoPlane::decode_plane(&enc_y, &self.qtable_intra_l);
106
107            let enc_u = frame.plane_u.encode_plane(&self.qtable_intra_c, 128);
108            let dec_u = VideoPlane::decode_plane(&enc_u, &self.qtable_intra_c);
109
110            let enc_v = frame.plane_v.encode_plane(&self.qtable_intra_c, 128);
111            let dec_v = VideoPlane::decode_plane(&enc_v, &self.qtable_intra_c);
112
113            let enc_frame = EncodedIFrame { y: enc_y, u: enc_u, v: enc_v };
114
115            self.prev_frame.plane_y.blit(&dec_y, 0, 0, 0, 0, dec_y.width, dec_y.height);
116            self.prev_frame.plane_u.blit(&dec_u, 0, 0, 0, 0, dec_u.width, dec_u.height);
117            self.prev_frame.plane_v.blit(&dec_v, 0, 0, 0, 0, dec_v.width, dec_v.height);
118
119            Encoder::<W>::write_iframe_packet(&enc_frame, &mut self.writer)?;
120        }
121
122        Ok(())
123    }
124
125    pub fn encode_pframe(self: &mut Encoder<W>, frame: &VideoFrame) -> Result<(), std::io::Error> {
126        assert!(frame.width == self.width && frame.height == self.height);
127        assert!(frame.plane_y.width == frame.width && frame.plane_y.height == frame.height);
128        assert!(frame.plane_u.width == frame.width / 2 && frame.plane_u.height == frame.height / 2);
129        assert!(frame.plane_v.width == frame.width / 2 && frame.plane_v.height == frame.height / 2);
130        assert!(!self.finished);
131
132        #[cfg(feature = "multithreading")]
133        {
134            let enc_y = frame.plane_y.encode_plane_delta(&self.prev_frame.plane_y, &self.qtable_inter_l, self.px_err, 0, &self.threadpool);
135            let dec_y = VideoPlane::decode_plane_delta(&enc_y, &self.prev_frame.plane_y, &self.qtable_inter_l, &self.threadpool);
136
137            let enc_u = frame.plane_u.encode_plane_delta(&self.prev_frame.plane_u, &self.qtable_inter_c, self.px_err, 128, &self.threadpool);
138            let dec_u = VideoPlane::decode_plane_delta(&enc_u, &self.prev_frame.plane_u, &self.qtable_inter_c, &self.threadpool);
139
140            let enc_v = frame.plane_v.encode_plane_delta(&self.prev_frame.plane_v, &self.qtable_inter_c, self.px_err, 128, &self.threadpool);
141            let dec_v = VideoPlane::decode_plane_delta(&enc_v, &self.prev_frame.plane_v, &self.qtable_inter_c, &self.threadpool);
142
143            let enc_frame = EncodedPFrame { y: enc_y, u: enc_u, v: enc_v };
144
145            self.prev_frame.plane_y.blit(&dec_y, 0, 0, 0, 0, dec_y.width, dec_y.height);
146            self.prev_frame.plane_u.blit(&dec_u, 0, 0, 0, 0, dec_u.width, dec_u.height);
147            self.prev_frame.plane_v.blit(&dec_v, 0, 0, 0, 0, dec_v.width, dec_v.height);
148            
149            Encoder::<W>::write_pframe_packet(&enc_frame, &mut self.writer)?;
150        }
151
152        #[cfg(not(feature = "multithreading"))]
153        {
154            let enc_y = frame.plane_y.encode_plane_delta(&self.prev_frame.plane_y, &self.qtable_inter_l, self.px_err, 0);
155            let dec_y = VideoPlane::decode_plane_delta(&enc_y, &self.prev_frame.plane_y, &self.qtable_inter_l);
156
157            let enc_u = frame.plane_u.encode_plane_delta(&self.prev_frame.plane_u, &self.qtable_inter_c, self.px_err, 128);
158            let dec_u = VideoPlane::decode_plane_delta(&enc_u, &self.prev_frame.plane_u, &self.qtable_inter_c);
159
160            let enc_v = frame.plane_v.encode_plane_delta(&self.prev_frame.plane_v, &self.qtable_inter_c, self.px_err, 128);
161            let dec_v = VideoPlane::decode_plane_delta(&enc_v, &self.prev_frame.plane_v, &self.qtable_inter_c);
162
163            let enc_frame = EncodedPFrame { y: enc_y, u: enc_u, v: enc_v };
164
165            self.prev_frame.plane_y.blit(&dec_y, 0, 0, 0, 0, dec_y.width, dec_y.height);
166            self.prev_frame.plane_u.blit(&dec_u, 0, 0, 0, 0, dec_u.width, dec_u.height);
167            self.prev_frame.plane_v.blit(&dec_v, 0, 0, 0, 0, dec_v.width, dec_v.height);
168
169            Encoder::<W>::write_pframe_packet(&enc_frame, &mut self.writer)?;
170        }
171
172        Ok(())
173    }
174
175    pub fn encode_dropframe(self: &mut Encoder<W>) -> Result<(), std::io::Error> {
176        assert!(!self.finished);
177
178        Encoder::<W>::write_drop_packet(&mut self.writer)?;
179        Ok(())
180    }
181
182    pub fn finish(self: &mut Encoder<W>) -> Result<(), std::io::Error> {
183        assert!(!self.finished);
184
185        self.finished = true;
186        Encoder::write_eof(&mut self.writer)?;
187        Ok(())
188    }
189
190    fn write_header(self: &mut Encoder<W>) -> Result<(), std::io::Error> {
191        // write PGV header
192        self.writer.write_all(PFV_MAGIC)?;
193        self.writer.write_u32::<LittleEndian>(PFV_VERSION)?;
194
195        self.writer.write_u16::<LittleEndian>(self.width as u16)?;
196        self.writer.write_u16::<LittleEndian>(self.height as u16)?;
197        self.writer.write_u16::<LittleEndian>(self.framerate as u16)?;
198
199        // write q-tables
200        self.writer.write_u16::<LittleEndian>(4)?;
201
202        for v in self.qtable_intra_l {
203            self.writer.write_u16::<LittleEndian>(v as u16)?;
204        }
205
206        for v in self.qtable_intra_c {
207            self.writer.write_u16::<LittleEndian>(v as u16)?;
208        }
209
210        for v in self.qtable_inter_l {
211            self.writer.write_u16::<LittleEndian>(v as u16)?;
212        }
213
214        for v in self.qtable_inter_c {
215            self.writer.write_u16::<LittleEndian>(v as u16)?;
216        }
217
218        Ok(())
219    }
220
221    fn write_eof(writer: &mut W) -> Result<(), std::io::Error> {
222        // write packet header
223        writer.write_u8(0)?; // packet type = EOF
224        writer.write_u32::<LittleEndian>(0)?;
225
226        Ok(())
227    }
228
229    fn write_drop_packet(writer: &mut W) -> Result<(), std::io::Error> {
230        // write packet header
231        writer.write_u8(1)?; // packet type = iframe
232        writer.write_u32::<LittleEndian>(0)?;
233
234        Ok(())
235    }
236
237    fn write_iframe_packet(f: &EncodedIFrame, writer: &mut W) -> Result<(), std::io::Error> {
238        // serialize packet data
239        let mut packet_data = Cursor::new(Vec::new());
240        let mut bitwriter = BitWriter::endian(&mut packet_data, bitstream_io::LittleEndian);
241
242        // gather RLE-encoded block coefficients for each plane
243        let mut block_coeff = Vec::new();
244        let mut symbol_table = [0;16];
245
246        for b in &f.y.blocks {
247            let mut coeff = Vec::new();
248            coeff.extend_from_slice(&b.subblocks[0].m);
249            coeff.extend_from_slice(&b.subblocks[1].m);
250            coeff.extend_from_slice(&b.subblocks[2].m);
251            coeff.extend_from_slice(&b.subblocks[3].m);
252            let mut rle_sequence = Vec::new();
253            rle_encode(&mut rle_sequence, &coeff);
254            update_table(&mut symbol_table, &rle_sequence);
255
256            block_coeff.push(rle_sequence);
257        }
258
259        for b in &f.u.blocks {
260            let mut coeff = Vec::new();
261            coeff.extend_from_slice(&b.subblocks[0].m);
262            coeff.extend_from_slice(&b.subblocks[1].m);
263            coeff.extend_from_slice(&b.subblocks[2].m);
264            coeff.extend_from_slice(&b.subblocks[3].m);
265            let mut rle_sequence = Vec::new();
266            rle_encode(&mut rle_sequence, &coeff);
267            update_table(&mut symbol_table, &rle_sequence);
268
269            block_coeff.push(rle_sequence);
270        }
271
272        for b in &f.v.blocks {
273            let mut coeff = Vec::new();
274            coeff.extend_from_slice(&b.subblocks[0].m);
275            coeff.extend_from_slice(&b.subblocks[1].m);
276            coeff.extend_from_slice(&b.subblocks[2].m);
277            coeff.extend_from_slice(&b.subblocks[3].m);
278            let mut rle_sequence = Vec::new();
279            rle_encode(&mut rle_sequence, &coeff);
280            update_table(&mut symbol_table, &rle_sequence);
281
282            block_coeff.push(rle_sequence);
283        }
284
285        // create huffman tree for encoding RLE results
286        let tree = rle_create_huffman(&symbol_table);
287        let tree_table = tree.get_table();
288
289        // write symbol frequency table
290        for i in 0..16 {
291            bitwriter.write(8, tree_table[i] as u8)?;
292        }
293
294        // we currently create four qtables: two for i-frames (0, 1) and two for p-frames (2, 3)
295        // note: (one qtable index per plane)
296        bitwriter.write(8, 0_u8)?;
297        bitwriter.write(8, 1_u8)?;
298        bitwriter.write(8, 1_u8)?;
299
300        // serialize blocks to bitstream
301        for block in &block_coeff {
302            for sq in block {
303                let num_zeroes = tree.get_code(sq.num_zeroes);
304                let num_bits = tree.get_code(sq.coeff_size);
305
306                debug_assert!(num_zeroes.len > 0 && num_bits.len > 0);
307
308                bitwriter.write(num_zeroes.len, num_zeroes.val)?;
309                bitwriter.write(num_bits.len, num_bits.val)?;
310
311                if sq.coeff_size > 0 {
312                    bitwriter.write_signed(sq.coeff_size as u32, sq.coeff)?;
313                }
314            }
315        }
316
317        // flush any partial bytes
318        bitwriter.byte_align()?;
319
320        // retrieve packet payload bytes
321        let packet_data = packet_data.into_inner();
322
323        // write packet header + data
324
325        writer.write_u8(1)?; // packet type = iframe
326        writer.write_u32::<LittleEndian>(packet_data.len() as u32)?;
327        writer.write_all(&packet_data)?;
328
329        Ok(())
330    }
331
332    fn write_pframe_packet(f: &EncodedPFrame, writer: &mut W) -> Result<(), std::io::Error> {
333        // serialize packet data
334        let mut packet_data = Cursor::new(Vec::new());
335        let mut bitwriter = BitWriter::endian(&mut packet_data, bitstream_io::LittleEndian);
336
337        // gather RLE-encoded block coefficients for each plane
338        let mut block_coeff = Vec::new();
339        let mut symbol_table = [0;16];
340
341        for b in &f.y.blocks {
342            match b.subblocks {
343                Some(subblocks) => {
344                    let mut coeff = Vec::new();
345                    coeff.extend_from_slice(&subblocks[0].m);
346                    coeff.extend_from_slice(&subblocks[1].m);
347                    coeff.extend_from_slice(&subblocks[2].m);
348                    coeff.extend_from_slice(&subblocks[3].m);
349                    let mut rle_sequence = Vec::new();
350                    rle_encode(&mut rle_sequence, &coeff);
351                    update_table(&mut symbol_table, &rle_sequence);
352
353                    block_coeff.push(rle_sequence);
354                }
355                None => {
356                }
357            }
358        }
359
360        for b in &f.u.blocks {
361            match b.subblocks {
362                Some(subblocks) => {
363                    let mut coeff = Vec::new();
364                    coeff.extend_from_slice(&subblocks[0].m);
365                    coeff.extend_from_slice(&subblocks[1].m);
366                    coeff.extend_from_slice(&subblocks[2].m);
367                    coeff.extend_from_slice(&subblocks[3].m);
368                    let mut rle_sequence = Vec::new();
369                    rle_encode(&mut rle_sequence, &coeff);
370                    update_table(&mut symbol_table, &rle_sequence);
371
372                    block_coeff.push(rle_sequence);
373                }
374                None => {
375                }
376            }
377        }
378
379        for b in &f.v.blocks {
380            match b.subblocks {
381                Some(subblocks) => {
382                    let mut coeff = Vec::new();
383                    coeff.extend_from_slice(&subblocks[0].m);
384                    coeff.extend_from_slice(&subblocks[1].m);
385                    coeff.extend_from_slice(&subblocks[2].m);
386                    coeff.extend_from_slice(&subblocks[3].m);
387                    let mut rle_sequence = Vec::new();
388                    rle_encode(&mut rle_sequence, &coeff);
389                    update_table(&mut symbol_table, &rle_sequence);
390
391                    block_coeff.push(rle_sequence);
392                }
393                None => {
394                }
395            }
396        }
397
398        // create huffman tree for encoding RLE results
399        let tree = rle_create_huffman(&symbol_table);
400        let tree_table = tree.get_table();
401
402        // write symbol frequency table
403        for i in 0..16 {
404            bitwriter.write(8, tree_table[i] as u8)?;
405        }
406
407        // we currently create four qtables: two for i-frames (0, 1) and two for p-frames (2, 3)
408        // note: (one qtable index per plane)
409        bitwriter.write(8, 2_u8)?;
410        bitwriter.write(8, 3_u8)?;
411        bitwriter.write(8, 3_u8)?;
412
413        // write block headers
414        for b in &f.y.blocks {
415            let has_mvec = b.motion_x != 0 || b.motion_y != 0;
416
417            bitwriter.write_bit(has_mvec)?;
418            bitwriter.write_bit(b.subblocks.is_some())?;
419
420            if has_mvec {
421                bitwriter.write_signed(7, b.motion_x as i32)?;
422                bitwriter.write_signed(7, b.motion_y as i32)?;
423            }
424        }
425
426        for b in &f.u.blocks {
427            let has_mvec = b.motion_x != 0 || b.motion_y != 0;
428
429            bitwriter.write_bit(has_mvec)?;
430            bitwriter.write_bit(b.subblocks.is_some())?;
431
432            if has_mvec {
433                bitwriter.write_signed(7, b.motion_x as i32)?;
434                bitwriter.write_signed(7, b.motion_y as i32)?;
435            }
436        }
437
438        for b in &f.v.blocks {
439            let has_mvec = b.motion_x != 0 || b.motion_y != 0;
440
441            bitwriter.write_bit(has_mvec)?;
442            bitwriter.write_bit(b.subblocks.is_some())?;
443
444            if has_mvec {
445                bitwriter.write_signed(7, b.motion_x as i32)?;
446                bitwriter.write_signed(7, b.motion_y as i32)?;
447
448                assert!(b.motion_x >= -16 && b.motion_x <= 16);
449                assert!(b.motion_y >= -16 && b.motion_y <= 16);
450            }
451        }
452
453        // serialize block data to bitstream
454        for block in &block_coeff {
455            for sq in block {
456                let num_zeroes = tree.get_code(sq.num_zeroes);
457                let num_bits = tree.get_code(sq.coeff_size);
458
459                bitwriter.write(num_zeroes.len, num_zeroes.val)?;
460                bitwriter.write(num_bits.len, num_bits.val)?;
461
462                if sq.coeff_size > 0 {
463                    bitwriter.write_signed(sq.coeff_size as u32, sq.coeff)?;
464                }
465            }
466        }
467
468        // flush any partial bytes
469        bitwriter.byte_align()?;
470
471        // retrieve packet payload bytes
472        let packet_data = packet_data.into_inner();
473
474        // write packet header + data
475
476        writer.write_u8(2)?; // packet type = pframe
477        writer.write_u32::<LittleEndian>(packet_data.len() as u32)?;
478        writer.write_all(&packet_data)?;
479
480        Ok(())
481    }
482}
pfv_rs/enc.rs

pfv_rs/
enc.rs
