aboutsummaryrefslogtreecommitdiff
path: root/remuxer/src/snippet.rs
blob: f9d386a210efdb8d13db537d15ef3cde93ab02c3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2023 metamuffin <metamuffin.org>
*/

use crate::{
    ebml_header, ebml_segment_info, ebml_track_entry, seek_index::get_seek_index,
    segment_extractor::SegmentExtractIter,
};
use anyhow::{anyhow, Context, Result};
use jellycommon::{LocalTrack, NodePublic};
use jellymatroska::{read::EbmlReader, write::EbmlWriter, Master, MatroskaTag};
use log::{debug, info};
use std::{fs::File, io::Write, ops::Range, path::Path};

const SNIPPET_LENGTH: f64 = 2.;

pub fn snippet_index(
    path_base: &Path,
    item: &NodePublic,
    track_sources: &Vec<LocalTrack>,
    track: usize,
) -> Result<Vec<Range<f64>>> {
    let media_info = item.media.as_ref().unwrap();
    let private = &track_sources[track];
    let source_path = path_base.join(&private.path);
    let index = get_seek_index(&source_path)?;
    let index = index
        .get(&(private.track as u64))
        .ok_or(anyhow!("track missing"))?;
    let average_kf_interval = media_info.duration / index.keyframes.len() as f64;
    let kf_per_snip = (SNIPPET_LENGTH / average_kf_interval).ceil() as usize;
    debug!("average keyframe interval: {average_kf_interval}");
    debug!(" => keyframes per snippet {kf_per_snip}");

    let n_snips = index.keyframes.len().div_ceil(kf_per_snip);
    Ok((0..n_snips)
        .map(|i| {
            let start = index.blocks[index.keyframes[i * kf_per_snip]].pts as f64 / 1000.;
            let end = index
                .keyframes
                .get((i + 1) * kf_per_snip)
                .map(|i| index.blocks[*i].pts as f64 / 1000.)
                .unwrap_or(media_info.duration);
            start..end
        })
        .collect())
}

pub fn write_snippet_into(
    writer: impl Write,
    path_base: &Path,
    item: &NodePublic,
    track_sources: Vec<LocalTrack>,
    track: usize,
    webm: bool,
    n: usize,
) -> anyhow::Result<()> {
    info!("writing snippet {n} of {:?} (track #{track})", item.title);
    let mut output = EbmlWriter::new(writer, 0);

    let media_info = item.media.as_ref().unwrap();
    let info = media_info
        .tracks
        .get(track)
        .ok_or(anyhow!("track not available"))?
        .to_owned();
    let private = &track_sources[track];
    let source_path = path_base.join(&private.path);
    let mapped = 1;
    info!(
        "\t- {track} {source_path:?} ({} => {mapped})",
        private.track
    );
    info!("\t    {}", info);
    let file = File::open(&source_path).context("opening source file")?;
    let index = get_seek_index(&source_path)?;
    let index = index
        .get(&(private.track as u64))
        .ok_or(anyhow!("track missing"))?
        .to_owned();
    debug!("\t    seek index: {} blocks loaded", index.blocks.len());
    let mut reader = EbmlReader::new(file);

    // TODO maybe refactor this to approx. contant time per snippet
    let average_kf_interval = media_info.duration / index.keyframes.len() as f64;
    let kf_per_snip = (SNIPPET_LENGTH / average_kf_interval).ceil() as usize;
    debug!("average keyframe interval: {average_kf_interval}");
    debug!(" => keyframes per snippet {kf_per_snip}");

    let start_block_index = *index
        .keyframes
        .get(n * kf_per_snip)
        .ok_or(anyhow!("snippet index out of range"))?;
    let end_block_index = *index
        .keyframes
        .get((n + 1) * kf_per_snip)
        .unwrap_or(&index.blocks.len());
    let start_block = &index.blocks[start_block_index];
    let last_block = &index.blocks[end_block_index - 1];

    reader.seek(start_block.source_off, MatroskaTag::Cluster(Master::Start))?;

    output.write_tag(&ebml_header(webm))?;
    output.write_tag(&MatroskaTag::Segment(Master::Start))?;
    output.write_tag(&ebml_segment_info(
        format!("{} (track {track}; snippet {n})", item.title),
        (last_block.pts - start_block.pts) as f64 / 1000.,
    ))?;
    output.write_tag(&MatroskaTag::Tags(Master::Collected(vec![])))?;
    output.write_tag(&MatroskaTag::Tracks(Master::Collected(vec![
        ebml_track_entry(mapped, &info, private.codec_private.clone()),
    ])))?;

    let mut reader = SegmentExtractIter::new(&mut reader, private.track as u64);

    let mut blocks = vec![MatroskaTag::Timestamp(0)];
    for i in start_block_index..end_block_index {
        let index_block = &index.blocks[i];
        let mut block = reader.next()?;

        assert_eq!(index_block.size, block.data.len(), "seek index is wrong");

        block.track = 1;
        block.timestamp_off = (index_block.pts - start_block.pts).try_into().unwrap();
        blocks.push(MatroskaTag::SimpleBlock(block.dump()))
    }
    output.write_tag(&MatroskaTag::Cluster(Master::Collected(blocks)))?;

    output.write_tag(&MatroskaTag::Segment(Master::End))?;
    Ok(())
}