mp4_video_samples_to_h264.rs

#![feature(int_to_from_bytes)]

extern crate mp4parse;

use std::fmt;
use std::fs::{ self, OpenOptions, };
use std::io::{ Read, Write, Seek, SeekFrom };

struct Chunks<'a> {
    track: &'a mp4parse::Track,
    stsc_sample_index: usize,
    index: usize,
    sample_index: usize,
}

impl<'a> Iterator for Chunks<'a> {
    type Item = Samples<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.track.stco {
            Some(ref chunk_offset_box) => {
                let index = self.index;
                
                if index >= chunk_offset_box.offsets.len() {
                    return None;
                }

                let offset = chunk_offset_box.offsets[index];
                
                let samples_per_chunk = match self.track.stsc {
                    Some(ref sample_to_chunk) => {
                        let chunk_idx = (self.index + 1) as u32;

                        for (i, ref stc) in sample_to_chunk.samples[self.stsc_sample_index..]
                                                .iter().enumerate() {
                            if stc.first_chunk == chunk_idx {
                                self.stsc_sample_index = i;
                            }
                        }

                        sample_to_chunk.samples[self.stsc_sample_index].samples_per_chunk as usize
                    },
                    None => panic!("Sample to Chunk 至少需要包含一条指示信息。"),
                };

                assert_eq!(samples_per_chunk > 0, true);

                let sample_index = self.sample_index;

                self.index += 1;
                self.sample_index += samples_per_chunk;

                if let Some(ref stsz) = self.track.stsz {
                    if stsz.sample_size == 0 {
                        if self.sample_index >= stsz.sample_sizes.len() {
                            println!("[WARN] 部分样本遭到遗弃 ...", );
                            return None;
                        }
                    }
                }
                
                Some(Samples {
                    track: &self.track,
                    chunk_index: index,
                    chunk_offset: offset,
                    sample_offset: offset,
                    sample_count: samples_per_chunk,
                    sample_index: sample_index,
                })
            },
            None => None,
        }
    }

    fn count(self) -> usize {
        match self.track.stco {
            Some(ref chunk_offset_box) => chunk_offset_box.offsets.len(),
            None => 0,
        }
    }
}

#[derive(Debug)]
pub struct Sample {
    pub chunk_index: usize,
    pub chunk_offset: u64,
    pub index: usize,
    pub offset: u64,
    pub size: u32,
    pub delta: u32, // 样本显示时间
}

struct Samples<'a> {
    track: &'a mp4parse::Track,
    
    chunk_index: usize,
    chunk_offset: u64,

    sample_offset: u64,
    sample_count: usize,
    sample_index: usize,
}


impl<'a> Iterator for Samples<'a> {
    type Item = Sample;

    fn next(&mut self) -> Option<Self::Item> {
        if self.sample_count == 0 {
            return None;
        }

        // stsz
        let sample_size = if let Some(ref stsz) = self.track.stsz {
            if stsz.sample_size > 0 {
                stsz.sample_size
            } else { 
                stsz.sample_sizes[self.sample_index]
            }
        } else {
            0
        };

        if sample_size == 0 {
            panic!("sample size 必须大于零!");
        }

        // stts 同步表
        let delta = if let Some(ref stts) = self.track.stts {
            let mut i = 0u32;
            let mut delta = 0u32;
            for item in stts.samples.iter() {
                i += item.sample_count;

                if i >= self.sample_index as u32 {
                    delta = item.sample_delta;
                }
            }

            delta
        } else {
            panic!("时间同步表缺失！");
        };

        // TODO: stss, ctts

        let ret = Some(Sample {
                chunk_index: self.chunk_index,
                chunk_offset: self.chunk_offset,
                index: self.sample_index,
                offset: self.sample_offset,
                size: sample_size,
                delta: delta,
        });

        self.sample_offset += sample_size as u64;
        self.sample_count -= 1;
        self.sample_index += 1;

        ret
    }
}




pub struct Mp4Track {
    pub id: usize,
    pub kind: mp4parse::TrackType,
    pub track_id: u32,
    pub empty_duration: Option<mp4parse::MediaScaledTime>,
    pub media_time: Option<mp4parse::TrackScaledTime<u64>>,
    pub timescale: Option<mp4parse::TrackTimeScale<u64>>,
    pub duration: Option<mp4parse::TrackScaledTime<u64>>,
    pub codec_type: mp4parse::CodecType,
    pub data: mp4parse::SampleEntry,
    pub samples: Vec<Sample>,
}

pub struct Mp4File {
    pub timescale: Option<mp4parse::MediaTimeScale>,
    pub mvex: Option<mp4parse::MovieExtendsBox>,
    pub psshs: Vec<mp4parse::ProtectionSystemSpecificHeaderBox>,
    pub tracks: Vec<Mp4Track>,
}

impl fmt::Debug for Mp4Track {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "\tTrack: #{}", self.track_id);
        writeln!(f, "\t\tKind: {:?} ", self.kind);
        writeln!(f, "\t\tEmpty duration: {:?} ", self.empty_duration);
        writeln!(f, "\t\tMedia Time: {:?} ", self.media_time);
        writeln!(f, "\t\tTimescale: {:?} ", self.timescale);
        writeln!(f, "\t\tDuration: {:?} ", self.duration);
        writeln!(f, "\t\tCodec: {:?} ", self.codec_type);
        writeln!(f, "\t\tData: {:?} ", self.data);
        writeln!(f, "\t\tSamples: {:?} ", self.samples.len());
        Ok(())
    }
}

impl fmt::Debug for Mp4File {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "Timescale: {:?} ", self.timescale);
        writeln!(f, "Mvex: {:?} ", self.mvex);
        writeln!(f, "Psshs: {:?} ", self.psshs);
        writeln!(f, "Tracks:");
        
        for track in &self.tracks {
            writeln!(f, "{:?}", track);
        }

        Ok(())
    }
}



fn parse<F: Read>(mut input_file: F) -> Result<Mp4File, mp4parse::Error> {
    let mut mp4_media_ctx = mp4parse::MediaContext::new();
    mp4parse::read_mp4(&mut input_file, &mut mp4_media_ctx).unwrap();

    let mut tracks = vec![];
    for track in mp4_media_ctx.tracks {
        let mut samples = vec![];
        {
            let chunks = Chunks { track: &track, index: 0, stsc_sample_index: 0, sample_index: 0 };
            for chunk_samples in chunks {
                for sample in chunk_samples {
                    samples.push(sample);
                }
            }
        }

        tracks.push(Mp4Track {
            id: track.id,
            track_id: track.track_id.unwrap(),
            kind: track.track_type,
            empty_duration: track.empty_duration,
            media_time: track.media_time,
            timescale: track.timescale,
            duration: track.duration,
            codec_type: track.codec_type,
            data: track.data.unwrap(),
            samples: samples,
        });
    }

    Ok(Mp4File {
        timescale: mp4_media_ctx.timescale,
        mvex: mp4_media_ctx.mvex,
        psshs: mp4_media_ctx.psshs,
        tracks: tracks,
    })
}



fn mp4_samples_to_h264<F: Read + Write + Seek>(mut input_file: F, mut output_file: F, samples: &[Sample]) {
    let mut buffer = Vec::new();

    for video_sample in samples.iter() {
        buffer.resize(video_sample.size as usize, 0u8);

        input_file.seek(SeekFrom::Start(video_sample.offset)).unwrap();
        input_file.read_exact(&mut buffer).unwrap();

        let mut start: usize = 0usize;
        loop {
            if start >= buffer.len() {
                break;
            }

            let size = u32::from_be_bytes([
                buffer[start + 0], buffer[start + 1],
                buffer[start + 2], buffer[start + 3]
            ]) as usize;

            start += 4;
            let end = start + size;

            let au = &buffer[start..end];
            
            if au[size-2] == 0 && au[size-1] == 0 {
                output_file.write(&[0u8, 0, 1]).unwrap();
                output_file.write_all(&au);
                output_file.write(&[3]).unwrap();
            } else {
                output_file.write(&[0u8, 0, 0, 1]).unwrap();
                output_file.write_all(&au);
            }

            start = end;
        }
    }
}

fn main() {
    let input_filepath = "a.mp4";
    let output_filepath = "a.h264";

    let mut mp4_input_file = fs::File::open(input_filepath).unwrap();
    let mut h264_output_file = {
        let _ = fs::remove_file(output_filepath);
        OpenOptions::new().create_new(true).write(true).open(output_filepath).unwrap()
    };

    let mp4 = parse(&mut mp4_input_file).unwrap();

    print!("{:?}", mp4);

    for track in &mp4.tracks {
        if track.kind == mp4parse::TrackType::Video {
            // NOTE: MP4 Video Track Samples TO H264 Byte Stream.
            mp4_samples_to_h264(&mut mp4_input_file, &mut h264_output_file, &track.samples);
        }
    }
}