From 74b8e965f0078c05cd5780c1352751fd5e546c20 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADtor=20Vasconcellos?= Date: Wed, 17 Apr 2024 19:16:36 -0300 Subject: [PATCH] Data models and start populating MediaInfo in rust --- Cargo.lock | 1 + core/prisma/schema.prisma | 2 +- crates/ffmpeg/Cargo.toml | 6 +- crates/ffmpeg/src/dict.rs | 21 +++--- crates/ffmpeg/src/format_ctx.rs | 98 +++++++++++++++++++++++++++- crates/ffmpeg/src/lib.rs | 1 + crates/ffmpeg/src/model.rs | 110 ++++++++++++++++++++++++++++++++ crates/ffmpeg/src/probe.rs | 17 ++++- 8 files changed, 237 insertions(+), 19 deletions(-) create mode 100644 crates/ffmpeg/src/model.rs diff --git a/Cargo.lock b/Cargo.lock index f7bd0ee3b00e..2d4c4d8e7e53 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8830,6 +8830,7 @@ dependencies = [ name = "sd-ffmpeg" version = "0.1.0" dependencies = [ + "chrono", "ffmpeg-sys-next", "libc", "tempfile", diff --git a/core/prisma/schema.prisma b/core/prisma/schema.prisma index c0ea61eea744..33cc7e744091 100644 --- a/core/prisma/schema.prisma +++ b/core/prisma/schema.prisma @@ -277,7 +277,7 @@ model MediaInfo { id Int @id @default(autoincrement()) // Internal FFmpeg properties - format String? + formats String? duration BigInt? start_time BigInt? bitrate BigInt? diff --git a/crates/ffmpeg/Cargo.toml b/crates/ffmpeg/Cargo.toml index 315eddacc55b..6abb6ed95143 100644 --- a/crates/ffmpeg/Cargo.toml +++ b/crates/ffmpeg/Cargo.toml @@ -10,13 +10,13 @@ repository = { workspace = true } edition = { workspace = true } [dependencies] +chrono = { workspace = true } +ffmpeg-sys-next = "6.0.1" +libc = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["fs", "rt"] } tracing = { workspace = true } webp = { workspace = true } -libc = { workspace = true } -ffmpeg-sys-next = "6.0.1" - [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/ffmpeg/src/dict.rs b/crates/ffmpeg/src/dict.rs index 931e243da3f1..6019d3dc26b4 100644 --- a/crates/ffmpeg/src/dict.rs +++ b/crates/ffmpeg/src/dict.rs @@ -73,19 +73,18 @@ impl<'a> Iterator for FFmpegDictIter<'a> { type Item = (String, String); fn next(&mut self) -> Option<(String, String)> { - while unsafe { + unsafe { self.prev = av_dict_iterate(self.dict, self.prev); - self.prev - } != ptr::null() - { - let key = unsafe { CStr::from_ptr((*self.prev).key) }; - let value = unsafe { CStr::from_ptr((*self.prev).value) }; - return Some(( - key.to_string_lossy().into_owned(), - value.to_string_lossy().into_owned(), - )); + } + if self.prev.is_null() { + return None; } - None + let key = unsafe { CStr::from_ptr((*self.prev).key) }; + let value = unsafe { CStr::from_ptr((*self.prev).value) }; + return Some(( + key.to_string_lossy().into_owned(), + value.to_string_lossy().into_owned(), + )); } } diff --git a/crates/ffmpeg/src/format_ctx.rs b/crates/ffmpeg/src/format_ctx.rs index 654ef91b7a97..70a60a501000 100644 --- a/crates/ffmpeg/src/format_ctx.rs +++ b/crates/ffmpeg/src/format_ctx.rs @@ -1,10 +1,13 @@ -use crate::{dict::FFmpegDict, error::Error, utils::check_error}; +use crate::{dict::FFmpegDict, error::Error, model::MediaMetadata, utils::check_error}; use ffmpeg_sys_next::{ avformat_close_input, avformat_find_stream_info, avformat_open_input, AVFormatContext, }; -use std::{ffi::CString, ptr}; +use std::{ + ffi::{CStr, CString}, + ptr, +}; #[derive(Debug)] pub(crate) struct FFmpegFormatContext { @@ -12,6 +15,97 @@ pub(crate) struct FFmpegFormatContext { } impl FFmpegFormatContext { + pub(crate) fn formats(&self) -> Option> { + if self.data.is_null() { + return None; + } + + let format: *const ffmpeg_sys_next::AVInputFormat = unsafe { (*self.data).iformat }; + if format.is_null() { + return None; + } + + let name = unsafe { (*format).name }; + if name.is_null() { + return None; + } + + let c_str = unsafe { CStr::from_ptr(name) }; + let string = c_str.to_string_lossy().into_owned(); + let entries: Vec = string + .split(',') + .map(|entry| entry.trim().to_string()) + .filter(|entry| !entry.is_empty()) + .collect(); + + Some(entries) + } + + pub(crate) fn metadata(&self) -> Option { + if self.data.is_null() { + return None; + } + + let metadata_ptr = unsafe { (*self.data).metadata }; + if metadata_ptr.is_null() { + return None; + } + + let mut media_metadata = MediaMetadata::default(); + + let metadata = FFmpegDict::new(Some(metadata_ptr)); + for (key, value) in metadata.into_iter() { + match key.as_str() { + "album" => media_metadata.album = Some(value.clone()), + "album_artist" => media_metadata.album_artist = Some(value.clone()), + "artist" => media_metadata.artist = Some(value.clone()), + "comment" => media_metadata.comment = Some(value.clone()), + "composer" => media_metadata.composer = Some(value.clone()), + "copyright" => media_metadata.copyright = Some(value.clone()), + "creation_time" => { + if let Ok(creation_time) = chrono::DateTime::parse_from_rfc3339(&value) { + media_metadata.creation_time = Some(creation_time.into()); + } + } + "date" => { + if let Ok(date) = chrono::DateTime::parse_from_rfc3339(&value) { + media_metadata.date = Some(date.into()); + } + } + "disc" => { + if let Ok(disc) = value.parse() { + media_metadata.disc = Some(disc); + } + } + "encoder" => media_metadata.encoder = Some(value.clone()), + "encoded_by" => media_metadata.encoded_by = Some(value.clone()), + "filename" => media_metadata.filename = Some(value.clone()), + "genre" => media_metadata.genre = Some(value.clone()), + "language" => media_metadata.language = Some(value.clone()), + "performer" => media_metadata.performer = Some(value.clone()), + "publisher" => media_metadata.publisher = Some(value.clone()), + "service_name" => media_metadata.service_name = Some(value.clone()), + "service_provider" => media_metadata.service_provider = Some(value.clone()), + "title" => media_metadata.title = Some(value.clone()), + "track" => { + if let Ok(track) = value.parse() { + media_metadata.track = Some(track); + } + } + "variant_bitrate" => { + if let Ok(variant_bitrate) = value.parse() { + media_metadata.variant_bitrate = Some(variant_bitrate); + } + } + _ => { + media_metadata.custom.insert(key.clone(), value.clone()); + } + } + } + + Some(media_metadata) + } + pub(crate) fn as_mut_ptr(&mut self) -> *mut AVFormatContext { self.data } diff --git a/crates/ffmpeg/src/lib.rs b/crates/ffmpeg/src/lib.rs index f81e63d955ff..5f3b023281cf 100644 --- a/crates/ffmpeg/src/lib.rs +++ b/crates/ffmpeg/src/lib.rs @@ -10,6 +10,7 @@ mod dict; mod error; mod film_strip; mod format_ctx; +mod model; mod movie_decoder; mod probe; mod thumbnailer; diff --git a/crates/ffmpeg/src/model.rs b/crates/ffmpeg/src/model.rs new file mode 100644 index 000000000000..1e4c21c24db5 --- /dev/null +++ b/crates/ffmpeg/src/model.rs @@ -0,0 +1,110 @@ +use std::collections::HashMap; + +#[derive(Default)] +pub struct MediaMetadata { + pub album: Option, + pub album_artist: Option, + pub artist: Option, + pub comment: Option, + pub composer: Option, + pub copyright: Option, + pub creation_time: Option>, + pub date: Option>, + pub disc: Option, + pub encoder: Option, + pub encoded_by: Option, + pub filename: Option, + pub genre: Option, + pub language: Option, + pub performer: Option, + pub publisher: Option, + pub service_name: Option, + pub service_provider: Option, + pub title: Option, + pub track: Option, + pub variant_bitrate: Option, + pub custom: HashMap, +} + +pub struct MediaChapter { + pub id: i32, + pub start: Option, + pub end: Option, + pub metadata: MediaMetadata, +} + +pub struct MediaVideoProps { + pub pixel_format: Option, + pub color_range: Option, + pub bits_per_channel: Option, + pub color_space: Option, + pub color_primaries: Option, + pub color_transfer: Option, + pub field_order: Option, + pub chroma_location: Option, + pub coded_width: Option, + pub coded_height: Option, + pub aspect_ratio_num: Option, + pub aspect_ratio_den: Option, + pub properties: Option, +} + +pub struct MediaAudioProps { + pub delay: Option, + pub padding: Option, + pub sample_rate: Option, + pub sample_format: Option, + pub bit_per_sample: Option, + pub channel_layout: Option, +} + +pub struct MediaSubtitleProps { + pub width: Option, + pub height: Option, +} + +enum Props { + MediaVideoProps, + MediaAudioProps, + MediaSubtitleProps, +} + +pub struct MediaCodec { + pub kind: Option, + pub tag: Option, + pub name: Option, + pub profile: Option, + pub bit_rate: Option, + pub props: Option, +} + +pub struct MediaStream { + pub id: i32, + pub name: Option, + pub codec: Option, + pub aspect_ratio_num: Option, + pub aspect_ratio_den: Option, + pub frames_per_second_num: Option, + pub frames_per_second_den: Option, + pub time_base_real_den: Option, + pub time_base_real_num: Option, + pub dispositions: Option, + pub metadata: MediaMetadata, +} + +pub struct MediaProgram { + pub id: i32, + pub name: Option, + pub streams: Vec, + pub metadata: MediaMetadata, +} + +pub struct MediaInfo { + pub formats: Option>, + pub duration: Option, + pub start_time: Option, + pub bitrate: Option, + pub chapters: Vec, + pub programs: Vec, + pub metadata: Option, +} diff --git a/crates/ffmpeg/src/probe.rs b/crates/ffmpeg/src/probe.rs index 730a9c40aa78..0f9a19838f52 100644 --- a/crates/ffmpeg/src/probe.rs +++ b/crates/ffmpeg/src/probe.rs @@ -2,9 +2,12 @@ use crate::{ dict::FFmpegDict, error::Error, format_ctx::FFmpegFormatContext, + model::MediaInfo, utils::{from_path, CSTRING_ERROR_MSG}, }; + use ffmpeg_sys_next::{av_log_set_level, AV_LOG_FATAL}; + use std::{ffi::CString, path::Path}; pub fn probe(filename: impl AsRef) -> Result<(), Error> { @@ -15,7 +18,7 @@ pub fn probe(filename: impl AsRef) -> Result<(), Error> { // Dictionary to store format options let mut format_opts = FFmpegDict::new(None); - // Some MPEGTS specific option (copied and pasted from ffprobe) + // Some MPEGTS specific option (copied from ffprobe) let scan_all_pmts = CString::new("scan_all_pmts").expect(CSTRING_ERROR_MSG); format_opts.set( scan_all_pmts.to_owned(), @@ -28,8 +31,18 @@ pub fn probe(filename: impl AsRef) -> Result<(), Error> { // Reset MPEGTS specific option format_opts.reset(scan_all_pmts)?; - // Read packets of a media file to get stream information. + // Read packets of media file to get stream information. fmt_ctx.find_stream_info()?; + let media_info = MediaInfo { + formats: fmt_ctx.formats(), + duration: "" , // Option + start_time: "" , // Option + bitrate: "" , // Option + chapters: "" , // Vec + programs: "" , // Vec + metadata: fmt_ctx.metadata(), // MediaMetadata + }; + Ok(()) }