pfv_rs/
dec.rs

1use std::{io::{Read, Seek, Cursor}, slice::{ChunksExact, Iter}};
2
3use bitstream_io::{BitReader, BitRead};
4use byteorder::{ReadBytesExt, LittleEndian};
5
6use crate::{common::{PFV_MAGIC, PFV_VERSION, EncodedMacroBlock, EncodedIPlane, DeltaEncodedMacroBlock, EncodedPPlane}, huffman::{HuffmanTree, HuffmanError}, frame::VideoFrame, plane::VideoPlane, dct::{DctQuantizedMatrix8x8}};
7
8#[derive(Debug, Clone, Copy)]
9struct DeltaBlockHeader {
10    mvec_x: i8,
11    mvec_y: i8,
12    has_coeff: bool,
13}
14
15pub struct Decoder<TReader: Read + Seek> {
16    reader: TReader,
17    width: usize,
18    height: usize,
19    framerate: u32,
20    qtables: Vec<[i32;64]>,
21    framebuffer: VideoFrame,
22    retframe: VideoFrame,
23    delta_accum: f64,
24    eof: bool,
25    reset_pos: u64,
26    #[cfg(feature = "multithreading")]
27    threadpool: rayon::ThreadPool
28}
29
30#[derive(Debug)]
31pub enum DecodeError {
32    FormatError,
33    VersionError,
34    IOError(std::io::Error)
35}
36
37impl<TReader: Read + Seek> Decoder<TReader> {
38    pub fn new(mut reader: TReader, #[cfg(feature = "multithreading")] num_threads: usize) -> Result<Decoder<TReader>, DecodeError> {
39        // read header
40        let mut magic = [0;8];
41        match reader.read_exact(&mut magic) {
42            Ok(_) => {}
43            Err(e) => {
44                return Err(DecodeError::IOError(e));
45            }
46        };
47
48        let magic_match = magic.iter().zip(PFV_MAGIC.iter()).all(|(a, b)| *a == *b);
49
50        if !magic_match {
51            return Err(DecodeError::FormatError);
52        }
53
54        // read version
55        match reader.read_u32::<LittleEndian>() {
56            Ok(ver) => {
57                if ver != PFV_VERSION {
58                    return Err(DecodeError::VersionError);
59                }
60
61                ver
62            }
63            Err(e) => {
64                return Err(DecodeError::IOError(e));
65            }
66        };
67
68        let width = match reader.read_u16::<LittleEndian>() {
69            Ok(v) => v,
70            Err(e) => {
71                return Err(DecodeError::IOError(e));
72            }
73        };
74
75        let height = match reader.read_u16::<LittleEndian>() {
76            Ok(v) => v,
77            Err(e) => {
78                return Err(DecodeError::IOError(e));
79            }
80        };
81
82        let framerate = match reader.read_u16::<LittleEndian>() {
83            Ok(v) => v,
84            Err(e) => {
85                return Err(DecodeError::IOError(e));
86            }
87        };
88
89        let num_qtable = match reader.read_u16::<LittleEndian>() {
90            Ok(v) => v,
91            Err(e) => {
92                return Err(DecodeError::IOError(e));
93            }
94        };
95
96        let mut qtables = Vec::new();
97
98        for _ in 0..num_qtable {
99            let mut qtable = [0;64];
100
101            for i in 0..64 {
102                qtable[i] = match reader.read_u16::<LittleEndian>() {
103                    Ok(v) => v as i32,
104                    Err(e) => {
105                        return Err(DecodeError::IOError(e));
106                    }
107                };
108            }
109
110            qtables.push(qtable);
111        }
112
113        let reset_pos = match reader.stream_position() {
114            Ok(v) => v,
115            Err(e) => {
116                return Err(DecodeError::IOError(e));
117            }
118        };
119
120        #[cfg(feature = "multithreading")]
121        {
122            Ok(Decoder { reader: reader, width: width as usize, height: height as usize, framerate: framerate as u32,
123                qtables: qtables, framebuffer: VideoFrame::new_padded(width as usize, height as usize),
124                retframe: VideoFrame::new(width as usize, height as usize), delta_accum: 0.0, eof: false, reset_pos: reset_pos,
125                threadpool: rayon::ThreadPoolBuilder::new().num_threads(num_threads).build().unwrap() })
126        }
127
128        #[cfg(not(feature = "multithreading"))]
129        {
130            Ok(Decoder { reader: reader, width: width as usize, height: height as usize, framerate: framerate as u32,
131                qtables: qtables, framebuffer: VideoFrame::new_padded(width as usize, height as usize),
132                retframe: VideoFrame::new(width as usize, height as usize), delta_accum: 0.0, eof: false, reset_pos: reset_pos, })
133        }
134    }
135
136    pub fn width(self: &Decoder<TReader>) -> usize {
137        return self.width;
138    }
139
140    pub fn height(self: &Decoder<TReader>) -> usize {
141        return self.height;
142    }
143
144    pub fn framerate(self: &Decoder<TReader>) -> u32 {
145        return self.framerate;
146    }
147
148    pub fn reset(self: &mut Decoder<TReader>) -> Result<(), std::io::Error> {
149        self.eof = false;
150        self.reader.seek(std::io::SeekFrom::Start(self.reset_pos))?;
151        Ok(())
152    }
153
154    pub fn advance_delta<FV>(self: &mut Decoder<TReader>, delta: f64, onvideo: &mut FV) -> Result<bool, std::io::Error>  where
155        FV: FnMut(&VideoFrame) {
156        self.delta_accum += delta;
157        let delta_per_frame = 1.0 / self.framerate as f64;
158
159        while self.delta_accum >= delta_per_frame {
160            if self.advance_frame(onvideo)? == false {
161                return Ok(false);
162            }
163            self.delta_accum -= delta_per_frame;
164        }
165
166        Ok(true)
167    }
168
169    pub fn advance_frame<FV>(self: &mut Decoder<TReader>, onvideo: &mut FV) -> Result<bool, std::io::Error> where
170        FV: FnMut(&VideoFrame) {
171        if self.eof {
172            return Ok(false);
173        }
174
175        loop {
176            // read next packet header
177            // if we hit EOF, return false
178
179            let packet_type = self.reader.read_u8()?;
180            let packet_len = self.reader.read_u32::<LittleEndian>()?;
181
182            match packet_type {
183                0 => {
184                    // EOF marker
185                    self.eof = true;
186                    return Ok(false);
187                }
188                1 => {
189                    // iframe. if payload length is zero, this is a drop frame (do nothing)
190                    if packet_len > 0 {
191                        let mut data = vec![0;packet_len as usize];
192                        self.reader.read_exact(&mut data)?;
193                        self.decode_iframe(&data)?;
194
195                        self.retframe.plane_y.blit(&self.framebuffer.plane_y, 0, 0, 0, 0, self.retframe.plane_y.width, self.retframe.plane_y.height);
196                        self.retframe.plane_u.blit(&self.framebuffer.plane_u, 0, 0, 0, 0, self.retframe.plane_u.width, self.retframe.plane_u.height);
197                        self.retframe.plane_v.blit(&self.framebuffer.plane_v, 0, 0, 0, 0, self.retframe.plane_v.width, self.retframe.plane_v.height);
198
199                        onvideo(&self.retframe);
200                    }
201                    break;
202                }
203                2 => {
204                    // pframe
205                    let mut data = vec![0;packet_len as usize];
206                    self.reader.read_exact(&mut data)?;
207                    self.decode_pframe(&data)?;
208
209                    self.retframe.plane_y.blit(&self.framebuffer.plane_y, 0, 0, 0, 0, self.retframe.plane_y.width, self.retframe.plane_y.height);
210                    self.retframe.plane_u.blit(&self.framebuffer.plane_u, 0, 0, 0, 0, self.retframe.plane_u.width, self.retframe.plane_u.height);
211                    self.retframe.plane_v.blit(&self.framebuffer.plane_v, 0, 0, 0, 0, self.retframe.plane_v.width, self.retframe.plane_v.height);
212
213                    onvideo(&self.retframe);
214                    break;
215                }
216                _ => {
217                    // unrecognized packet type, just skip over packet payload
218                    self.reader.seek(std::io::SeekFrom::Current(packet_len as i64))?;
219                }
220            }
221        }
222
223        Ok(true)
224    }
225
226    fn decode_iframe(self: &mut Decoder<TReader>, payload: &[u8]) -> Result<(), std::io::Error> {
227        let reader = Cursor::new(payload);
228        let mut bitreader = BitReader::endian(reader, bitstream_io::LittleEndian);
229
230        let bitstream_length = bitreader.seek_bits(std::io::SeekFrom::End(0))?;
231        bitreader.seek_bits(std::io::SeekFrom::Start(0))?;
232
233        // read symbol frequency table
234        let mut table = [0;16];
235
236        for i in 0..16 {
237            table[i] = bitreader.read::<u8>(8).unwrap();
238        }
239
240        // construct huffman tree
241        let tree = HuffmanTree::from_table(&table);
242
243        // fetch qtables
244        let qtable_y = &self.qtables[bitreader.read::<u8>(8).unwrap() as usize];
245        let qtable_u = &self.qtables[bitreader.read::<u8>(8).unwrap() as usize];
246        let qtable_v = &self.qtables[bitreader.read::<u8>(8).unwrap() as usize];
247
248        // decode RLE coefficients
249        let blocks_wide = self.framebuffer.plane_y.width / 16;
250        let blocks_high = self.framebuffer.plane_y.height / 16;
251
252        let chroma_blocks_wide = self.framebuffer.plane_u.width / 16;
253        let chroma_blocks_high = self.framebuffer.plane_u.height / 16;
254
255        let total_blocks = (blocks_wide * blocks_high) + (chroma_blocks_wide * chroma_blocks_high * 2);
256        let total_subblocks = total_blocks * 4;
257
258        let mut coefficients = vec![0;total_subblocks * 64 as usize];
259
260        let mut out_idx = 0;
261        while out_idx < coefficients.len() {
262            let num_zeroes = match tree.read(&mut bitreader, bitstream_length) {
263                Ok(v) => v,
264                Err(e) => match e {
265                    HuffmanError::DecodeError => unreachable!(),
266                    HuffmanError::IOError(e2) => {
267                        return Err(e2);
268                    },
269                }
270            } as usize;
271
272            out_idx += num_zeroes;
273
274            let num_bits = match tree.read(&mut bitreader, bitstream_length) {
275                Ok(v) => v,
276                Err(e) => match e {
277                    HuffmanError::DecodeError => unreachable!(),
278                    HuffmanError::IOError(e2) => {
279                        return Err(e2);
280                    },
281                }
282            };
283
284            // if num_bits is 0, then this is only a run of 0s with no value
285            if num_bits > 0 {
286                let coeff = match bitreader.read_signed::<i16>(num_bits as u32) {
287                    Ok(v) => v,
288                    Err(e) => {
289                        return Err(e);
290                    }
291                };
292                coefficients[out_idx] = coeff;
293
294                out_idx += 1;
295            }
296        }
297
298        let mut subblocks = coefficients.chunks_exact(64);
299
300        // deserialize each plane
301        #[cfg(feature = "multithreading")]
302        {
303            Decoder::<TReader>::deserialize_plane(self.framebuffer.plane_y.width, self.framebuffer.plane_y.height,
304                &mut subblocks, qtable_y, &mut self.framebuffer.plane_y, &self.threadpool);
305                
306            Decoder::<TReader>::deserialize_plane(self.framebuffer.plane_u.width, self.framebuffer.plane_u.height,
307                &mut subblocks, qtable_u, &mut self.framebuffer.plane_u, &self.threadpool);
308                
309            Decoder::<TReader>::deserialize_plane(self.framebuffer.plane_v.width, self.framebuffer.plane_v.height,
310                &mut subblocks, qtable_v, &mut self.framebuffer.plane_v, &self.threadpool);
311        }
312
313        #[cfg(not(feature = "multithreading"))]
314        {
315            Decoder::<TReader>::deserialize_plane(self.framebuffer.plane_y.width, self.framebuffer.plane_y.height,
316                &mut subblocks, qtable_y, &mut self.framebuffer.plane_y);
317                
318            Decoder::<TReader>::deserialize_plane(self.framebuffer.plane_u.width, self.framebuffer.plane_u.height,
319                &mut subblocks, qtable_u, &mut self.framebuffer.plane_u);
320                
321            Decoder::<TReader>::deserialize_plane(self.framebuffer.plane_v.width, self.framebuffer.plane_v.height,
322                &mut subblocks, qtable_v, &mut self.framebuffer.plane_v);
323        }
324
325        Ok(())
326    }
327
328    fn decode_pframe(self: &mut Decoder<TReader>, payload: &[u8]) -> Result<(), std::io::Error> {
329        let reader = Cursor::new(payload);
330        let mut bitreader = BitReader::endian(reader, bitstream_io::LittleEndian);
331
332        let bitstream_length = bitreader.seek_bits(std::io::SeekFrom::End(0))?;
333        bitreader.seek_bits(std::io::SeekFrom::Start(0))?;
334
335        // read symbol frequency table
336        let mut table = [0;16];
337
338        for i in 0..16 {
339            table[i] = bitreader.read::<u8>(8).unwrap();
340        }
341
342        // construct huffman tree
343        let tree = HuffmanTree::from_table(&table);
344
345        // fetch qtables
346        let qtable_y = &self.qtables[bitreader.read::<u8>(8)? as usize];
347        let qtable_u = &self.qtables[bitreader.read::<u8>(8)? as usize];
348        let qtable_v = &self.qtables[bitreader.read::<u8>(8)? as usize];
349
350        // read block headers
351        let blocks_wide = self.framebuffer.plane_y.width / 16;
352        let blocks_high = self.framebuffer.plane_y.height / 16;
353
354        let chroma_blocks_wide = self.framebuffer.plane_u.width / 16;
355        let chroma_blocks_high = self.framebuffer.plane_u.height / 16;
356
357        let total_blocks = (blocks_wide * blocks_high) + (chroma_blocks_wide * chroma_blocks_high * 2);
358
359        let mut block_headers = Vec::with_capacity(total_blocks);
360
361        for _ in 0..total_blocks {
362            let mut header = DeltaBlockHeader { mvec_x: 0, mvec_y: 0, has_coeff: false };
363            let has_mvec = bitreader.read_bit()?;
364            header.has_coeff = bitreader.read_bit()?;
365
366            if has_mvec {
367                header.mvec_x = bitreader.read_signed(7)?;
368                header.mvec_y = bitreader.read_signed(7)?;
369            }
370
371            block_headers.push(header);
372        }
373
374        // decode block coefficients
375
376        let mut coefficients = vec![0;total_blocks * 256];
377
378        for (idx, header) in block_headers.iter().enumerate() {
379            let mut block_coeff = [0;256];
380            let block_offset = idx * 256;
381            if header.has_coeff {
382                // read 256 coefficients from bit stream
383                let mut out_idx = 0;
384                while out_idx < 256 {
385                    let num_zeroes = match tree.read(&mut bitreader, bitstream_length) {
386                        Ok(v) => v,
387                        Err(e) => match e {
388                            HuffmanError::DecodeError => unreachable!(),
389                            HuffmanError::IOError(e2) => {
390                                return Err(e2);
391                            },
392                        }
393                    } as usize;
394
395                    out_idx += num_zeroes;
396
397                    let num_bits = match tree.read(&mut bitreader, bitstream_length) {
398                        Ok(v) => v,
399                        Err(e) => match e {
400                            HuffmanError::DecodeError => unreachable!(),
401                            HuffmanError::IOError(e2) => {
402                                return Err(e2);
403                            },
404                        }
405                    };
406
407                    // if num_bits is 0, then this is only a run of 0s with no value
408                    if num_bits > 0 {
409                        let coeff = bitreader.read_signed::<i16>(num_bits as u32)?;
410                        block_coeff[out_idx] = coeff;
411
412                        out_idx += 1;
413                    }
414                }
415            }
416            coefficients[block_offset..block_offset+256].copy_from_slice(&block_coeff);
417        }
418
419        let mut subblocks = coefficients.chunks_exact(64);
420        let mut headers = block_headers.iter();
421
422        // deserialize each plane
423        #[cfg(feature = "multithreading")]
424        {
425            Decoder::<TReader>::deserialize_plane_delta(self.framebuffer.plane_y.width, self.framebuffer.plane_y.height,
426                &mut headers, &mut subblocks, qtable_y, &mut self.framebuffer.plane_y, &self.threadpool);
427                
428            Decoder::<TReader>::deserialize_plane_delta(self.framebuffer.plane_u.width, self.framebuffer.plane_u.height,
429                &mut headers, &mut subblocks, qtable_u, &mut self.framebuffer.plane_u, &self.threadpool);
430                
431            Decoder::<TReader>::deserialize_plane_delta(self.framebuffer.plane_v.width, self.framebuffer.plane_v.height,
432                &mut headers, &mut subblocks, qtable_v, &mut self.framebuffer.plane_v, &self.threadpool);
433        }
434
435        #[cfg(not(feature = "multithreading"))]
436        {
437            Decoder::<TReader>::deserialize_plane_delta(self.framebuffer.plane_y.width, self.framebuffer.plane_y.height,
438                &mut headers, &mut subblocks, qtable_y, &mut self.framebuffer.plane_y);
439                
440            Decoder::<TReader>::deserialize_plane_delta(self.framebuffer.plane_u.width, self.framebuffer.plane_u.height,
441                &mut headers, &mut subblocks, qtable_u, &mut self.framebuffer.plane_u);
442                
443            Decoder::<TReader>::deserialize_plane_delta(self.framebuffer.plane_v.width, self.framebuffer.plane_v.height,
444                &mut headers, &mut subblocks, qtable_v, &mut self.framebuffer.plane_v);
445        }
446
447        Ok(())
448    }
449
450    fn deserialize_plane(width: usize, height: usize, subblocks: &mut ChunksExact<i16>, q_table: &[i32;64], target: &mut VideoPlane, #[cfg(feature = "multithreading")] tp: &rayon::ThreadPool) {
451        let blocks_wide = width / 16;
452        let blocks_high = height / 16;
453        let total_blocks = blocks_wide * blocks_high;
454
455        let mut enc_plane = EncodedIPlane { blocks_wide: blocks_wide, blocks_high: blocks_high, width: width, height: height,
456            blocks: Vec::with_capacity(total_blocks) };
457
458        for _ in 0..total_blocks {
459            let s0 = subblocks.next().unwrap();
460            let s1 = subblocks.next().unwrap();
461            let s2 = subblocks.next().unwrap();
462            let s3 = subblocks.next().unwrap();
463
464            let block = EncodedMacroBlock { subblocks: [
465                DctQuantizedMatrix8x8::from_slice(s0),
466                DctQuantizedMatrix8x8::from_slice(s1),
467                DctQuantizedMatrix8x8::from_slice(s2),
468                DctQuantizedMatrix8x8::from_slice(s3),
469            ] };
470
471            enc_plane.blocks.push(block);
472        }
473
474        #[cfg(feature = "multithreading")]
475        VideoPlane::decode_plane_into(&enc_plane, q_table, target, tp);
476
477        #[cfg(not(feature = "multithreading"))]
478        VideoPlane::decode_plane_into(&enc_plane, q_table, target);
479    }
480
481    fn deserialize_plane_delta(width: usize, height: usize, headers: &mut Iter<DeltaBlockHeader>, subblocks: &mut ChunksExact<i16>, q_table: &[i32;64], target: &mut VideoPlane,
482        #[cfg(feature = "multithreading")] tp: &rayon::ThreadPool) {
483        let blocks_wide = width / 16;
484        let blocks_high = height / 16;
485        let total_blocks = blocks_wide * blocks_high;
486
487        let mut enc_plane = EncodedPPlane { blocks_wide: blocks_wide, blocks_high: blocks_high, width: width, height: height,
488            blocks: Vec::with_capacity(total_blocks) };
489
490        for _ in 0..total_blocks {
491            let header = headers.next().unwrap();
492
493            let s0 = subblocks.next().unwrap();
494            let s1 = subblocks.next().unwrap();
495            let s2 = subblocks.next().unwrap();
496            let s3 = subblocks.next().unwrap();
497
498            let block = DeltaEncodedMacroBlock {
499                motion_x: header.mvec_x,
500                motion_y: header.mvec_y,
501                subblocks: if header.has_coeff { Some([
502                    DctQuantizedMatrix8x8::from_slice(s0),
503                    DctQuantizedMatrix8x8::from_slice(s1),
504                    DctQuantizedMatrix8x8::from_slice(s2),
505                    DctQuantizedMatrix8x8::from_slice(s3),
506                ]) } else { None }
507            };
508
509            enc_plane.blocks.push(block);
510        }
511
512        #[cfg(feature = "multithreading")]
513        VideoPlane::decode_plane_delta_into(&enc_plane, target, q_table, tp);
514
515        #[cfg(not(feature = "multithreading"))]
516        VideoPlane::decode_plane_delta_into(&enc_plane, target, q_table);
517    }
518}
pfv_rs/dec.rs

pfv_rs/
dec.rs
