use std::{path::PathBuf, time::Duration};
use crate::{
decoder::{AudioDecoder, VideoDecoder, VideoDecoderOptions},
media::MediaStreamId,
metadata::{ContainerFormat, SourceId},
processor::{
MediaProcessor, MediaProcessorInput, MediaProcessorOutput, MediaProcessorSpec,
MediaProcessorWorkloadHint,
},
reader::{AudioReader, VideoReader},
scheduler::Scheduler,
stats::ProcessorStats,
types::CodecName,
video::{VideoFormat, VideoFrame},
video_h264::H264AnnexBNalUnits,
};
use orfail::OrFail;
use shiguredo_openh264::Openh264Library;
const AUDIO_ENCODED_STREAM_ID: MediaStreamId = MediaStreamId::new(0);
const VIDEO_ENCODED_STREAM_ID: MediaStreamId = MediaStreamId::new(1);
const AUDIO_DECODED_STREAM_ID: MediaStreamId = MediaStreamId::new(2);
const VIDEO_DECODED_STREAM_ID: MediaStreamId = MediaStreamId::new(3);
pub fn run(mut args: noargs::RawArgs) -> noargs::Result<()> {
let decode: bool = noargs::flag("decode")
.doc("指定された場合にはデコードまで行います")
.take(&mut args)
.is_present();
let openh264: Option<PathBuf> = noargs::opt("openh264")
.ty("PATH")
.env("HISUI_OPENH264_PATH")
.doc("OpenH264 の共有ライブラリのパス")
.take(&mut args)
.present_and_then(|a| a.value().parse())?;
let input_file_path: PathBuf = noargs::arg("INPUT_FILE")
.example("/path/to/archive.mp4")
.doc("情報取得対象の録画ファイル(.mp4|.webm)")
.take(&mut args)
.then(|a| a.value().parse())?;
if let Some(help) = args.finish()? {
print!("{help}");
return Ok(());
}
let format = match input_file_path
.extension()
.unwrap_or_default()
.to_string_lossy()
.as_ref()
{
"mp4" => ContainerFormat::Mp4,
"webm" => ContainerFormat::Webm,
ext => {
return Err(
orfail::Failure::new(format!("unsupported container format: {ext}")).into(),
);
}
};
let mut scheduler = Scheduler::new();
let dummy_source_id = SourceId::new("inspect");
let reader = AudioReader::new(
AUDIO_ENCODED_STREAM_ID,
dummy_source_id.clone(),
format,
Duration::ZERO,
vec![input_file_path.clone()],
)
.or_fail()?;
scheduler.register(reader).or_fail()?;
let reader = VideoReader::new(
VIDEO_ENCODED_STREAM_ID,
dummy_source_id.clone(),
format,
Duration::ZERO,
vec![input_file_path.clone()],
)
.or_fail()?;
scheduler.register(reader).or_fail()?;
if decode {
let decoder =
AudioDecoder::new_opus(AUDIO_ENCODED_STREAM_ID, AUDIO_DECODED_STREAM_ID).or_fail()?;
scheduler.register(decoder).or_fail()?;
}
if decode {
let options = VideoDecoderOptions {
openh264_lib: openh264
.clone()
.map(Openh264Library::load)
.transpose()
.or_fail()?,
decode_params: Default::default(),
engines: None,
};
let decoder = VideoDecoder::new(VIDEO_ENCODED_STREAM_ID, VIDEO_DECODED_STREAM_ID, options);
scheduler.register(decoder).or_fail()?;
}
scheduler
.register(OutputPrinter::new(input_file_path.clone(), format, decode))
.or_fail()?;
scheduler.run().or_fail()?;
Ok(())
}
#[derive(Debug)]
struct AudioSampleInfo {
timestamp: Duration,
duration: Duration,
data_size: usize,
decoded_data_size: Option<usize>,
}
impl nojson::DisplayJson for AudioSampleInfo {
fn fmt(&self, f: &mut nojson::JsonFormatter<'_, '_>) -> std::fmt::Result {
f.set_indent_size(0);
f.object(|f| {
f.member("timestamp_us", self.timestamp.as_micros())?;
f.member("duration_us", self.duration.as_micros())?;
f.member("data_size", self.data_size)?;
if let Some(v) = self.decoded_data_size {
f.member("decoded_data_size", v)?;
}
Ok(())
})?;
f.set_indent_size(2);
Ok(())
}
}
#[derive(Debug)]
struct VideoSampleInfo {
timestamp: Duration,
duration: Duration,
data_size: usize,
keyframe: bool,
codec_specific_info: Option<VideoCodecSpecificInfo>,
decoded_data_size: Option<usize>,
width: Option<usize>,
height: Option<usize>,
}
impl VideoSampleInfo {
fn update(&mut self, decoded: &VideoFrame) {
self.decoded_data_size = Some(decoded.data.len());
self.width = Some(decoded.width);
self.height = Some(decoded.height);
}
}
impl nojson::DisplayJson for VideoSampleInfo {
fn fmt(&self, f: &mut nojson::JsonFormatter<'_, '_>) -> std::fmt::Result {
f.set_indent_size(0);
f.object(|f| {
f.member("timestamp_us", self.timestamp.as_micros())?;
f.member("duration_us", self.duration.as_micros())?;
f.member("data_size", self.data_size)?;
f.member("keyframe", self.keyframe)?;
match &self.codec_specific_info {
None => {}
Some(VideoCodecSpecificInfo::H264 { nalus }) => {
f.member("nalus", nalus)?;
}
}
if let Some(v) = self.decoded_data_size {
f.member("decoded_data_size", v)?;
}
if let Some(v) = self.width {
f.member("width", v)?;
}
if let Some(v) = self.height {
f.member("height", v)?;
}
Ok(())
})?;
f.set_indent_size(2);
Ok(())
}
}
#[derive(Debug)]
struct H264NalUnitInfo {
ty: u8,
nri: u8,
}
impl nojson::DisplayJson for H264NalUnitInfo {
fn fmt(&self, f: &mut nojson::JsonFormatter<'_, '_>) -> std::fmt::Result {
f.object(|f| {
f.member("type", self.ty)?;
f.member("nri", self.nri)
})
}
}
#[derive(Debug)]
enum VideoCodecSpecificInfo {
H264 { nalus: Vec<H264NalUnitInfo> },
}
impl VideoCodecSpecificInfo {
fn new(sample: &VideoFrame) -> Option<Self> {
match sample.format {
VideoFormat::H264AnnexB => {
let mut nalus = Vec::new();
for nalu in H264AnnexBNalUnits::new(&sample.data) {
match nalu {
Ok(nalu) => {
let header_byte = nalu.data.first()?;
let nri = (header_byte >> 5) & 0b11;
nalus.push(H264NalUnitInfo { ty: nalu.ty, nri });
}
Err(_) => return None, }
}
Some(VideoCodecSpecificInfo::H264 { nalus })
}
VideoFormat::H264 => {
let mut nalus = Vec::new();
let mut data = &sample.data[..];
while data.len() > 4 {
let length = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize;
data = &data[4..];
if data.len() < length || length == 0 {
return None; }
let header_byte = data[0];
let nalu_type = header_byte & 0b0001_1111;
let nri = (header_byte >> 5) & 0b11;
nalus.push(H264NalUnitInfo { ty: nalu_type, nri });
data = &data[length..];
}
Some(VideoCodecSpecificInfo::H264 { nalus })
}
_ => None,
}
}
}
#[derive(Debug)]
pub struct OutputPrinter {
path: PathBuf,
format: ContainerFormat,
audio_codec: Option<CodecName>,
video_codec: Option<CodecName>,
audio_samples: Vec<AudioSampleInfo>,
video_samples: Vec<VideoSampleInfo>,
input_stream_ids: Vec<MediaStreamId>,
next_input_stream_index: usize,
}
impl OutputPrinter {
fn new(path: PathBuf, format: ContainerFormat, decode: bool) -> Self {
Self {
path,
format,
audio_codec: None,
video_codec: None,
audio_samples: Vec::new(),
video_samples: Vec::new(),
input_stream_ids: if decode {
vec![
AUDIO_ENCODED_STREAM_ID,
VIDEO_ENCODED_STREAM_ID,
AUDIO_DECODED_STREAM_ID,
VIDEO_DECODED_STREAM_ID,
]
} else {
vec![AUDIO_ENCODED_STREAM_ID, VIDEO_ENCODED_STREAM_ID]
},
next_input_stream_index: 0,
}
}
}
impl MediaProcessor for OutputPrinter {
fn spec(&self) -> MediaProcessorSpec {
MediaProcessorSpec {
input_stream_ids: self.input_stream_ids.clone(),
output_stream_ids: Vec::new(),
stats: ProcessorStats::other("output_printer"),
workload_hint: MediaProcessorWorkloadHint::WRITER,
}
}
fn process_input(&mut self, input: MediaProcessorInput) -> orfail::Result<()> {
let Some(sample) = input.sample else {
self.input_stream_ids.retain(|id| *id != input.stream_id);
self.next_input_stream_index = 0;
return Ok(());
};
match input.stream_id {
AUDIO_ENCODED_STREAM_ID => {
let sample = sample.expect_audio_data().or_fail()?;
if self.audio_codec.is_none() {
self.audio_codec = sample.format.codec_name();
}
self.audio_samples.push(AudioSampleInfo {
timestamp: sample.timestamp,
duration: sample.duration,
data_size: sample.data.len(),
decoded_data_size: None,
});
}
AUDIO_DECODED_STREAM_ID => {
let sample = sample.expect_audio_data().or_fail()?;
let info = self
.audio_samples
.iter_mut()
.rfind(|s| s.decoded_data_size.is_none())
.or_fail()?;
info.decoded_data_size = Some(sample.data.len());
}
VIDEO_ENCODED_STREAM_ID => {
let sample = sample.expect_video_frame().or_fail()?;
if self.video_codec.is_none() {
self.video_codec = sample.format.codec_name();
}
self.video_samples.push(VideoSampleInfo {
timestamp: sample.timestamp,
duration: sample.duration,
data_size: sample.data.len(),
keyframe: sample.keyframe,
codec_specific_info: VideoCodecSpecificInfo::new(&sample),
decoded_data_size: None,
width: None,
height: None,
});
}
VIDEO_DECODED_STREAM_ID => {
let sample = sample.expect_video_frame().or_fail()?;
let info = self
.video_samples
.iter_mut()
.rfind(|s| s.decoded_data_size.is_none())
.or_fail()?;
info.update(&sample);
}
_ => return Err(orfail::Failure::new("BUG: unexpected stream ID")),
}
Ok(())
}
fn process_output(&mut self) -> orfail::Result<MediaProcessorOutput> {
if self.input_stream_ids.is_empty() {
crate::json::pretty_print(self).or_fail()?;
Ok(MediaProcessorOutput::Finished)
} else {
let awaiting_stream_id = self.input_stream_ids[self.next_input_stream_index];
self.next_input_stream_index =
(self.next_input_stream_index + 1) % self.input_stream_ids.len();
Ok(MediaProcessorOutput::pending(awaiting_stream_id))
}
}
}
impl nojson::DisplayJson for OutputPrinter {
fn fmt(&self, f: &mut nojson::JsonFormatter<'_, '_>) -> std::fmt::Result {
f.object(|f| {
f.member("path", &self.path)?;
f.member("format", self.format)?;
if let Some(c) = self.audio_codec {
f.member("audio_codec", c)?;
f.member(
"audio_duration_us",
self.audio_samples
.iter()
.map(|s| s.duration)
.sum::<Duration>()
.as_micros(),
)?;
f.member("audio_sample_count", self.audio_samples.len())?;
f.member("audio_samples", &self.audio_samples)?;
}
if let Some(c) = self.video_codec {
f.member("video_codec", c)?;
f.member(
"video_duration_us",
self.video_samples
.iter()
.map(|s| s.duration)
.sum::<Duration>()
.as_micros(),
)?;
f.member("video_sample_count", self.video_samples.len())?;
f.member(
"video_keyframe_sample_count",
self.video_samples.iter().filter(|s| s.keyframe).count(),
)?;
f.member("video_samples", &self.video_samples)?;
}
Ok(())
})
}
}