├── media └── .gitignore ├── models └── .gitignore ├── rustfmt.toml ├── .gitignore ├── docs ├── infur_crab.png ├── infur_crabs.png ├── infur_onstreet_0.5.png └── infur_onstreet_1.0.png ├── image-ext ├── src │ ├── lib.rs │ └── image_bgr.rs └── Cargo.toml ├── ff-video ├── src │ ├── lib.rs │ ├── error.rs │ ├── decoder.rs │ └── parse.rs └── Cargo.toml ├── infur-test-gen ├── Cargo.toml ├── src │ └── lib.rs └── build.rs ├── .pre-commit-config.yaml ├── infur ├── Cargo.toml └── src │ ├── decode_predict.rs │ ├── main.rs │ ├── app.rs │ ├── processing.rs │ ├── gui.rs │ └── predict_onnx.rs ├── Cargo.toml ├── LICENSE ├── .github └── workflows │ └── test.yaml └── README.md /media/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /models/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | use_small_heuristics = "Max" 2 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | 4 | /.idea 5 | -------------------------------------------------------------------------------- /docs/infur_crab.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ahirner/infur/HEAD/docs/infur_crab.png -------------------------------------------------------------------------------- /docs/infur_crabs.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ahirner/infur/HEAD/docs/infur_crabs.png -------------------------------------------------------------------------------- /docs/infur_onstreet_0.5.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ahirner/infur/HEAD/docs/infur_onstreet_0.5.png -------------------------------------------------------------------------------- /docs/infur_onstreet_1.0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ahirner/infur/HEAD/docs/infur_onstreet_1.0.png -------------------------------------------------------------------------------- /image-ext/src/lib.rs: -------------------------------------------------------------------------------- 1 | mod image_bgr; 2 | 3 | pub use image::imageops; 4 | pub use image::*; 5 | pub use image_bgr::{Bgr, BgrImage}; 6 | -------------------------------------------------------------------------------- /ff-video/src/lib.rs: -------------------------------------------------------------------------------- 1 | mod decoder; 2 | mod error; 3 | mod parse; 4 | 5 | pub use crate::error::{FFVideoError, VideoProcError, VideoResult}; 6 | pub use decoder::{FFMpegDecoder, FFMpegDecoderBuilder}; 7 | -------------------------------------------------------------------------------- /image-ext/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "image-ext" 3 | description = "Extend image crate with BGR pixel type" 4 | version.workspace = true 5 | edition.workspace = true 6 | 7 | [dependencies] 8 | image.workspace = true 9 | -------------------------------------------------------------------------------- /infur-test-gen/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "infur-test-gen" 3 | description = "generate synthetic test media" 4 | version.workspace = true 5 | edition.workspace = true 6 | 7 | [build-dependencies] 8 | ureq = "2" 9 | filetime = "0.2" 10 | -------------------------------------------------------------------------------- /ff-video/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ff-video" 3 | description = "Slim image decoding with ffmpeg." 4 | version.workspace = true 5 | edition.workspace = true 6 | 7 | [dependencies] 8 | image-ext = { path = "../image-ext" } 9 | thiserror.workspace = true 10 | tracing.workspace = true 11 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v4.3.0 4 | hooks: 5 | - id: trailing-whitespace 6 | - id: end-of-file-fixer 7 | - repo: https://github.com/Lucas-C/pre-commit-hooks 8 | rev: v1.3.1 9 | hooks: 10 | - id: forbid-tabs 11 | - repo: local 12 | hooks: 13 | - id: rustfmt 14 | name: rustfmt 15 | language: system 16 | entry: cargo fmt -- 17 | types: [rust] 18 | - id: clippy 19 | name: clippy 20 | language: system 21 | entry: cargo clippy --all-targets --all-features -- -D warnings 22 | types: [rust] 23 | pass_filenames: false 24 | -------------------------------------------------------------------------------- /infur/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "infur" 3 | description.workspace = true 4 | version.workspace = true 5 | edition.workspace = true 6 | 7 | [features] 8 | default = ["persistence"] 9 | persistence = ["eframe/persistence"] 10 | 11 | [dependencies] 12 | fast_image_resize.workspace = true 13 | onnxruntime.workspace = true 14 | thiserror.workspace = true 15 | tracing.workspace = true 16 | serde.workspace = true 17 | once_cell = "1" 18 | eframe = { version = "0.19", features = ["wgpu", "default_fonts"], default-features = false } 19 | tracing-subscriber = { version = "0.3", features = ["ansi", "env-filter", "fmt"], default-features = false } 20 | stable-eyre = "0.2" 21 | image-ext = { path = "../image-ext" } 22 | ff-video = { path = "../ff-video" } 23 | 24 | [dev-dependencies] 25 | infur-test-gen = { "path" = "../infur-test-gen" } 26 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = ["infur", "infur-test-gen", "ff-video", "image-ext"] 3 | default-members = ["infur"] 4 | resolver = "2" 5 | 6 | [workspace.package] 7 | version = "0.2.0" 8 | edition = "2021" 9 | description = "ONNX model inference on video and images" 10 | authors = ["Alexander Hirner"] 11 | license = "MIT" 12 | keywords = ["ONNX", "Neural Networks", "Inference", "Segmentation", "GUI", "Prediction", "Video"] 13 | readme = "README.md" 14 | repository = "https://github.com/ahirner/infur" 15 | homepage = "https://github.com/ahirner/infur" 16 | 17 | [workspace.dependencies] 18 | image = "0.24" 19 | fast_image_resize = { version = "1" } 20 | # need onnxruntime .14 for 0-dim input tolerance (not in .13), 21 | # then furthermore need master to resolve ndarray with tract-core 22 | onnxruntime = { git = "https://github.com/nbigaouette/onnxruntime-rs" } 23 | serde = { version = "1", features = ["derive"] } 24 | thiserror = "1" 25 | tracing = "0.1" 26 | 27 | [profile.dev.package] 28 | image-ext = { opt-level = 3 } 29 | backtrace = { opt-level = 3 } 30 | -------------------------------------------------------------------------------- /infur-test-gen/src/lib.rs: -------------------------------------------------------------------------------- 1 | use std::{env, path::PathBuf}; 2 | 3 | fn media_root() -> PathBuf { 4 | let gen_root = PathBuf::from(&env::var("CARGO_MANIFEST_DIR").unwrap()); 5 | gen_root.parent().unwrap().join("media") 6 | } 7 | 8 | pub fn long_small_video() -> PathBuf { 9 | media_root().join("synth_640x480_40secs_10fps.mp4") 10 | } 11 | 12 | pub fn short_large_video() -> PathBuf { 13 | media_root().join("synth_1280x720_5secs_30fps.mp4") 14 | } 15 | 16 | pub fn fcn_resnet50_12_int8_onnx() -> PathBuf { 17 | let gen_root = PathBuf::from(&env::var("CARGO_MANIFEST_DIR").unwrap()); 18 | gen_root.parent().unwrap().join("models").join("fcn-resnet50-12-int8.onnx") 19 | } 20 | 21 | #[cfg(test)] 22 | mod tests { 23 | use super::*; 24 | 25 | #[test] 26 | fn long_small_video_exists() { 27 | assert!(long_small_video().is_file()) 28 | } 29 | 30 | #[test] 31 | fn short_large_exists() { 32 | assert!(short_large_video().is_file()) 33 | } 34 | 35 | #[test] 36 | fn fcn_resnet50_12_int8_onnx_exists() { 37 | assert!(fcn_resnet50_12_int8_onnx().is_file()) 38 | } 39 | } 40 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 MoonVision 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.github/workflows/test.yaml: -------------------------------------------------------------------------------- 1 | name: test 2 | 3 | on: 4 | push: 5 | 6 | jobs: 7 | lint: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v3 11 | - name: install pre-commit 12 | run: 13 | pip install pre-commit 14 | pre-commit install 15 | - run: pre-commit run --all-files --show-diff-on-failure 16 | env: 17 | INFUR_NO_TEST_GEN: 1 18 | test: 19 | strategy: 20 | fail-fast: false 21 | matrix: 22 | os: [ubuntu-latest, windows-latest, macos-latest] 23 | runs-on: ${{ matrix.os }} 24 | steps: 25 | - name: setup-ffmpeg 26 | uses: FedericoCarboni/setup-ffmpeg@v2 27 | - run: cargo version -v 28 | - run: ffmpeg -version 29 | - uses: actions/checkout@v3 30 | - name: cargo test (*nix) 31 | if: startsWith(matrix.os, 'windows') != true 32 | run: cargo test 33 | - name: cargo test (windows) 34 | if: startsWith(matrix.os, 'windows') 35 | shell: bash 36 | run: | 37 | # workaround so that downloaded dll gets in fact loaded by OS, 38 | # the native link directive in onnxruntime-sys build.rs doesn't work 39 | # https://github.com/nbigaouette/onnxruntime-rs/issues/83 40 | cargo build 41 | # copy .dll next to test .exe 42 | find target -name onnxruntime.dll -exec cp "{}" target/debug/deps \; 43 | cargo test 44 | -------------------------------------------------------------------------------- /ff-video/src/error.rs: -------------------------------------------------------------------------------- 1 | use crate::parse::ParseError; 2 | pub use thiserror::Error; 3 | 4 | /// Results from processing a video. 5 | pub type VideoResult = std::result::Result; 6 | /// Results from parsing video info while processing it. 7 | pub type InfoResult = std::result::Result; 8 | 9 | #[derive(Error, Debug)] 10 | pub enum VideoProcError { 11 | #[error("video IO failed with {msg}")] 12 | IO { 13 | msg: String, 14 | #[source] 15 | source: std::io::Error, 16 | }, 17 | #[error("finished normally")] 18 | FinishedNormally { 19 | #[source] 20 | source: std::io::Error, 21 | }, 22 | #[error("couldn't read an entire image")] 23 | ExactReadError { 24 | #[source] 25 | source: std::io::Error, 26 | }, 27 | #[error("couldn't parse stream info in time ({0})")] 28 | Start(String), 29 | #[error("couldn't obtain {0}")] 30 | MissingValue(String), 31 | #[error("video process exit code: {0}")] 32 | ExitCode(i32), 33 | #[error("other error: {0}")] 34 | Other(String), 35 | } 36 | 37 | impl VideoProcError { 38 | pub(crate) fn is_missing(msg: impl ToString) -> Self { 39 | Self::MissingValue(msg.to_string()) 40 | } 41 | pub(crate) fn explain_io(msg: impl ToString, e: std::io::Error) -> Self { 42 | Self::IO { msg: msg.to_string(), source: e } 43 | } 44 | } 45 | 46 | #[derive(Error, Debug)] 47 | pub enum FFVideoError { 48 | #[error("video processing error: {0}")] 49 | Processing(#[from] VideoProcError), 50 | #[error("parsing error: {0}")] 51 | Parsing(#[from] ParseError), 52 | } 53 | 54 | #[cfg(test)] 55 | mod tests { 56 | // test lifted from image crate 57 | use super::*; 58 | use std::mem; 59 | 60 | #[allow(dead_code)] 61 | // This will fail to compile if the size of this type is large. 62 | const ASSERT_SMALLISH: usize = [0][(mem::size_of::() >= 200) as usize]; 63 | 64 | #[test] 65 | fn test_send_sync_stability() { 66 | fn assert_send_sync() {} 67 | 68 | assert_send_sync::(); 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![test](https://github.com/ahirner/infur/actions/workflows/test.yaml/badge.svg)](https://github.com/ahirner/infur/actions/workflows/test.yaml) 2 | 3 | ## InFur 4 | 5 | `InFur` showcases [ONNX](https://onnx.ai) dense model inference on videos (and images). 6 | 7 | [![Red crab invasion](docs/infur_crabs.png)](https://www.ocregister.com/2016/05/16/red-crab-invasion-beaches-at-newport-beach-laguna-beach-covered-by-the-tiny-crustaceans/) 8 | 9 | ### Requirements 10 | 11 | `Infur` should compile on any Tier 1 platform with a recent [Rust](https://www.rust-lang.org) toolchain. 12 | 13 | You must also have an [ffmpeg](https://ffmpeg.org) executable in your `PATH` (>=4.3 recommended). 14 | 15 | ### Test 16 | 17 | ``` 18 | cargo test 19 | ``` 20 | 21 | `cargo test` will ensure synthetic test videos exist in [./media](./media) 22 | and download a [quantized segmentation model](https://github.com/onnx/models/tree/main/vision/object_detection_segmentation/fcn) 23 | to [./models](./models). 24 | 25 | #### Windows 26 | 27 | The first test run will fail if no `onnxruntime` with Opset support >= 8 is 28 | [on the system path](https://github.com/nbigaouette/onnxruntime-rs/issues/83). 29 | One fix is to copy the `.dll` downloaded by [onnxruntime-sys](https://github.com/nbigaouette/onnxruntime-rs) next to the target `.exe`: 30 | 31 | ```bash 32 | find target -name onnxruntime.dll -exec cp "{}" target/debug/deps \; 33 | cargo test 34 | ``` 35 | 36 | Copy it also to `debug` and `release` to run the main application 37 | (e.g. `cp target/debug/deps/onnxrutime.dll target/release)`. 38 | 39 | ### Use 40 | 41 | You can provide a video URL to start with: 42 | 43 | ``` 44 | cargo run --release -- http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/SubaruOutbackOnStreetAndDirt.mp4 45 | ``` 46 | 47 | Next, paste the path of the segmentation model from the test fixture into the `Inference` text box: 48 | 49 | ``` 50 | models/fcn-resnet50-12-int8.onnx 51 | ``` 52 | 53 | The model's dense multi-class prediction, i.e. a segmentation mask is color-coded (argmax) and 54 | shaded (by confidence): 55 | 56 | ![](docs/infur_onstreet_0.5.png) 57 | 58 | A model's output often varies greatly with the scale of the input image. Thus, you can 59 | tune its scale factor on `Pause`: 60 | 61 | ![](docs/infur_onstreet_1.0.png) 62 | 63 | By default, the app's settings are persisted after closing. 64 | 65 | ### Todos 66 | 67 | The purpose of this crate is to study tradeoffs regarding model inference, native GUIs and 68 | video decoding approaches, in Rust :crab:. 69 | 70 | There are a couple of Todos will make `InFur` more interesting beyond exploring 71 | production-readiness as now: 72 | 73 | - [ ] GATify `type Output` in `trait Processor` 74 | - [ ] bi-linear image scaling 75 | - [ ] [meta-data aware](https://github.com/onnx/onnx/blob/main/docs/MetadataProps.md#image-category-definition) image pre-processing choices 76 | - [ ] softmax if model predictions are logits (and/or clamp confidence shading) 77 | - [ ] class label captions 78 | - [ ] file-picker for model and video input 79 | - [ ] video fast-forward/backward 80 | - [ ] video seeking 81 | -------------------------------------------------------------------------------- /infur-test-gen/build.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | env, 3 | ffi::OsStr, 4 | fs, 5 | io::{self, Write}, 6 | path::{Path, PathBuf}, 7 | process::Command, 8 | }; 9 | 10 | use filetime::{set_file_mtime, FileTime}; 11 | 12 | fn run_ffmpeg_synth( 13 | out_file: impl AsRef, 14 | width: usize, 15 | height: usize, 16 | rate: usize, 17 | duration: usize, 18 | ) { 19 | let mut cmd = Command::new("ffmpeg"); 20 | cmd.args(["-f", "lavfi", "-i"]) 21 | .arg(format!("testsrc=duration={duration}:size={width}x{height}:rate={rate}")) 22 | .args(["-pix_fmt", "yuv420p", "-y"]) 23 | .arg(out_file); 24 | 25 | let status = cmd 26 | .spawn() 27 | .expect("synthesizing video couldn't start, do you have ffmpeg in PATH?") 28 | .wait() 29 | .expect("synthesizing video didn't finish"); 30 | assert!(status.success(), "synthesizing videos didn't finish successfully"); 31 | } 32 | 33 | fn download(source_url: &str, target_file: impl AsRef) { 34 | // borrowed from onnxruntime 35 | let resp = ureq::get(source_url) 36 | .timeout(std::time::Duration::from_secs(300)) 37 | .call() 38 | .unwrap_or_else(|err| panic!("ERROR: Failed to download {}: {:?}", source_url, err)); 39 | 40 | let len = resp.header("Content-Length").and_then(|s| s.parse::().ok()).unwrap(); 41 | let mut reader = resp.into_reader(); 42 | // FIXME: Save directly to the file 43 | let mut buffer = vec![]; 44 | let read_len = reader.read_to_end(&mut buffer).unwrap(); 45 | assert_eq!(buffer.len(), len); 46 | assert_eq!(buffer.len(), read_len); 47 | 48 | let f = fs::File::create(&target_file).unwrap(); 49 | let mut writer = io::BufWriter::new(f); 50 | writer.write_all(&buffer).unwrap(); 51 | } 52 | 53 | /// Makes files look like they were there 60 seconds earlier. 54 | /// 55 | /// We need that for generated files since cargo tests that 56 | /// they are not stale iff. mtime of build artifacts > rerun-if-changed. 57 | fn make_younger(file: impl AsRef) { 58 | let file_meta = fs::metadata(&file).unwrap(); 59 | let mtime = FileTime::from_last_modification_time(&file_meta); 60 | let mtime_before = mtime.unix_seconds() - 60; 61 | set_file_mtime(&file, filetime::FileTime::from_unix_time(mtime_before, 0)).unwrap(); 62 | } 63 | 64 | pub fn main() { 65 | println!("cargo:rerun-if-changed=build.rs"); 66 | 67 | // set in CI to avoid ffmpeg dependency on clippy runs 68 | println!("cargo:rerun-if-env-changed=INFUR_NO_TEST_GEN"); 69 | if std::env::var("INFUR_NO_TEST_GEN").ok().as_deref() == Some("1") { 70 | return; 71 | }; 72 | 73 | let gen_root = PathBuf::from(&env::var("CARGO_MANIFEST_DIR").unwrap()); 74 | let gen_root = 75 | Path::new(&gen_root).parent().expect("wanted parent of manifest for generating test files"); 76 | 77 | // video files 78 | for (width, height, rate, dur) in [(1280, 720, 30, 5), (640, 480, 10, 40)] { 79 | let file = format!("synth_{width}x{height}_{dur}secs_{rate}fps.mp4"); 80 | let dest_path = gen_root.join("media").join(file); 81 | 82 | run_ffmpeg_synth(&dest_path, width, height, rate, dur); 83 | make_younger(&dest_path); 84 | println!("cargo:rerun-if-changed={}", &dest_path.to_string_lossy()); 85 | } 86 | 87 | // models 88 | // segmentation model, see: https://github.com/onnx/models/tree/main/vision/object_detection_segmentation/fcn 89 | let fcn_resnet50_12_int8 = gen_root.join("models").join("fcn-resnet50-12-int8.onnx"); 90 | download("https://github.com/onnx/models/raw/main/vision/object_detection_segmentation/fcn/model/fcn-resnet50-12-int8.onnx", 91 | &fcn_resnet50_12_int8); 92 | make_younger(&fcn_resnet50_12_int8); 93 | println!("cargo:rerun-if-changed={}", &fcn_resnet50_12_int8.to_string_lossy()); 94 | } 95 | -------------------------------------------------------------------------------- /infur/src/decode_predict.rs: -------------------------------------------------------------------------------- 1 | use crate::app::Processor; 2 | use eframe::epaint::{Color32, ColorImage}; 3 | use onnxruntime::ndarray::Array3; 4 | 5 | /// 20 RGB high-contrast BGR/RGB triplets 6 | /// 7 | /// adapted from: 8 | /// and: 9 | const COLORS_PALETTE: [(u8, u8, u8); 20] = [ 10 | (75, 180, 60), 11 | (75, 25, 230), 12 | (25, 225, 255), 13 | (200, 130, 0), 14 | (48, 130, 245), 15 | (240, 240, 70), 16 | (230, 50, 240), 17 | (60, 245, 210), 18 | (180, 30, 145), 19 | (190, 190, 250), 20 | (128, 128, 0), 21 | (255, 190, 230), 22 | (40, 110, 170), 23 | (200, 250, 255), 24 | (0, 0, 128), 25 | (195, 255, 170), 26 | (0, 128, 128), 27 | (180, 215, 255), 28 | (128, 0, 0), 29 | (128, 128, 128), 30 | ]; 31 | 32 | fn color_code(klass: usize, alpha: f32) -> Color32 { 33 | // todo: pre-transform COLORS into linear space 34 | let (r, g, b) = COLORS_PALETTE[klass % COLORS_PALETTE.len()]; 35 | Color32::from_rgba_unmultiplied(r, g, b, (alpha * 255.0f32) as u8) 36 | } 37 | 38 | #[derive(Default)] 39 | pub(crate) struct ColorCode; 40 | 41 | impl Processor for ColorCode { 42 | type Command = (); 43 | type ControlError = (); 44 | /// KxHxW confidences 45 | type Input = Array3; 46 | type Output = Option; 47 | type ProcessResult = (); 48 | 49 | fn control(&mut self, _cmd: Self::Command) -> Result<&mut Self, Self::ControlError> { 50 | Ok(self) 51 | } 52 | 53 | fn advance(&mut self, inp: &Self::Input, out: &mut Self::Output) -> Self::ProcessResult { 54 | let shape = inp.shape(); 55 | let (k, h, w) = (shape[0], shape[1], shape[2]); 56 | 57 | // get or re-create output image 58 | let img = if let Some(ref mut img) = out { 59 | if img.width() != w || img.height() != h { 60 | *img = ColorImage::new([w, h], Color32::BLACK); 61 | } 62 | img 63 | } else { 64 | out.get_or_insert_with(|| ColorImage::new([w, h], Color32::BLACK)) 65 | }; 66 | 67 | let inp_flat = inp.exact_chunks([k, 1, 1]); 68 | img.pixels.iter_mut().zip(inp_flat).for_each(|(col, klasses)| { 69 | let mut k_max = 0; 70 | let mut c_max = 0f32; 71 | klasses.iter().enumerate().for_each(|(i, confidence)| { 72 | if confidence > &c_max { 73 | k_max = i; 74 | c_max = *confidence; 75 | } 76 | }); 77 | *col = color_code(k_max, c_max); 78 | }); 79 | } 80 | 81 | fn is_dirty(&self) -> bool { 82 | false 83 | } 84 | } 85 | 86 | #[cfg(test)] 87 | mod test { 88 | 89 | use onnxruntime::ndarray::Array1; 90 | 91 | use super::*; 92 | 93 | #[test] 94 | fn color_2() { 95 | let c = COLORS_PALETTE[2]; 96 | assert_eq!(color_code(2, 0.5), Color32::from_rgba_unmultiplied(c.0, c.1, c.2, 127)); 97 | } 98 | 99 | #[test] 100 | fn decode_0to1() { 101 | let hm = >::linspace(0., 1., 22 * 24 * 32).into_shape([22, 24, 32]).unwrap(); 102 | let mut img = None; 103 | let mut decoder = ColorCode; 104 | decoder.advance(&hm, &mut img); 105 | 106 | let img = img.unwrap(); 107 | assert_eq!(img.width(), 32); 108 | assert_eq!(img.height(), 24); 109 | let mut conf = 0; 110 | for p in img.pixels { 111 | assert_eq!(p, color_code(21, p.a() as f32 / 255f32)); 112 | assert!(conf <= p.a(), "expected monotically rising confidence/alpha"); 113 | conf = p.a(); 114 | } 115 | assert_eq!(conf, 255); 116 | } 117 | } 118 | -------------------------------------------------------------------------------- /image-ext/src/image_bgr.rs: -------------------------------------------------------------------------------- 1 | use image::{ImageBuffer, Rgb}; 2 | 3 | /// For Bgr ImageBuffer. 4 | /// 5 | /// The native BGR PixelFormat was removed in 0.24: 6 | /// 7 | #[repr(C)] 8 | #[derive(Copy, Clone)] 9 | pub struct Bgr(pub [u8; 3]); 10 | 11 | pub type BgrImage = ImageBuffer>; 12 | 13 | impl image::Pixel for Bgr { 14 | type Subpixel = u8; 15 | 16 | const CHANNEL_COUNT: u8 = 3; 17 | const COLOR_MODEL: &'static str = "BGR"; 18 | 19 | fn channels(&self) -> &[Self::Subpixel] { 20 | &self.0 21 | } 22 | 23 | fn channels_mut(&mut self) -> &mut [Self::Subpixel] { 24 | &mut self.0 25 | } 26 | 27 | fn channels4(&self) -> (Self::Subpixel, Self::Subpixel, Self::Subpixel, Self::Subpixel) { 28 | let mut channels = [Self::Subpixel::MAX; 4]; 29 | channels[0..3].copy_from_slice(&self.0); 30 | (channels[0], channels[1], channels[2], channels[3]) 31 | } 32 | 33 | fn from_channels( 34 | a: Self::Subpixel, 35 | b: Self::Subpixel, 36 | c: Self::Subpixel, 37 | _d: Self::Subpixel, 38 | ) -> Self { 39 | Self([a, b, c]) 40 | } 41 | 42 | fn from_slice(slice: &[Self::Subpixel]) -> &Self { 43 | let slice3 = slice.get(..3).unwrap(); 44 | // overheard that RangeTo is const on nightly... 45 | unsafe { &*(slice3.as_ptr() as *const Bgr) } 46 | } 47 | 48 | fn from_slice_mut(slice: &mut [Self::Subpixel]) -> &mut Self { 49 | let slice3 = slice.get_mut(..3).unwrap(); 50 | // overheard that RangeTo is const on nightly... 51 | unsafe { &mut *(slice3.as_mut_ptr() as *mut Bgr) } 52 | } 53 | 54 | fn to_rgb(&self) -> Rgb { 55 | #[allow(deprecated)] 56 | Rgb::from_channels(self.0[2], self.0[1], self.0[0], 0) 57 | } 58 | 59 | fn to_rgba(&self) -> image::Rgba { 60 | #[allow(deprecated)] 61 | image::Rgba::from_channels(self.0[2], self.0[1], self.0[0], Self::Subpixel::MAX) 62 | } 63 | 64 | fn to_luma(&self) -> image::Luma { 65 | todo!() 66 | } 67 | 68 | fn to_luma_alpha(&self) -> image::LumaA { 69 | todo!() 70 | } 71 | 72 | fn map(&self, f: F) -> Self 73 | where 74 | F: FnMut(Self::Subpixel) -> Self::Subpixel, 75 | { 76 | Self(self.0.map(f)) 77 | } 78 | 79 | fn apply(&mut self, mut f: F) 80 | where 81 | F: FnMut(Self::Subpixel) -> Self::Subpixel, 82 | { 83 | for v in &mut self.0 { 84 | *v = f(*v) 85 | } 86 | } 87 | 88 | fn map_with_alpha(&self, f: F, _g: G) -> Self 89 | where 90 | F: FnMut(Self::Subpixel) -> Self::Subpixel, 91 | G: FnMut(Self::Subpixel) -> Self::Subpixel, 92 | { 93 | self.map(f) 94 | } 95 | 96 | fn apply_with_alpha(&mut self, f: F, _g: G) 97 | where 98 | F: FnMut(Self::Subpixel) -> Self::Subpixel, 99 | G: FnMut(Self::Subpixel) -> Self::Subpixel, 100 | { 101 | self.apply(f) 102 | } 103 | 104 | fn map2(&self, other: &Self, mut f: F) -> Self 105 | where 106 | F: FnMut(Self::Subpixel, Self::Subpixel) -> Self::Subpixel, 107 | { 108 | let pixels = [f(self.0[0], other.0[0]), f(self.0[1], other.0[1]), f(self.0[1], other.0[1])]; 109 | Self(pixels) 110 | } 111 | 112 | fn apply2(&mut self, other: &Self, mut f: F) 113 | where 114 | F: FnMut(Self::Subpixel, Self::Subpixel) -> Self::Subpixel, 115 | { 116 | self.0[0] = f(self.0[0], other.0[0]); 117 | self.0[1] = f(self.0[1], other.0[1]); 118 | self.0[2] = f(self.0[1], other.0[2]); 119 | } 120 | 121 | fn invert(&mut self) { 122 | self.0[0] = Self::Subpixel::MAX - self.0[0]; 123 | self.0[1] = Self::Subpixel::MAX - self.0[1]; 124 | self.0[2] = Self::Subpixel::MAX - self.0[2]; 125 | } 126 | 127 | #[allow(unused_variables)] 128 | fn blend(&mut self, other: &Self) { 129 | todo!() 130 | } 131 | } 132 | -------------------------------------------------------------------------------- /infur/src/main.rs: -------------------------------------------------------------------------------- 1 | mod app; 2 | mod decode_predict; 3 | mod gui; 4 | mod predict_onnx; 5 | mod processing; 6 | 7 | use std::sync::mpsc::{Receiver, SyncSender, TryRecvError}; 8 | 9 | use app::{AppCmd, ProcessingApp, Processor}; 10 | use gui::{CtrlResult, FrameResult}; 11 | use stable_eyre::eyre::{eyre, Report}; 12 | use tracing::{debug, warn}; 13 | use tracing_subscriber::{fmt, EnvFilter}; 14 | 15 | /// Result with user facing error 16 | type Result = std::result::Result; 17 | 18 | fn init_logs() -> Result<()> { 19 | stable_eyre::install()?; 20 | let format = fmt::format().with_thread_names(true).with_target(false).compact(); 21 | let filter = EnvFilter::try_from_default_env().or_else(|_| EnvFilter::try_new("info")).unwrap(); 22 | tracing_subscriber::fmt::fmt().with_env_filter(filter).event_format(format).init(); 23 | Ok(()) 24 | } 25 | 26 | /// Channel events from and processing results to GUI 27 | fn proc_loop( 28 | ctrl_rx: Receiver, 29 | frame_tx: SyncSender, 30 | app_tx: SyncSender, 31 | ) -> Result<()> { 32 | fn send_app_info(app: &ProcessingApp, app_tx: &SyncSender) { 33 | let app_info = app.info(); 34 | debug!("sending updated app info {:?}", &app_info); 35 | let _ = app_tx.send(Ok(app_info)); 36 | } 37 | 38 | // instantiate app in processing thread, 39 | // since ort session can't be moved/sent 40 | let mut app = ProcessingApp::default(); 41 | 42 | loop { 43 | // todo: exit on closed channel? 44 | let mut state_change = false; 45 | loop { 46 | let cmd = if !app.is_dirty() { 47 | // video is not playing, block 48 | debug!("blocking on new command"); 49 | if state_change { 50 | send_app_info(&app, &app_tx); 51 | state_change = false; 52 | }; 53 | match ctrl_rx.recv() { 54 | Ok(c) => Some(c), 55 | // unfixable (hung-up) 56 | Err(e) => return Err(eyre!(e)), 57 | } 58 | } else { 59 | // video is playing, don't block 60 | match ctrl_rx.try_recv() { 61 | Ok(c) => Some(c), 62 | Err(TryRecvError::Empty) => break, 63 | // unfixable (hung-up) 64 | Err(e) => return Err(eyre!(e)), 65 | } 66 | }; 67 | if let Some(cmd) = cmd { 68 | debug!("relaying command: {:?}", cmd); 69 | if let Err(e) = app.control(cmd) { 70 | // Control Error 71 | let _ = app_tx.send(Err(e)); 72 | } else { 73 | state_change = true; 74 | } 75 | }; 76 | if app.to_exit { 77 | return Ok(()); 78 | }; 79 | } 80 | 81 | if state_change { 82 | send_app_info(&app, &app_tx); 83 | } 84 | 85 | match app.generate() { 86 | Ok(Some(frame)) => { 87 | // block on GUI backpressure 88 | let _ = frame_tx.send(Ok(frame)); 89 | } 90 | Ok(None) => { 91 | warn!("Nothing to process yet") 92 | } 93 | 94 | Err(e) => { 95 | let _ = frame_tx.send(Err(e)); 96 | } 97 | }; 98 | } 99 | } 100 | 101 | fn main() -> Result<()> { 102 | init_logs()?; 103 | let args = std::env::args().skip(1).collect::>(); 104 | 105 | let (frame_tx, frame_rx) = std::sync::mpsc::sync_channel(2); 106 | let (ctrl_tx, ctrl_rx) = std::sync::mpsc::channel(); 107 | let (ctrl_result_tx, ctrl_result_rx) = std::sync::mpsc::sync_channel(2); 108 | 109 | debug!("spawning Proc thread"); 110 | let infur_thread = std::thread::Builder::new() 111 | .name("Proc".to_string()) 112 | .spawn(move || proc_loop(ctrl_rx, frame_tx, ctrl_result_tx))?; 113 | 114 | debug!("starting InFur GUI"); 115 | let window_opts = eframe::NativeOptions { vsync: true, ..Default::default() }; 116 | let ctrl_tx_gui = ctrl_tx; 117 | eframe::run_native( 118 | "InFur", 119 | window_opts, 120 | Box::new(|cc| { 121 | let mut config = match cc.storage { 122 | #[cfg(feature = "persistence")] 123 | Some(storage) => eframe::get_value(storage, eframe::APP_KEY).unwrap_or_default(), 124 | _ => gui::ProcConfig::default(), 125 | }; 126 | // still override video from args 127 | if !args.is_empty() { 128 | config.video_input = args; 129 | } 130 | let app_gui = gui::InFur::new(config, ctrl_tx_gui, frame_rx, ctrl_result_rx); 131 | Box::new(app_gui) 132 | }), 133 | ); 134 | 135 | // ensure exit code 136 | infur_thread.join().unwrap().unwrap(); 137 | Ok(()) 138 | } 139 | -------------------------------------------------------------------------------- /infur/src/app.rs: -------------------------------------------------------------------------------- 1 | /// Example Application 2 | use eframe::epaint::ColorImage; 3 | use ff_video::{FFVideoError, VideoProcError}; 4 | use image_ext::Pixel; 5 | use onnxruntime::ndarray::Array3; 6 | use thiserror::Error; 7 | 8 | use crate::{ 9 | decode_predict::ColorCode, 10 | predict_onnx::{Model, ModelCmd, ModelCmdError, ModelInfo, ModelProcError}, 11 | processing::{Frame, Scale, ScaleProcError, ValidScaleError, VideoCmd, VideoPlayer}, 12 | }; 13 | 14 | pub(crate) use crate::processing::Processor; 15 | 16 | /// Application processing error 17 | #[derive(Error, Debug)] 18 | pub(crate) enum AppProcError { 19 | #[error(transparent)] 20 | Video(#[from] VideoProcError), 21 | #[error(transparent)] 22 | Scale(#[from] ScaleProcError), 23 | #[error(transparent)] 24 | Model(#[from] ModelProcError), 25 | } 26 | 27 | /// Application command processing error 28 | #[derive(Error, Debug)] 29 | pub(crate) enum AppCmdError { 30 | #[error(transparent)] 31 | Scale(#[from] ValidScaleError), 32 | #[error(transparent)] 33 | Video(#[from] FFVideoError), 34 | #[error(transparent)] 35 | Model(#[from] ModelCmdError), 36 | } 37 | 38 | /// Control entire application 39 | #[derive(Clone, Debug)] 40 | pub(crate) enum AppCmd { 41 | /// Control video input 42 | Video(VideoCmd), 43 | /// Control scale factor 44 | Scale(f32), 45 | /// Control loaded model, empty disables it 46 | Model(ModelCmd), 47 | /// Exit App 48 | Exit, 49 | } 50 | 51 | /// Example app 52 | #[derive(Default)] 53 | pub(crate) struct ProcessingApp<'m> { 54 | vid: VideoPlayer, 55 | scale: Scale, 56 | frame: Option, 57 | scaled_frame: Option, 58 | model: Model<'m>, 59 | decoder: ColorCode, 60 | decoded_img: Option, 61 | pub(crate) to_exit: bool, 62 | } 63 | 64 | /// Frame transmitted to GUI 65 | pub(crate) struct GUIFrame { 66 | pub(crate) id: u64, 67 | pub(crate) buffer: ColorImage, 68 | pub(crate) decoded_buffer: Option, 69 | } 70 | 71 | #[derive(Clone, Debug)] 72 | /// Information on current state of app 73 | pub(crate) struct AppInfo { 74 | pub(crate) model_info: Option, 75 | } 76 | 77 | impl ProcessingApp<'_> { 78 | pub(crate) fn info(&self) -> AppInfo { 79 | let model_info = self.model.get_info().cloned(); 80 | AppInfo { model_info } 81 | } 82 | } 83 | 84 | impl Processor for ProcessingApp<'_> { 85 | type Command = AppCmd; 86 | type ControlError = AppCmdError; 87 | type Input = (); 88 | type Output = (); 89 | type ProcessResult = Result, AppProcError>; 90 | 91 | fn control(&mut self, cmd: Self::Command) -> Result<&mut Self, Self::ControlError> { 92 | match cmd { 93 | AppCmd::Video(cmd) => { 94 | self.vid.control(cmd)?; 95 | } 96 | AppCmd::Scale(cmd) => { 97 | self.scale.control(cmd)?; 98 | } 99 | AppCmd::Exit => self.to_exit = true, 100 | AppCmd::Model(cmd) => { 101 | self.model.control(cmd)?; 102 | } 103 | }; 104 | Ok(self) 105 | } 106 | 107 | fn advance(&mut self, input: &(), _out: &mut ()) -> Self::ProcessResult { 108 | self.vid.advance(input, &mut self.frame)?; 109 | if self.is_dirty() { 110 | self.scale.advance(&self.frame, &mut self.scaled_frame)?; 111 | }; 112 | if let Some(scaled_frame) = &self.scaled_frame { 113 | let mut out = vec![]; 114 | self.model.advance(&scaled_frame.img, &mut out)?; 115 | if !out.is_empty() { 116 | let out = &out[0]; 117 | let shape = out.shape(); 118 | let shape = [shape[0], shape[1], shape[2]]; 119 | // todo: find way to not clone 120 | let hm: Array3 = 121 | Array3::from_shape_vec(shape, out.clone().into_raw_vec()).unwrap(); 122 | 123 | self.decoder.advance(&hm, &mut self.decoded_img); 124 | 125 | // todo: send input image and color coded img 126 | //return Ok(Some(GUIFrame { id: scaled_frame.id, buffer: img.unwrap() })); 127 | } else { 128 | self.decoded_img = None; 129 | } 130 | 131 | // todo: trait and/or processor 132 | let rgba_pixels = scaled_frame 133 | .img 134 | .pixels() 135 | .map(|p| { 136 | let cs = p.channels(); 137 | eframe::epaint::Color32::from_rgb(cs[2], cs[1], cs[0]) 138 | }) 139 | .collect::>(); 140 | 141 | let col_img = ColorImage { 142 | size: [scaled_frame.img.width() as usize, scaled_frame.img.height() as usize], 143 | pixels: rgba_pixels, 144 | }; 145 | Ok(Some(GUIFrame { 146 | id: scaled_frame.id, 147 | buffer: col_img, 148 | decoded_buffer: self.decoded_img.clone(), 149 | })) 150 | } else { 151 | Ok(None) 152 | } 153 | } 154 | 155 | fn is_dirty(&self) -> bool { 156 | self.vid.is_dirty() || self.scale.is_dirty() 157 | } 158 | } 159 | 160 | #[cfg(test)] 161 | mod test { 162 | use super::*; 163 | use infur_test_gen::{long_small_video, short_large_video}; 164 | 165 | /// 640x480 166 | fn short_large_input() -> Vec { 167 | vec![short_large_video().to_string_lossy().to_string()] 168 | } 169 | /// 1280x720 170 | fn long_small_input() -> Vec { 171 | vec![long_small_video().to_string_lossy().to_string()] 172 | } 173 | 174 | #[test] 175 | fn void() { 176 | let mut app = ProcessingApp::default(); 177 | assert!(app.generate().unwrap().is_none()); 178 | assert!(app.generate().unwrap().is_none()); 179 | } 180 | 181 | #[test] 182 | fn scale() { 183 | let mut app = ProcessingApp::default(); 184 | app.control(AppCmd::Video(VideoCmd::Play(short_large_input()))).unwrap(); 185 | app.control(AppCmd::Scale(0.5)).unwrap(); 186 | let f2 = app.generate().unwrap().expect("video should already play"); 187 | assert_eq!(f2.buffer.size, [1280 / 2, 720 / 2]); 188 | } 189 | 190 | #[test] 191 | fn switch_scale() { 192 | let mut app = ProcessingApp::default(); 193 | app.control(AppCmd::Video(VideoCmd::Play(long_small_input()))).unwrap(); 194 | let f1 = app.generate().unwrap().expect("video should already play"); 195 | assert_eq!(f1.buffer.size, [640, 480]); 196 | 197 | app.control(AppCmd::Scale(0.5)).unwrap(); 198 | let f2 = app.generate().unwrap().expect("video should keep playing"); 199 | assert_eq!(f2.buffer.size, [640 / 2, 480 / 2]); 200 | } 201 | 202 | #[test] 203 | fn switch_video_then_scale() { 204 | let mut app = ProcessingApp::default(); 205 | // in 640x480 out same 206 | app.control(AppCmd::Video(VideoCmd::Play(long_small_input()))).unwrap(); 207 | let f1 = app.generate().unwrap().unwrap(); 208 | assert_eq!(f1.buffer.size, [640, 480]); 209 | // in 1280x720 out same 210 | app.control(AppCmd::Video(VideoCmd::Play(short_large_input()))).unwrap(); 211 | let f2 = app.generate().unwrap().unwrap(); 212 | assert_eq!(f2.buffer.size, [1280, 720]); 213 | // in 1280x720 out twice 214 | app.control(AppCmd::Scale(2.0)).unwrap(); 215 | let f3 = app.generate().unwrap().unwrap(); 216 | assert_eq!(f3.buffer.size, [1280 * 2, 720 * 2]); 217 | } 218 | 219 | #[test] 220 | fn scaled_frame_after_stopped_video() { 221 | let mut app = ProcessingApp::default(); 222 | app.control(AppCmd::Video(VideoCmd::Play(short_large_input()))).unwrap(); 223 | let f1 = app.generate().unwrap().unwrap(); 224 | assert_eq!(f1.buffer.size, [1280, 720]); 225 | app.control(AppCmd::Video(VideoCmd::Stop)).unwrap(); 226 | let f2 = app.generate().unwrap().unwrap(); 227 | assert_eq!(f1.id, f2.id); 228 | assert!(!app.is_dirty()); 229 | 230 | app.control(AppCmd::Scale(0.5)).unwrap(); 231 | assert!(app.is_dirty()); 232 | let f3 = app.generate().unwrap().unwrap(); 233 | assert_eq!(f2.id, f3.id); 234 | assert_eq!(f3.buffer.size, [1280 / 2, 720 / 2]); 235 | } 236 | 237 | #[test] 238 | fn pause_video() { 239 | let mut app = ProcessingApp::default(); 240 | app.control(AppCmd::Video(VideoCmd::Play(long_small_input()))).unwrap(); 241 | let f1 = app.generate().unwrap().unwrap(); 242 | app.control(AppCmd::Video(VideoCmd::Pause(true))).unwrap(); 243 | assert!(!app.is_dirty()); 244 | let f2 = app.generate().unwrap().unwrap(); 245 | assert_eq!(f1.id, f2.id); 246 | assert!(!app.is_dirty()); 247 | 248 | app.control(AppCmd::Video(VideoCmd::Pause(false))).unwrap(); 249 | assert!(app.is_dirty()); 250 | let f3 = app.generate().unwrap().unwrap(); 251 | assert_ne!(f2.id, f3.id); 252 | } 253 | } 254 | -------------------------------------------------------------------------------- /ff-video/src/decoder.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | io::{ErrorKind, Read, Write}, 3 | process::Command, 4 | sync::mpsc::{Receiver, RecvTimeoutError}, 5 | thread::{self, JoinHandle}, 6 | time::Duration, 7 | }; 8 | 9 | use image_ext::{BgrImage, ImageBuffer}; 10 | use tracing::{debug, error, info, warn}; 11 | 12 | use crate::{error::VideoResult, parse::FFMpegLineIter}; 13 | use crate::{ 14 | error::{InfoResult, VideoProcError}, 15 | parse::{InfoParser, Stream, StreamInfo, VideoInfo}, 16 | }; 17 | 18 | pub struct FFMpegDecoderBuilder { 19 | cmd: Command, 20 | } 21 | 22 | pub struct FFMpegDecoder { 23 | child: std::process::Child, 24 | stdout: std::process::ChildStdout, 25 | info_thread: JoinHandle, 26 | pub frame_counter: u64, 27 | pub video_output: Stream, 28 | } 29 | 30 | impl Default for FFMpegDecoderBuilder { 31 | fn default() -> Self { 32 | let mut cmd = Command::new("ffmpeg"); 33 | // options 34 | cmd.arg("-hide_banner"); 35 | // escape input 36 | cmd.arg("-i"); 37 | Self { cmd } 38 | } 39 | } 40 | 41 | impl FFMpegDecoderBuilder { 42 | pub fn input(mut self, input: I) -> Self 43 | where 44 | I: IntoIterator, 45 | S: AsRef, 46 | { 47 | self.cmd.args(input); 48 | self 49 | } 50 | 51 | fn cmd(mut self) -> Command { 52 | // output 53 | self.cmd.args([ 54 | "-an", 55 | "-f", 56 | "image2pipe", 57 | "-fflags", 58 | "nobuffer", 59 | "-pix_fmt", 60 | "bgr24", 61 | "-c:v", 62 | "rawvideo", 63 | "pipe:1", 64 | ]); 65 | // piping 66 | use std::process::Stdio; 67 | self.cmd.stderr(Stdio::piped()).stdout(Stdio::piped()).stdin(Stdio::piped()); 68 | // todo: rm 69 | eprintln!( 70 | "ffmpeg {}", 71 | self.cmd.get_args().map(|s| s.to_string_lossy()).collect::>().join(" ") 72 | ); 73 | self.cmd 74 | } 75 | } 76 | 77 | impl FFMpegDecoder { 78 | pub fn try_new(builder: FFMpegDecoderBuilder) -> VideoResult { 79 | let mut cmd = builder.cmd(); 80 | let mut child = cmd 81 | .spawn() 82 | .map_err(|e| VideoProcError::explain_io("couldn't spawn video process", e))?; 83 | let stderr = 84 | child.stderr.take().ok_or_else(|| VideoProcError::is_missing("stderr pipe"))?; 85 | let (stream_info_rx, info_thread) = spawn_info_thread(stderr)?; 86 | 87 | // determine output 88 | let mut final_line = None; 89 | let video_output = loop { 90 | let msg = match stream_info_rx.recv_timeout(Duration::from_secs(10)) { 91 | Ok(msg) => msg, 92 | Err(e) => { 93 | let why = match e { 94 | RecvTimeoutError::Timeout => "timeout", 95 | RecvTimeoutError::Disconnected => "disconnected", 96 | }; 97 | 98 | let explanation = match final_line { 99 | None => why.to_string(), 100 | Some(line) => format!("{why} - {line}"), 101 | }; 102 | return Err(VideoProcError::Start(explanation)); 103 | } 104 | }; 105 | if let Ok(StreamInfoTerm::Info(StreamInfo::Output { stream, .. })) = msg { 106 | break stream; 107 | }; 108 | if let Ok(StreamInfoTerm::Final(line)) = msg { 109 | final_line = Some(line); 110 | } 111 | }; 112 | 113 | let stdout = 114 | child.stdout.take().ok_or_else(|| VideoProcError::is_missing("stdout pipe"))?; 115 | Ok(Self { child, stdout, info_thread, video_output, frame_counter: 0 }) 116 | } 117 | 118 | /// stop process gracefully and await exit code 119 | pub fn close(mut self) -> VideoResult<()> { 120 | let mut stdin = 121 | self.child.stdin.take().ok_or_else(|| VideoProcError::is_missing("stdin pipe"))?; 122 | // to close, we first send quit message (stdin is available since we encode nothing from it) 123 | // any stdin buffer is flushed in the drop() of wait() 124 | // thus it should break output pipe since ffmpeg has already hung up 125 | match stdin.write_all(b"q") { 126 | Ok(_) => {} 127 | Err(e) if e.kind() == ErrorKind::BrokenPipe => {} // process probably already exited 128 | Err(e) => { 129 | return Err(VideoProcError::explain_io("couldn't send q to process", e)); 130 | } 131 | }; 132 | // ... unless we don't drain stdout as well, which we do here 133 | self.stdout.bytes().for_each(|_| {}); 134 | 135 | let exit_code = self 136 | .child 137 | .wait() 138 | .map_err(|e| VideoProcError::explain_io("waiting on video process", e))?; 139 | _ = self 140 | .info_thread 141 | .join() 142 | .map_err(|_| VideoProcError::Other("error joining meta data thread".to_string()))?; 143 | match exit_code.code() { 144 | Some(c) if c > 0 => Err(VideoProcError::ExitCode(c)), 145 | None => Err(VideoProcError::Other("video child process killed by signal".to_string())), 146 | _ => Ok(()), 147 | } 148 | } 149 | 150 | pub fn empty_image(&self) -> BgrImage { 151 | let (width, height) = (self.video_output.width, self.video_output.height); 152 | ImageBuffer::new(width, height) 153 | } 154 | 155 | /// Write new image and return its frame id. 156 | pub fn read_frame(&mut self, image: &mut BgrImage) -> VideoResult { 157 | self.stdout.read_exact(image.as_mut()).map_err(|e| match self.child.try_wait() { 158 | Ok(Some(status)) if status.code() == Some(0) => { 159 | VideoProcError::FinishedNormally { source: e } 160 | } 161 | _ => VideoProcError::ExactReadError { source: e }, 162 | })?; 163 | self.frame_counter += 1; 164 | Ok(self.frame_counter) 165 | } 166 | } 167 | 168 | /// StreamInfo or variant to signal EOF 169 | enum StreamInfoTerm { 170 | /// Video IO stream infos 171 | Info(StreamInfo), 172 | /// Last line read without parsing after which 173 | /// no more messages will be sent 174 | Final(String), 175 | } 176 | 177 | /// Deliver infos about an ffmpeg video process through its stderr file 178 | /// 179 | /// The receiver can be read until satisfying info was obtained and dropped anytime. 180 | /// By default, frame updates and other infos are logged as tracing event. 181 | /// The last line is returned if the thread joins without errors. 182 | /// 183 | /// todo: offer a custom callback for info messages 184 | fn spawn_info_thread( 185 | stderr: R, 186 | ) -> VideoResult<(Receiver>, JoinHandle)> 187 | where 188 | R: Read + Send + 'static, 189 | { 190 | let (stream_info_tx, stream_info_rx) = 191 | std::sync::mpsc::sync_channel::>(2); 192 | 193 | let info_thread = thread::Builder::new() 194 | .name("Video".to_string()) 195 | .spawn(move || { 196 | let reader = std::io::BufReader::new(stderr); 197 | let mut ffmpeg_lines = reader.bytes().ffmpeg_lines(); 198 | let lines = ffmpeg_lines.by_ref().filter_map(|r| match r { 199 | Err(e) => { 200 | error!("couldn't read stderr {:?}", e); 201 | None 202 | } 203 | Ok(r) => Some(r), 204 | }); 205 | //.inspect(|l| println!("!!{}", l)); 206 | 207 | // Delivery semantics depend on the message type: 208 | // - Stream and Error: must be delivered until the recv hangs up (thus isn't interested anymore) 209 | // - Any info: logged 210 | for msg in InfoParser::default().iter_on(lines) { 211 | match msg { 212 | Ok(VideoInfo::Stream(msg)) => { 213 | _ = stream_info_tx 214 | .send(Ok(StreamInfoTerm::Info(msg.clone()))) 215 | .map_err(|e| warn!("could not send stream info: {:?}", e)); 216 | log_info_handler(Ok(VideoInfo::Stream(msg))); 217 | } 218 | Ok(msg) => { 219 | log_info_handler(Ok(msg)); 220 | } 221 | Err(e) => { 222 | _ = stream_info_tx.send(Err(e.clone())); 223 | log_info_handler(Err(e)); 224 | } 225 | }; 226 | } 227 | let last_line = String::from_utf8_lossy(ffmpeg_lines.state()).to_string(); 228 | info!("finished reading stderr: {}", &last_line); 229 | _ = stream_info_tx.send(Ok(StreamInfoTerm::Final(last_line.clone()))); 230 | last_line 231 | }) 232 | .map_err(|e| VideoProcError::explain_io("couldn't spawn info parsing thread", e))?; 233 | Ok((stream_info_rx, info_thread)) 234 | } 235 | 236 | fn log_info_handler(msg: crate::parse::Result) { 237 | match msg { 238 | Ok(msg) => match msg { 239 | VideoInfo::Stream(msg) => { 240 | info!("new stream info: {:?}", msg); 241 | } 242 | VideoInfo::Frame(msg) => { 243 | debug!("frame: {:?}", msg); 244 | } 245 | VideoInfo::Codec(msg) => { 246 | info!("codec: {}", msg); 247 | } 248 | }, 249 | Err(e) => { 250 | error!("parsing video update: {:?}", e); 251 | } 252 | } 253 | } 254 | -------------------------------------------------------------------------------- /infur/src/processing.rs: -------------------------------------------------------------------------------- 1 | use std::{error::Error as StdError, fmt::Display, num::NonZeroU32, ops::Deref}; 2 | 3 | use fast_image_resize as fr; 4 | use ff_video::{FFMpegDecoder, FFMpegDecoderBuilder, FFVideoError, VideoProcError, VideoResult}; 5 | use image_ext::BgrImage; 6 | use thiserror::Error; 7 | 8 | /// Frame produced and processed 9 | pub(crate) struct Frame { 10 | pub(crate) id: u64, 11 | pub(crate) img: BgrImage, 12 | } 13 | 14 | impl PartialEq for Frame { 15 | fn eq(&self, other: &Self) -> bool { 16 | self.id == other.id 17 | } 18 | } 19 | 20 | /// Transform expensive to clone data 21 | /// 22 | /// Other parameters affecting output and results are controlled by a single message type. 23 | pub(crate) trait Processor { 24 | /// Type of commands to control the processor 25 | type Command; 26 | 27 | /// Error to expect when controlling the processor 28 | type ControlError; 29 | 30 | /// Input to advance with 31 | type Input; 32 | 33 | /// Item mutated on advance 34 | /// 35 | /// For why we can't impl Iterator see: 36 | /// 37 | /// > 38 | type Output; 39 | 40 | /// Return value of processing input 41 | type ProcessResult; 42 | 43 | /// Affect processor parameters 44 | fn control(&mut self, cmd: Self::Command) -> Result<&mut Self, Self::ControlError>; 45 | 46 | /// Process and store a new result 47 | fn advance(&mut self, inp: &Self::Input, out: &mut Self::Output) -> Self::ProcessResult; 48 | 49 | /// True if passing the same input into advance writes new results 50 | fn is_dirty(&self) -> bool; 51 | 52 | /// Generate results for default input/output (usually final processing nodes) 53 | fn generate(&mut self) -> ::ProcessResult 54 | where 55 | Self::Input: Default, 56 | Self::Output: Default, 57 | { 58 | self.advance(&Self::Input::default(), &mut Self::Output::default()) 59 | } 60 | } 61 | 62 | /// Commands that control VideoPlayer 63 | #[derive(Debug, Clone)] 64 | pub(crate) enum VideoCmd { 65 | /// Start or restart playing video from this ffmpeg input 66 | Play(Vec), 67 | /// Pause generating new frames 68 | Pause(bool), 69 | /// Stop whenever 70 | Stop, 71 | } 72 | 73 | /// Writes video frames at command 74 | #[derive(Default)] 75 | pub(crate) struct VideoPlayer { 76 | vid: Option, 77 | input: Vec, 78 | paused: bool, 79 | } 80 | 81 | impl VideoPlayer { 82 | fn close_video(&mut self) -> VideoResult<()> { 83 | self.vid.take().map_or(Ok(()), |vid| vid.close()) 84 | } 85 | } 86 | 87 | impl Processor for VideoPlayer { 88 | type Command = VideoCmd; 89 | type ControlError = FFVideoError; 90 | type Input = (); 91 | type Output = Option; 92 | type ProcessResult = VideoResult<()>; 93 | 94 | fn control(&mut self, cmd: Self::Command) -> Result<&mut Self, Self::ControlError> { 95 | match cmd { 96 | Self::Command::Play(input) => { 97 | self.close_video()?; 98 | self.input = input; 99 | let builder = FFMpegDecoderBuilder::default().input(self.input.clone()); 100 | self.vid = Some(FFMpegDecoder::try_new(builder)?); 101 | } 102 | Self::Command::Pause(paused) => { 103 | self.paused = paused; 104 | } 105 | Self::Command::Stop => { 106 | self.close_video()?; 107 | } 108 | } 109 | Ok(self) 110 | } 111 | 112 | fn is_dirty(&self) -> bool { 113 | !self.paused && self.vid.is_some() 114 | } 115 | 116 | fn advance(&mut self, _inp: &(), out: &mut Self::Output) -> Self::ProcessResult { 117 | if self.paused { 118 | return Ok(()); 119 | } 120 | if let Some(vid) = self.vid.as_mut() { 121 | // either reuse, re-create (on size change) or create a frame with suitable buffer 122 | let frame = if let Some(ref mut frame) = out { 123 | if frame.img.width() != vid.video_output.width 124 | || frame.img.height() != vid.video_output.height 125 | { 126 | frame.img = vid.empty_image(); 127 | } 128 | frame 129 | } else { 130 | out.get_or_insert_with(|| Frame { id: 0, img: vid.empty_image() }) 131 | }; 132 | let id = vid.read_frame(&mut frame.img); 133 | if let Err(VideoProcError::FinishedNormally { .. }) = id { 134 | self.close_video()?; 135 | } 136 | frame.id = id?; 137 | }; 138 | Ok(()) 139 | } 140 | } 141 | 142 | #[derive(PartialEq, Debug)] 143 | pub(crate) struct ValidScale(f32); 144 | 145 | #[derive(Debug, Clone)] 146 | pub(crate) struct ValidScaleError { 147 | msg: &'static str, 148 | } 149 | 150 | impl Display for ValidScaleError { 151 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 152 | f.write_str(self.msg) 153 | } 154 | } 155 | 156 | impl StdError for ValidScaleError {} 157 | 158 | impl TryFrom for ValidScale { 159 | type Error = ValidScaleError; 160 | 161 | fn try_from(value: f32) -> Result { 162 | if value <= 0.0f32 { 163 | Err(ValidScaleError { msg: "Cannot scale by negative number" }) 164 | } else { 165 | Ok(Self(value)) 166 | } 167 | } 168 | } 169 | 170 | impl Deref for ValidScale { 171 | type Target = f32; 172 | 173 | fn deref(&self) -> &Self::Target { 174 | &self.0 175 | } 176 | } 177 | 178 | /// Scale frames by a constant factor 179 | pub(crate) struct Scale { 180 | factor: ValidScale, 181 | resizer: fr::Resizer, 182 | dirty: bool, 183 | } 184 | 185 | impl Default for Scale { 186 | fn default() -> Self { 187 | Self { 188 | factor: ValidScale(1.0f32), 189 | resizer: fr::Resizer::new(fr::ResizeAlg::Nearest), 190 | dirty: true, 191 | } 192 | } 193 | } 194 | 195 | impl Scale { 196 | fn is_unit_scale(&self) -> bool { 197 | self.factor.0 == 1.0f32 198 | } 199 | } 200 | /// Error processing scale 201 | #[derive(Error, Debug)] 202 | pub(crate) enum ScaleProcError { 203 | #[error("scaling from 0-sized input")] 204 | ZeroSizeIn, 205 | #[error("scaling to 0-sized output")] 206 | ZeroSizeOut, 207 | #[error(transparent)] 208 | PixelType(#[from] fr::DifferentTypesOfPixelsError), 209 | #[error(transparent)] 210 | BufferError(#[from] fr::ImageBufferError), 211 | } 212 | 213 | impl Processor for Scale { 214 | type Command = f32; 215 | type ControlError = >::Error; 216 | type Input = Option; 217 | type Output = Option; 218 | type ProcessResult = Result<(), ScaleProcError>; 219 | 220 | fn control(&mut self, cmd: Self::Command) -> Result<&mut Self, Self::ControlError> { 221 | let factor = cmd.try_into()?; 222 | self.dirty = factor != self.factor; 223 | self.factor = factor; 224 | // todo: change resizer to bilinear for some factors? 225 | Ok(self) 226 | } 227 | 228 | fn is_dirty(&self) -> bool { 229 | self.dirty 230 | } 231 | 232 | fn advance(&mut self, input: &Self::Input, out: &mut Self::Output) -> Self::ProcessResult { 233 | self.dirty = false; 234 | let input = match input { 235 | Some(i) => i, 236 | None => return Ok(()), 237 | }; 238 | if self.is_unit_scale() { 239 | // todo: can we at all avoid a clone? 240 | *out = Some(Frame { id: input.id, img: input.img.clone() }); 241 | return Ok(()); 242 | } 243 | 244 | // todo: some conversion trait 245 | // get input view 246 | let img_view = fr::ImageView::from_buffer( 247 | NonZeroU32::new(input.img.width()).ok_or(ScaleProcError::ZeroSizeIn)?, 248 | NonZeroU32::new(input.img.height()).ok_or(ScaleProcError::ZeroSizeIn)?, 249 | input.img.as_raw(), 250 | fr::PixelType::U8x3, 251 | )?; 252 | 253 | let nwidth = (input.img.width() as f32 * self.factor.0) as _; 254 | let nheight = (input.img.height() as f32 * self.factor.0) as _; 255 | let nwidth0 = NonZeroU32::new(nwidth).ok_or(ScaleProcError::ZeroSizeOut)?; 256 | let nheight0 = NonZeroU32::new(nheight).ok_or(ScaleProcError::ZeroSizeOut)?; 257 | 258 | // todo: some conversion trait 259 | // get or create new frame 260 | let frame = if let Some(ref mut frame) = out { 261 | if frame.img.width() != nwidth || frame.img.height() != nheight { 262 | frame.img = BgrImage::new(nwidth, nheight); 263 | } 264 | frame.id = input.id; 265 | frame 266 | } else { 267 | out.get_or_insert_with(|| Frame { id: input.id, img: BgrImage::new(nwidth, nheight) }) 268 | }; 269 | 270 | // get output view 271 | let mut img_view_mut = fr::ImageViewMut::from_buffer( 272 | nwidth0, 273 | nheight0, 274 | frame.img.as_mut(), 275 | fr::PixelType::U8x3, 276 | )?; 277 | 278 | self.resizer.resize(&img_view, &mut img_view_mut)?; 279 | self.dirty = false; 280 | Ok(()) 281 | } 282 | } 283 | 284 | #[cfg(test)] 285 | mod test { 286 | use super::*; 287 | 288 | #[test] 289 | fn scale_from_size0() { 290 | let zero = Frame { id: 0, img: BgrImage::new(0, 10) }; 291 | let mut out = None; 292 | let mut scale = Scale::default(); 293 | scale.control(0.99).unwrap(); 294 | assert!(matches!(scale.advance(&Some(zero), &mut out), Err(ScaleProcError::ZeroSizeIn))); 295 | } 296 | #[test] 297 | fn scale_to_size0() { 298 | let img = Frame { id: 0, img: BgrImage::new(10, 10) }; 299 | let mut out = None; 300 | let mut scale = Scale::default(); 301 | scale.control(0.00000001).unwrap(); 302 | assert!(matches!(scale.advance(&Some(img), &mut out), Err(ScaleProcError::ZeroSizeOut))); 303 | } 304 | } 305 | -------------------------------------------------------------------------------- /infur/src/gui.rs: -------------------------------------------------------------------------------- 1 | use std::collections::VecDeque; 2 | use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, TryRecvError}; 3 | use std::time::{Duration, Instant}; 4 | 5 | use crate::app::{AppCmd, AppCmdError, AppInfo, AppProcError, GUIFrame}; 6 | use crate::predict_onnx::ModelCmd; 7 | use crate::processing::VideoCmd; 8 | use eframe::{ 9 | egui::{self, CentralPanel, RichText, SidePanel, Slider, TextureFilter, TextureHandle}, 10 | epaint::FontId, 11 | }; 12 | 13 | /// Result from processing a frame 14 | pub(crate) type FrameResult = std::result::Result; 15 | 16 | /// Result from processing commands 17 | pub(crate) type CtrlResult = std::result::Result; 18 | 19 | /// GUI textures of model in-/output 20 | pub(crate) struct TextureFrame { 21 | pub(crate) id: u64, 22 | pub(crate) handle: TextureHandle, 23 | pub(crate) decoded_handle: Option, 24 | } 25 | 26 | /// Count frames and time between set points 27 | pub(crate) struct FrameCounter { 28 | pub(crate) recvd_id: Option, 29 | pub(crate) shown_id: u64, 30 | pub(crate) since: Instant, 31 | pub(crate) elapsed_since: Duration, 32 | pub(crate) shown_since: u64, 33 | pub(crate) recvd_since: Option, 34 | } 35 | 36 | impl FrameCounter { 37 | // set start of new time strip (call per measurement) 38 | pub(crate) fn set_on(&mut self, now: Instant, shown_id: u64, recvd_id: Option) { 39 | // deltas 40 | self.shown_since = shown_id - self.shown_id; 41 | self.recvd_since = match (self.recvd_id, recvd_id) { 42 | (None, _) => None, 43 | (_, None) => None, 44 | (Some(r0), Some(r1)) if r0 > r1 => None, 45 | (Some(r0), Some(r1)) => Some(r1 - r0), 46 | }; 47 | self.elapsed_since = self.elapsed(now); 48 | // new 0 49 | self.recvd_id = recvd_id; 50 | self.shown_id = shown_id; 51 | self.since = now; 52 | } 53 | 54 | // time elapsed since last setting 55 | pub(crate) fn elapsed(&self, now: Instant) -> Duration { 56 | now - self.since 57 | } 58 | 59 | // fps of shown frames with respect to last time 60 | pub(crate) fn shown_fps(&self) -> f64 { 61 | self.shown_since as f64 / self.elapsed_since.as_secs_f64() 62 | } 63 | 64 | // fps of received frames with respect to last time 65 | pub(crate) fn recvd_fps(&self) -> f64 { 66 | match self.recvd_since { 67 | Some(r) => r as f64 / self.elapsed_since.as_secs_f64(), 68 | None => f64::NAN, 69 | } 70 | } 71 | 72 | // frames dropped (not shown) or skipped (also not shown) 73 | pub(crate) fn dropped_since(&self) -> i64 { 74 | self.recvd_since.unwrap_or_default() as i64 - self.shown_since as i64 75 | } 76 | } 77 | 78 | impl Default for FrameCounter { 79 | fn default() -> Self { 80 | Self { 81 | recvd_id: None, 82 | shown_id: 0, 83 | since: Instant::now(), 84 | elapsed_since: Duration::ZERO, 85 | shown_since: 0, 86 | recvd_since: None, 87 | } 88 | } 89 | } 90 | 91 | #[derive(serde::Deserialize, serde::Serialize)] 92 | pub(crate) struct ProcConfig { 93 | pub(crate) video_input: Vec, 94 | pub(crate) scale: f32, 95 | pub(crate) paused: bool, 96 | pub(crate) model_input: String, 97 | } 98 | 99 | impl Default for ProcConfig { 100 | fn default() -> Self { 101 | Self { video_input: vec![], scale: 0.5, paused: false, model_input: String::default() } 102 | } 103 | } 104 | 105 | #[derive(Default, Clone)] 106 | pub(crate) struct ProcStatus { 107 | pub(crate) video: String, 108 | pub(crate) scale: String, 109 | pub(crate) model: String, 110 | } 111 | 112 | pub(crate) struct InFur { 113 | pub(crate) ctrl_tx: Sender, 114 | pub(crate) frame_rx: Receiver, 115 | pub(crate) proc_result: Option, 116 | pub(crate) ctrl_rx: Receiver, 117 | pub(crate) main_texture: Option, 118 | pub(crate) config: ProcConfig, 119 | pub(crate) closing: bool, 120 | pub(crate) allow_closing: bool, 121 | pub(crate) error_history: VecDeque, 122 | pub(crate) counter: FrameCounter, 123 | pub(crate) show_count: u64, 124 | pub(crate) proc_status: ProcStatus, 125 | } 126 | 127 | impl InFur { 128 | pub(crate) fn new( 129 | config: ProcConfig, 130 | ctrl_tx: Sender, 131 | frame_rx: Receiver, 132 | ctrl_rx: Receiver, 133 | ) -> Self { 134 | let mut app = Self { 135 | ctrl_tx, 136 | frame_rx, 137 | ctrl_rx, 138 | proc_result: None, 139 | main_texture: None, 140 | config, 141 | closing: false, 142 | allow_closing: false, 143 | error_history: VecDeque::with_capacity(3), 144 | counter: FrameCounter::default(), 145 | show_count: 0, 146 | proc_status: ProcStatus::default(), 147 | }; 148 | // send initial config 149 | app.send(AppCmd::Scale(app.config.scale)); 150 | app.send(AppCmd::Video(VideoCmd::Play( 151 | app.config.video_input.iter().cloned().filter(|s| !s.is_empty()).collect(), 152 | ))); 153 | app.send(AppCmd::Video(VideoCmd::Pause(app.config.paused))); 154 | app.send(AppCmd::Model(ModelCmd::Load(app.config.model_input.clone()))); 155 | app 156 | } 157 | 158 | pub(crate) fn send(&mut self, cmd: AppCmd) { 159 | self.error_history.truncate(2); 160 | _ = self.ctrl_tx.send(cmd).map_err(|e| self.error_history.push_front(e.to_string())); 161 | } 162 | } 163 | 164 | impl eframe::App for InFur { 165 | fn update(&mut self, ctx: &eframe::egui::Context, frame_: &mut eframe::Frame) { 166 | // update texture from new frame or close if disconnected 167 | // this limits UI updates if no frames are sent to ca. 30fps 168 | let mut new_frame = false; 169 | match self.frame_rx.recv_timeout(Duration::from_millis(30)) { 170 | Ok(Ok(frame)) => { 171 | let decoded_handle = frame.decoded_buffer.map(|decoded_img| { 172 | ctx.load_texture("decoded_texture", decoded_img, TextureFilter::Linear) 173 | }); 174 | 175 | let tex = TextureFrame { 176 | id: frame.id, 177 | handle: ctx.load_texture("main_texture", frame.buffer, TextureFilter::Linear), 178 | decoded_handle, 179 | }; 180 | new_frame = true; 181 | self.main_texture = Some(tex); 182 | self.proc_result = None; 183 | } 184 | Ok(Err(e)) => { 185 | self.proc_result = Some(e); 186 | } 187 | Err(RecvTimeoutError::Timeout) => {} 188 | Err(RecvTimeoutError::Disconnected) => { 189 | self.allow_closing = true; 190 | frame_.close(); 191 | } 192 | } 193 | if self.closing && !self.allow_closing { 194 | self.send(AppCmd::Video(VideoCmd::Stop)); 195 | self.send(AppCmd::Exit); 196 | } 197 | 198 | // advance and reset counters every second 199 | self.show_count += 1; 200 | let now = std::time::Instant::now(); 201 | if self.counter.elapsed(now) > Duration::from_secs(1) { 202 | self.counter.set_on(now, self.show_count, self.main_texture.as_ref().map(|t| t.id)); 203 | } 204 | 205 | // stringify last frame's statuses 206 | match (new_frame, &self.main_texture, &self.proc_result) { 207 | (true, _, Some(AppProcError::Video(e))) => self.proc_status.video = e.to_string(), 208 | (true, Some(tex), _) => self.proc_status.video = tex.id.to_string(), 209 | _ => {} 210 | }; 211 | match &self.proc_result { 212 | Some(AppProcError::Scale(e)) => self.proc_status.scale = e.to_string(), 213 | None => self.proc_status.scale = String::default(), 214 | _ => {} 215 | } 216 | match &self.proc_result { 217 | Some(AppProcError::Model(e)) => self.proc_status.model = e.to_string(), 218 | None if self.config.model_input.is_empty() => { 219 | self.proc_status.model = String::default() 220 | } 221 | _ => {} 222 | } 223 | 224 | // stringify control errors or app infos, may override frame status 225 | match self.ctrl_rx.try_recv() { 226 | Ok(info) => match info { 227 | Ok(info) => { 228 | if let Some(model_info) = info.model_info { 229 | self.proc_status.model = format!( 230 | "Model loaded: {} -> {}", 231 | model_info.input_names.join(","), 232 | model_info.output_names.join(",") 233 | ); 234 | } 235 | } 236 | Err(AppCmdError::Video(e)) => { 237 | self.proc_status.video = e.to_string(); 238 | } 239 | Err(AppCmdError::Scale(e)) => { 240 | self.proc_status.scale = e.to_string(); 241 | } 242 | Err(AppCmdError::Model(e)) => { 243 | self.proc_status.model = e.to_string(); 244 | } 245 | }, 246 | Err(TryRecvError::Disconnected) => { 247 | self.error_history.push_front("lost processing control".to_string()); 248 | frame_.close(); 249 | } 250 | Err(_) => {} 251 | } 252 | 253 | SidePanel::left("Options").show(ctx, |ui| { 254 | ui.spacing_mut().item_spacing.y = 10.0; 255 | // video input 256 | ui.label(RichText::new("Video").font(FontId::proportional(30.0))); 257 | // (un-)pause video 258 | if ui.checkbox(&mut self.config.paused, "Pause").changed { 259 | self.send(AppCmd::Video(VideoCmd::Pause(self.config.paused))) 260 | }; 261 | // (re-)play video 262 | if self.config.video_input.is_empty() { 263 | self.config.video_input.push(String::default()); 264 | } 265 | let mut vid_input_changed = false; 266 | for inp in self.config.video_input.iter_mut() { 267 | let textbox = ui.text_edit_singleline(inp); 268 | vid_input_changed = vid_input_changed || textbox.lost_focus(); 269 | } 270 | if vid_input_changed { 271 | self.send(AppCmd::Video(VideoCmd::Play( 272 | self.config.video_input.iter().cloned().filter(|s| !s.is_empty()).collect(), 273 | ))); 274 | } 275 | ui.label(&self.proc_status.video); 276 | 277 | ui.label(RichText::new("Inference").font(FontId::proportional(30.0))); 278 | let scale = Slider::new(&mut self.config.scale, 0.1f32..=1.0) 279 | .step_by(0.01f64) 280 | .text("scale") 281 | .clamp_to_range(true); 282 | let scale_response = ui.add(scale); 283 | if scale_response.changed { 284 | self.send(AppCmd::Scale(self.config.scale)); 285 | }; 286 | if !self.proc_status.scale.is_empty() { 287 | ui.label(&self.proc_status.model); 288 | } 289 | 290 | // (re-)load model 291 | let model_input = ui.text_edit_singleline(&mut self.config.model_input); 292 | if model_input.lost_focus() { 293 | self.send(AppCmd::Model(ModelCmd::Load(self.config.model_input.clone()))); 294 | } 295 | ui.label(&self.proc_status.model); 296 | 297 | // frame stats 298 | ui.label(RichText::new("Stats").font(FontId::proportional(30.0))); 299 | let frame_stats = format!( 300 | "fps UI: {:>3.1}\nprocessed: {:>3.1}\ndrops/skips: {}", 301 | self.counter.shown_fps(), 302 | self.counter.recvd_fps(), 303 | self.counter.dropped_since() 304 | ); 305 | ui.label(frame_stats); 306 | 307 | // rather fatal errors or final messages 308 | ui.with_layout(egui::Layout::bottom_up(egui::Align::LEFT), |ui| { 309 | for (i, err) in self.error_history.iter().cloned().enumerate() { 310 | let col = egui::Color32::RED.linear_multiply(1.0 - (i as f32 / 4.0)); 311 | ui.colored_label(col, err); 312 | } 313 | }); 314 | }); 315 | 316 | // show last_texture 317 | if let Some(tex_frame) = &self.main_texture { 318 | CentralPanel::default().show(ctx, |ui| { 319 | // occupy max width with constant aspect ratio 320 | let max_width = ui.available_width(); 321 | let [w, h] = tex_frame.handle.size(); 322 | let w_scale = max_width / w as f32; 323 | let (w, h) = (w as f32 * w_scale, h as f32 * w_scale); 324 | ui.image(&tex_frame.handle, [w, h]); 325 | // prop decoded image underneath 326 | // todo: blend somehow? 327 | if let Some(ref handle) = tex_frame.decoded_handle { 328 | ui.image(handle, [w, h]); 329 | }; 330 | }); 331 | }; 332 | 333 | ctx.request_repaint(); 334 | } 335 | 336 | fn on_close_event(&mut self) -> bool { 337 | self.error_history.push_front("exiting...".to_string()); 338 | // send exit once 339 | if !self.closing { 340 | // we could not close the video to exit faster, but 341 | // would end with an ffmpeg error 342 | self.send(AppCmd::Video(VideoCmd::Stop)); 343 | self.send(AppCmd::Exit); 344 | } 345 | self.closing = true; 346 | self.allow_closing 347 | } 348 | 349 | #[cfg(feature = "persistence")] 350 | fn save(&mut self, storage: &mut dyn eframe::Storage) { 351 | eframe::set_value(storage, eframe::APP_KEY, &self.config); 352 | } 353 | } 354 | -------------------------------------------------------------------------------- /infur/src/predict_onnx.rs: -------------------------------------------------------------------------------- 1 | use std::marker::PhantomData; 2 | 3 | use image_ext::BgrImage; 4 | use once_cell::sync::Lazy; 5 | use onnxruntime::{ 6 | environment::Environment, 7 | ndarray, 8 | ndarray::{arr1, Array4, ArrayD, ArrayView4, Axis, IxDyn}, 9 | session::{Input, Session}, 10 | tensor::OrtOwnedTensor, 11 | GraphOptimizationLevel, LoggingLevel, OrtError, TensorElementDataType, 12 | TypeToTensorElementDataType, 13 | }; 14 | use thiserror::Error; 15 | 16 | use crate::app::Processor; 17 | 18 | // ONNX global environment to provide 'static to any session 19 | static ENVIRONMENT: Lazy = Lazy::new(|| { 20 | #[cfg(debug_assertions)] 21 | const LOGGING_LEVEL: LoggingLevel = LoggingLevel::Verbose; 22 | #[cfg(not(debug_assertions))] 23 | const LOGGING_LEVEL: LoggingLevel = LoggingLevel::Warning; 24 | 25 | Environment::builder() 26 | .with_name(env!("CARGO_PKG_NAME")) 27 | .with_log_level(LOGGING_LEVEL) 28 | .build() 29 | .unwrap() 30 | }); 31 | 32 | /// Error processing model 33 | #[derive(Error, Debug)] 34 | pub(crate) enum ModelProcError { 35 | #[error("couldn't transform image")] 36 | ShapeError(#[from] ndarray::ShapeError), 37 | #[error("scaling to 0-sized output")] 38 | RuntimeError(#[from] OrtError), 39 | } 40 | 41 | /// Error loading model 42 | #[derive(Error, Debug)] 43 | pub(crate) enum ModelCmdError { 44 | #[error(transparent)] 45 | OrtError(#[from] OrtError), 46 | #[error(transparent)] 47 | RuntimeError(#[from] ModelInputFormatError), 48 | } 49 | 50 | #[derive(Error, Debug)] 51 | pub(crate) enum ModelInputFormatError { 52 | #[error("couldn't infer image input")] 53 | Infer(String), 54 | } 55 | 56 | #[derive(Debug, Clone)] 57 | #[allow(dead_code)] 58 | pub(crate) struct ModelInfo { 59 | pub(crate) input_names: Vec, 60 | pub(crate) input0_dtype: String, 61 | pub(crate) output_names: Vec, 62 | } 63 | 64 | #[derive(Debug)] 65 | struct ImageSession<'s> { 66 | session: Session<'s>, 67 | img_proc: ImgPreProc, 68 | model_info: ModelInfo, 69 | } 70 | 71 | /// ONNX session with pre-processing u8 images. 72 | impl<'s> ImageSession<'s> { 73 | /// Construct an `ImageSession` by inferring some required image input meta data. 74 | /// 75 | /// The basic assumption is that images are passed as batches at position 0. 76 | /// 77 | /// #Arguments 78 | /// 79 | /// * `session` - ONNX session with desired runtime behavior 80 | /// * `color_seq` - Order of color channels in the color dimension 81 | /// * `norm_float` - Whether `f32` input should not be scaled from 0-1 but around another mean and std deviation 82 | fn try_from_session( 83 | session: Session<'s>, 84 | color_seq: ColorSeq, 85 | norm_float: Option>, 86 | ) -> Result { 87 | let (dim_seq, color_range) = infer_img_pre_proc(&session.inputs[0], norm_float)?; 88 | let img_proc = ImgPreProc { dim_seq, color_seq, color_range }; 89 | let input_names = session.inputs.iter().map(|i| i.name.clone()).collect(); 90 | let input0_dtype = format!("{:?}", session.inputs[0].input_type); 91 | let output_names = session.outputs.iter().map(|o| o.name.clone()).collect(); 92 | let model_info = ModelInfo { input_names, input0_dtype, output_names }; 93 | Ok(Self { session, img_proc, model_info }) 94 | } 95 | 96 | /// Forward pass an NHWC(BGR) image batch 97 | fn forward( 98 | &mut self, 99 | mut img_tensor: ArrayView4<'_, u8>, 100 | ) -> Result>, ModelProcError> { 101 | let pre = &self.img_proc; 102 | 103 | match pre.color_seq { 104 | ColorSeq::BGR => {} 105 | ColorSeq::RGB => { 106 | img_tensor.invert_axis(Axis(3)); 107 | } 108 | }; 109 | let (col_axis, img_tensor) = match pre.dim_seq { 110 | DimSeq::NHWC => (Axis(3), img_tensor), 111 | DimSeq::NCHW => (Axis(1), img_tensor.permuted_axes([0, 3, 1, 2])), 112 | }; 113 | 114 | // todo: can at least keep an input vec around 115 | let model_tensors: Vec> = match &pre.color_range { 116 | ColorRange::Uint8 => { 117 | // todo: why .run() doesn't accept a view? 118 | // if it actually requires contiguity, to_owned may not provide that 119 | // and we may get BGR flipped if RGB is required (negative stride ignored) or a segfault.. 120 | let owned_img = img_tensor.to_owned(); 121 | self.session.run(vec![owned_img])? 122 | } 123 | ColorRange::Float32(norm) => { 124 | // instead of mapv, we have to recollect to ensure c contiguity 125 | // given potential prior permutations, otherwise onnxruntime segfaults 126 | let mut img_tensor_float = Array4::from_shape_vec( 127 | img_tensor.raw_dim(), 128 | img_tensor.iter().cloned().map(|v| f32::from(v) * 1f32 / 255f32).collect(), 129 | )?; 130 | if let Some(norm) = norm { 131 | let mean = arr1(&norm.mean); 132 | let std1 = 1.0f32 / arr1(&norm.std); 133 | for mut lane in img_tensor_float.lanes_mut(col_axis) { 134 | lane -= &mean; 135 | lane *= &std1; 136 | } 137 | }; 138 | self.session.run(vec![img_tensor_float])? 139 | } 140 | }; 141 | Ok(model_tensors) 142 | } 143 | } 144 | 145 | /// ONNX model session 146 | pub(crate) struct Model<'s, T = f32> { 147 | img_session: Option>, 148 | _marker: PhantomData, 149 | } 150 | 151 | impl Default for Model<'_, T> { 152 | fn default() -> Self { 153 | Self { img_session: None, _marker: PhantomData } 154 | } 155 | } 156 | 157 | /// Order of color channels of a model's image input 158 | #[derive(Debug, Clone)] 159 | #[allow(clippy::upper_case_acronyms)] 160 | pub(crate) enum ColorSeq { 161 | RGB, 162 | BGR, 163 | } 164 | 165 | /// Relative to a 0-1 range, subtract each target channel by mean and divide by std 166 | #[derive(Debug, Clone)] 167 | pub(crate) struct ColorNorm { 168 | mean: [T; 3], 169 | std: [T; 3], 170 | } 171 | 172 | // todo: numtraits 173 | impl + Copy> ColorNorm { 174 | /// Default of torchvision's imagenet and many other pre-trained models 175 | fn new_torchvision_rgb() -> Self { 176 | Self { 177 | mean: [0.485.into(), 0.456.into(), 0.406.into()], 178 | std: [0.229.into(), 0.224.into(), 0.225.into()], 179 | } 180 | } 181 | /// Return a version with color channels flipped (last becomes first) 182 | fn flip(&self) -> Self { 183 | Self { 184 | mean: [self.mean[2], self.mean[1], self.mean[0]], 185 | std: [self.std[2], self.std[1], self.std[0]], 186 | } 187 | } 188 | } 189 | 190 | /// DType and nominal color range of a model's image input 191 | #[derive(Debug, Clone)] 192 | pub(crate) enum ColorRange { 193 | // byte normalized to u8::MIN - u8::MAX 194 | Uint8, 195 | // f32 normalized to 0-1 if None or by ColorNorm 196 | Float32(Option>), 197 | } 198 | 199 | /// Order of semantic dimensions of a model's image input 200 | #[derive(Debug, Clone)] 201 | #[allow(clippy::upper_case_acronyms)] 202 | pub(crate) enum DimSeq { 203 | /// Batch + TorchVision convention e.g. 204 | NHWC, 205 | /// Batch + OpenCV convention e.g. 206 | NCHW, 207 | } 208 | 209 | /// Specification of a model's image input 210 | #[derive(Debug, Clone)] 211 | pub(crate) struct ImgPreProc { 212 | dim_seq: DimSeq, 213 | color_seq: ColorSeq, 214 | color_range: ColorRange, 215 | } 216 | 217 | /// Determine partially model's image input requirements heuristically 218 | /// 219 | /// # Arguments 220 | /// 221 | /// * `input` - The session model's description of an image input tensor 222 | /// * `norm_float` - Normalization to apply if and only if input is Float32 (ignored otherwise) 223 | fn infer_img_pre_proc( 224 | input: &Input, 225 | norm_float: Option>, 226 | ) -> Result<(DimSeq, ColorRange), ModelInputFormatError> { 227 | // find first dim with length 3 228 | let col_dim = 229 | input.dimensions.iter().position(|d| d.as_ref() == Some(&3)).ok_or_else(|| { 230 | ModelInputFormatError::Infer( 231 | "couldn't locate model's color input by dimension length 3".to_string(), 232 | ) 233 | })?; 234 | 235 | if input.dimensions.len() != 4 { 236 | return Err(ModelInputFormatError::Infer(format!( 237 | "only 4 dimensions supported got {}", 238 | input.dimensions.len() 239 | ))); 240 | }; 241 | 242 | let dim_seq = match col_dim { 243 | 1 => DimSeq::NCHW, 244 | 3 => DimSeq::NHWC, 245 | p => { 246 | return Err(ModelInputFormatError::Infer(format!( 247 | "color dimension only at NCHW or NHWC but not in position {} supported", 248 | p 249 | ))); 250 | } 251 | }; 252 | 253 | let color_range = match &input.input_type { 254 | TensorElementDataType::Float => ColorRange::Float32(norm_float), 255 | TensorElementDataType::Uint8 => ColorRange::Uint8, 256 | dtype => { 257 | return Err(ModelInputFormatError::Infer(format!( 258 | "only Float (f32) and Uint8 (u8) input supported, got {:?}", 259 | dtype 260 | ))); 261 | } 262 | }; 263 | 264 | Ok((dim_seq, color_range)) 265 | } 266 | 267 | #[derive(Clone, Debug)] 268 | pub(crate) enum ModelCmd { 269 | Load(String), 270 | } 271 | 272 | impl<'s, 'session, T: TypeToTensorElementDataType + std::fmt::Debug + Clone> Processor 273 | for Model<'session, T> 274 | where 275 | 's: 'session, 276 | { 277 | type Command = ModelCmd; 278 | type ControlError = ModelCmdError; 279 | type Input = BgrImage; 280 | type Output = Vec>; 281 | type ProcessResult = Result<(), ModelProcError>; 282 | 283 | fn control(&mut self, cmd: Self::Command) -> Result<&mut Self, Self::ControlError> { 284 | match cmd { 285 | // todo: could use a more advanced fork to control intra vs. inter threads 286 | // e.g.: https://github.com/VOICEVOX/onnruntime-rs 287 | // discussion to migrate to official org: https://github.com/nbigaouette/onnxruntime-rs/issues/112 288 | ModelCmd::Load(string_path) if !string_path.is_empty() => { 289 | let session = ENVIRONMENT 290 | .new_session_builder()? 291 | .with_optimization_level(GraphOptimizationLevel::Extended)? 292 | .with_number_threads(3)? 293 | .with_model_from_file(string_path)?; 294 | 295 | // todo: control col_seq properly instead of hardcoding our conventions 296 | let col_seq = 297 | if matches!(session.inputs[0].input_type, TensorElementDataType::Float) { 298 | ColorSeq::RGB 299 | } else { 300 | ColorSeq::BGR 301 | }; 302 | // todo: control color_norm properly instead of hardcoding our conventions 303 | let norm_float = match col_seq { 304 | ColorSeq::RGB => ColorNorm::new_torchvision_rgb(), 305 | ColorSeq::BGR => ColorNorm::new_torchvision_rgb().flip(), 306 | }; 307 | self.img_session = 308 | Some(ImageSession::try_from_session(session, col_seq, Some(norm_float))?); 309 | } 310 | ModelCmd::Load(_) => { 311 | self.img_session = None; 312 | } 313 | } 314 | Ok(self) 315 | } 316 | 317 | fn advance(&mut self, img: &Self::Input, out: &mut Self::Output) -> Self::ProcessResult { 318 | if let Some(ref mut session) = self.img_session { 319 | let img_shape = [1, img.height() as _, img.width() as _, 3]; 320 | let img_tensor = ArrayView4::from_shape(img_shape, img)?; 321 | 322 | // todo: to return a Deref ArrayViewD with &session from &self, we'd need 323 | // maybe some Rc or GATs: https://github.com/rust-lang/rust/pull/96709 324 | // set cloned output without batch dim 325 | let model_tensors = session.forward(img_tensor)?; 326 | out.clear(); 327 | // strip batch dim and clone 328 | for t in model_tensors { 329 | out.push(t.index_axis(ndarray::Axis(0), 0).into_owned()); 330 | } 331 | } 332 | 333 | Ok(()) 334 | } 335 | 336 | fn is_dirty(&self) -> bool { 337 | false 338 | } 339 | } 340 | 341 | impl Model<'_, T> { 342 | pub(crate) fn get_info(&self) -> Option<&ModelInfo> { 343 | self.img_session.as_ref().map(|s| &s.model_info) 344 | } 345 | } 346 | 347 | #[cfg(test)] 348 | mod test { 349 | use super::*; 350 | use infur_test_gen::fcn_resnet50_12_int8_onnx; 351 | 352 | fn fcn_seg_int8() -> String { 353 | fcn_resnet50_12_int8_onnx().to_string_lossy().to_string() 354 | } 355 | 356 | #[test] 357 | fn load_seg_model() { 358 | let mut m = Model::::default(); 359 | m.control(ModelCmd::Load(fcn_seg_int8())).unwrap(); 360 | let session = m.img_session.unwrap(); 361 | eprintln!("model {} session {:?}", fcn_seg_int8(), session); 362 | for (i, inp) in session.session.inputs.iter().enumerate() { 363 | eprintln!("input {}: {:?} {} {:?}", i, inp.dimensions, inp.name, inp.input_type); 364 | } 365 | for (i, out) in session.session.outputs.iter().enumerate() { 366 | eprintln!("output {}: {:?} {} {:?}", i, out.dimensions, out.name, out.output_type); 367 | } 368 | } 369 | 370 | #[test] 371 | fn infer_seg_model() { 372 | let mut m = Model::::default(); 373 | m.control(ModelCmd::Load(fcn_seg_int8())).unwrap(); 374 | let img = BgrImage::new(320, 240); 375 | let mut tensors = vec![]; 376 | m.advance(&img, &mut tensors).unwrap(); 377 | 378 | assert_eq!(tensors.len(), 2, "this segmentation model should return two tensors"); 379 | assert_eq!(tensors[0].shape(), [21, 240, 320], "out should be 21 classes upscaled"); 380 | assert_eq!(tensors[1].shape(), [21, 240, 320], "aux should be 21 classes upscaled"); 381 | } 382 | } 383 | -------------------------------------------------------------------------------- /ff-video/src/parse.rs: -------------------------------------------------------------------------------- 1 | #![allow(dead_code)] 2 | 3 | use std::{error::Error, fmt::Display}; 4 | 5 | /// Describes one video stream 6 | #[derive(Debug, Clone, PartialEq)] 7 | pub struct Stream { 8 | num: u32, 9 | pub width: u32, 10 | pub height: u32, 11 | pub fps: Option, 12 | } 13 | 14 | /// Describe in- or output video stream 15 | #[derive(Debug, Clone, PartialEq)] 16 | pub enum StreamInfo { 17 | Input { from: String, stream: Stream }, 18 | Output { to: String, stream: Stream }, 19 | } 20 | 21 | /// Describes a stream's update 22 | #[derive(Debug, Clone, PartialEq)] 23 | pub struct FrameUpdate { 24 | frame: u64, 25 | fps: Option, 26 | dup: Option, 27 | drop: Option, 28 | } 29 | 30 | /// Describes a video's stream updates 31 | #[derive(Debug, Clone, PartialEq)] 32 | pub enum VideoInfo { 33 | Stream(StreamInfo), 34 | Frame(FrameUpdate), 35 | Codec(String), 36 | } 37 | 38 | #[derive(Debug, Clone, PartialEq)] 39 | enum ParseContext { 40 | Stateless, 41 | Output(u32, String), 42 | Input(u32, String), 43 | } 44 | 45 | /// Parses ffmpeg's stdout into VideoInfo 46 | #[derive(Debug, Clone)] 47 | pub(crate) struct InfoParser { 48 | mode: ParseContext, 49 | } 50 | 51 | #[derive(Debug, Clone, PartialEq)] 52 | pub struct ParseError { 53 | context: ParseContext, 54 | line: String, 55 | reason: String, 56 | } 57 | 58 | impl Error for ParseError {} 59 | 60 | impl Display for ParseError { 61 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 62 | write!(f, "{:?}", self) 63 | } 64 | } 65 | 66 | type InfoResult = std::result::Result, ParseError>; 67 | pub type Result = std::result::Result; 68 | 69 | impl InfoParser { 70 | pub fn default() -> Self { 71 | InfoParser { mode: ParseContext::Stateless } 72 | } 73 | 74 | fn error_on(&self, reason: impl Into, line: &str) -> ParseError { 75 | ParseError { context: self.mode.clone(), line: line.to_string(), reason: reason.into() } 76 | } 77 | 78 | pub fn push(&mut self, line: &str) -> InfoResult { 79 | let error_on = |reason| self.error_on(reason, line); 80 | let error_on_ = |reason| self.error_on(reason, line); // no generic closures 81 | 82 | // Begin Stream 83 | let output = line.strip_prefix("Output #").unwrap_or(line); 84 | let input = line.strip_prefix("Input #").unwrap_or(line); 85 | let in_out = match (input.len(), output.len()) { 86 | (i, _) if i < line.len() => Some((true, input)), 87 | (_, o) if o < line.len() => Some((false, output)), 88 | _ => None, 89 | }; 90 | if let Some((is_input, remaining)) = in_out { 91 | let mut parts = remaining.split(','); 92 | let num_stream = parts 93 | .next() 94 | .ok_or_else(|| error_on("no delimiter after output number"))? 95 | .trim() 96 | .parse::() 97 | .map_err(|e| error_on_(format!("# not a number but {:?}", e)))?; 98 | let to_from = 99 | parts.last().ok_or_else(|| error_on("no last stream element (from or to)"))?.trim(); 100 | 101 | // unquote and extract if possible 102 | let to_from = 103 | to_from.strip_prefix(if is_input { "from '" } else { "to '" }).unwrap_or(to_from); 104 | let to_from = to_from.strip_suffix("':").unwrap_or(to_from); 105 | 106 | self.mode = if is_input { 107 | ParseContext::Input(num_stream, to_from.to_string()) 108 | } else { 109 | ParseContext::Output(num_stream, to_from.to_string()) 110 | }; 111 | return Ok(None); 112 | } 113 | 114 | // Codec 115 | if line.starts_with('[') && line.contains(']') { 116 | return Ok(Some(VideoInfo::Codec(line.into()))); 117 | } 118 | 119 | let line_trimmed = line.trim(); 120 | let frame_str = line_trimmed.strip_prefix("frame=").unwrap_or(line_trimmed); 121 | 122 | // reset if some other header comes up 123 | if line_trimmed.len() == line.len() && !frame_str.len() < line.len() { 124 | self.mode = ParseContext::Stateless; 125 | return Ok(None); 126 | } 127 | 128 | // VideoInfos 129 | let stream_str = line_trimmed.strip_prefix("Stream #").unwrap_or(line_trimmed); 130 | if !matches!(self.mode, ParseContext::Stateless) && stream_str.len() < line_trimmed.len() { 131 | // let chains ftw: https://github.com/rust-lang/rust/issues/53667 132 | let (is_input, num_stream, to_from) = match self.mode { 133 | ParseContext::Input(num_stream, ref from) => (true, num_stream, from), 134 | ParseContext::Output(num_stream, ref to) => (false, num_stream, to), 135 | _ => return Err(error_on("found Stream while not looking for it")), 136 | }; 137 | let mut parts = stream_str.split(':'); 138 | let parse_num_stream = parts 139 | .next() 140 | .ok_or_else(|| error_on("no delimiter after stream number"))? 141 | .parse::() 142 | .map_err(|e| error_on_(format!("Stream # not a number {:?}", e)))?; 143 | 144 | if num_stream != parse_num_stream { 145 | return Err(error_on_(format!("Stream {} didn't match Output", parse_num_stream))); 146 | }; 147 | 148 | let mut is_video = false; 149 | let mut width_height = None; 150 | let mut fps = None; 151 | for p in parts { 152 | if !is_video && p.trim() == "Video" { 153 | is_video = true; 154 | } 155 | if is_video { 156 | for key_vals in p.split(',') { 157 | let key_vals = key_vals.trim(); 158 | let fps_vals = key_vals.trim_end_matches(" fps"); 159 | if fps_vals.len() < key_vals.len() { 160 | fps = fps_vals 161 | .parse::() 162 | .map_err(|_| error_on("fps not a number"))? 163 | .into(); 164 | } else { 165 | let mut dim_vals = key_vals.splitn(2, 'x'); 166 | if let (Some(width_str), Some(height_str)) = 167 | (dim_vals.next(), dim_vals.next()) 168 | { 169 | // some annoying trailing whitespace 170 | let height_str = 171 | height_str.split_once(' ').map_or_else(|| height_str, |v| v.0); 172 | if let (Ok(w), Ok(h)) = 173 | (width_str.parse::(), height_str.parse::()) 174 | { 175 | width_height = Some((w, h)) 176 | }; 177 | } 178 | } 179 | } 180 | } 181 | } 182 | if !is_video { 183 | return Ok(None); 184 | } 185 | return if let Some((width, height)) = width_height { 186 | let stream = Stream { num: num_stream, width, height, fps }; 187 | let info = if is_input { 188 | VideoInfo::Stream(StreamInfo::Input { from: to_from.clone(), stream }) 189 | } else { 190 | VideoInfo::Stream(StreamInfo::Output { to: to_from.clone(), stream }) 191 | }; 192 | self.mode = ParseContext::Stateless; 193 | Ok(Some(info)) 194 | } else { 195 | Err(error_on("didn't find x in video stream")) 196 | }; 197 | } 198 | 199 | // Frame message 200 | if frame_str.len() < line_trimmed.len() { 201 | // frame required 202 | if let Some((frame_num_str, mut frame_rest)) = frame_str.trim().split_once(' ') { 203 | let frame = frame_num_str 204 | .trim() 205 | .parse::() 206 | .map_err(|_| error_on("frame is no number"))?; 207 | 208 | // other key values if available 209 | let (mut fps, mut dup, mut drop) = (None, None, None); 210 | while let Some((key, rest)) = frame_rest.split_once('=') { 211 | if let Some((value, rest)) = rest.trim_start().split_once(' ') { 212 | match key { 213 | "fps" => fps = value.parse().ok(), 214 | "dup" => dup = value.parse().ok(), 215 | "drop" => drop = value.parse().ok(), 216 | _ => {} 217 | } 218 | frame_rest = rest; 219 | } else { 220 | frame_rest = rest; 221 | } 222 | } 223 | let frame_upd = FrameUpdate { frame, fps, dup, drop }; 224 | Ok(Some(VideoInfo::Frame(frame_upd))) 225 | } else { 226 | Ok(None) 227 | } 228 | } else { 229 | Ok(None) 230 | } 231 | } 232 | 233 | pub fn iter_on<'a, I, T: AsRef>( 234 | &'a mut self, 235 | lines: I, 236 | ) -> impl Iterator + 'a 237 | where 238 | I: IntoIterator + 'a, 239 | { 240 | fn un_opt(info: InfoResult) -> Option { 241 | match info { 242 | Err(e) => Some(Err(e)), 243 | Ok(Some(m)) => Some(Ok(m)), 244 | Ok(None) => None, 245 | } 246 | } 247 | 248 | lines.into_iter().map(|l| self.push(l.as_ref())).filter_map(un_opt) 249 | } 250 | } 251 | 252 | /// Blanket implementation for lines of ffmpeg's default stderr bytes. 253 | pub(crate) trait FFMpegLineIter: Iterator { 254 | /// Emit lines on \n, \r (CR) or both but never empty lines. 255 | fn ffmpeg_lines(self) -> FFMpegLines 256 | where 257 | Self: Sized, 258 | { 259 | FFMpegLines { inner: self, state: Vec::::new(), to_clear: false } 260 | } 261 | } 262 | 263 | impl FFMpegLineIter for I {} 264 | 265 | pub(crate) struct FFMpegLines 266 | where 267 | I: Sized, 268 | { 269 | inner: I, 270 | state: Vec, 271 | to_clear: bool, 272 | } 273 | 274 | impl FFMpegLines { 275 | pub(crate) fn state(&self) -> &[u8] { 276 | &self.state 277 | } 278 | } 279 | 280 | /// Since ffmpeg terminates frame= lines only by \r without the -progress flag. 281 | /// Alternatives considered: 282 | /// 1) -progress makes fields appear across lines, thus makes parsing them stateful 283 | /// 2) somehow switch terminator from std-lib lines() to CR after stream header, 284 | /// but seems finicky 285 | /// 3) byte-based splits could probably be done more efficiently on BufReader 286 | impl Iterator for FFMpegLines 287 | where 288 | I: Iterator>, 289 | { 290 | type Item = std::io::Result; 291 | fn next(&mut self) -> Option { 292 | for next_byte in self.inner.by_ref() { 293 | match next_byte { 294 | Ok(b) => { 295 | if self.to_clear { 296 | self.state.clear(); 297 | self.to_clear = false; 298 | } 299 | if b == b'\n' || b == b'\r' { 300 | if !self.state.is_empty() { 301 | let line = String::from_utf8_lossy(&self.state).into_owned(); 302 | // defear clearing state until next() time to preserve state 303 | self.to_clear = true; 304 | return Some(Ok(line)); 305 | } 306 | } else { 307 | self.state.push(b); 308 | } 309 | } 310 | // ignore such Err like .lines() 311 | Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {} 312 | Err(e) => { 313 | return Some(Err(e)); 314 | } 315 | }; 316 | } 317 | None 318 | } 319 | } 320 | 321 | #[cfg(test)] 322 | mod test { 323 | use super::{FrameUpdate, InfoParser, Stream, StreamInfo, VideoInfo}; 324 | 325 | static TEST_INFO: &str = r#"Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'media/huhu_test.mp4': 326 | Metadata: 327 | major_brand : isom 328 | minor_version : 512 329 | compatible_brands: isomiso2avc1mp41 330 | title : Session streamed with GStreamer 331 | encoder : Lavf58.45.100 332 | comment : rtsp-server 333 | Duration: 00:29:58.68, start: 0.000000, bitrate: 650 kb/s 334 | Stream #0:0(und): Video: h264 (Main) (avc1 / 0x31637661), yuvj420p(pc, bt709), 1280x720 [SAR 1:1 DAR 16:9], 647 kb/s, 29.59 fps, 30 tbr, 90k tbn, 180k tbc (default) 335 | Metadata: 336 | handler_name : VideoHandler 337 | vendor_id : [0][0][0][0] 338 | Stream mapping: 339 | Stream #0:0 -> #0:0 (h264 (native) -> rawvideo (native)) 340 | Press [q] to stop, [?] for help 341 | [swscaler @ 0x7fb0ac4dc000] deprecated pixel format used, make sure you did set range correctly 342 | Output #0, image2pipe, to 'pipe:': 343 | Metadata: 344 | major_brand : isom 345 | minor_version : 512 346 | compatible_brands: isomiso2avc1mp41 347 | title : Session streamed with GStreamer 348 | comment : rtsp-server 349 | encoder : Lavf58.76.100 350 | Stream #0:0(und): Video: rawvideo (BGR[24] / 0x18524742), bgr24(pc, gbr/bt709/bt709, progressive), 1280x720 [SAR 1:1 DAR 16:9], q=2-31, 663552 kb/s, 30 fps, 30 tbn (default) 351 | Metadata: 352 | handler_name : VideoHandler 353 | vendor_id : [0][0][0][0] 354 | encoder : Lavc58.134.100 rawvideo 355 | 356 | frame= 3926 fps=978 q=-0.0 size=10600200kB time=00:02:10.86 bitrate=663552.0kbits/s speed=32.6x 357 | frame= 4026 fps=1002 q=-0.0 size=10870200kB time=00:02:14.20 bitrate=663552.0kbits/s speed=33.4x 358 | frame=27045 fps= 1019.6 q=-0.0 size=73021500kB time=00:15:01.50 bitrate=663552.0kbits/s dup= 0 drop=5 speed= 34x"#; 359 | 360 | #[test] 361 | fn test_all_info() { 362 | let mut parser = InfoParser::default(); 363 | let mut infos = parser.iter_on(TEST_INFO.lines()); 364 | 365 | // input 366 | assert_eq!( 367 | infos.next().unwrap(), 368 | Ok(VideoInfo::Stream(StreamInfo::Input { 369 | stream: Stream { num: 0, width: 1280, height: 720, fps: Some(29.59f32) }, 370 | from: "media/huhu_test.mp4".to_string(), 371 | })) 372 | ); 373 | // codec 374 | assert_eq!( 375 | infos.next().unwrap(), 376 | Ok(VideoInfo::Codec("[swscaler @ 0x7fb0ac4dc000] deprecated pixel format used, make sure you did set range correctly".into())) 377 | ); 378 | 379 | // output 380 | assert_eq!( 381 | infos.next().unwrap(), 382 | Ok(VideoInfo::Stream(StreamInfo::Output { 383 | stream: Stream { num: 0, width: 1280, height: 720, fps: Some(30f32) }, 384 | to: "pipe:".to_string(), 385 | })) 386 | ); 387 | 388 | // frames 389 | assert_eq!( 390 | infos.next().unwrap(), 391 | Ok(VideoInfo::Frame(FrameUpdate { 392 | frame: 3926, 393 | fps: Some(978f32), 394 | dup: None, 395 | drop: None, 396 | })) 397 | ); 398 | assert_eq!( 399 | infos.next().unwrap(), 400 | Ok(VideoInfo::Frame(FrameUpdate { 401 | frame: 4026, 402 | fps: Some(1002f32), 403 | dup: None, 404 | drop: None, 405 | })) 406 | ); 407 | assert_eq!( 408 | infos.next().unwrap(), 409 | Ok(VideoInfo::Frame(FrameUpdate { 410 | frame: 27045, 411 | fps: Some(1019.6f32), 412 | dup: Some(0), 413 | drop: Some(5), 414 | })) 415 | ); 416 | } 417 | #[test] 418 | fn test_illegal_input() { 419 | assert!(InfoParser::default() 420 | .iter_on(["Input #1, from 'x':\n]", " Stream #1: Video: abc, 1X01x42 , 20 fps"]) 421 | .next() 422 | .unwrap() 423 | .is_err()); 424 | } 425 | #[test] 426 | fn test_illegal_output() { 427 | assert!(InfoParser::default() 428 | .iter_on(["Output #2, from 'x':\n]", " Stream #1: Video: abc, 100x100 , 20 fps"]) 429 | .next() 430 | .unwrap() 431 | .is_err()); 432 | } 433 | #[test] 434 | fn test_illegal_frame() { 435 | assert_eq!(InfoParser::default() 436 | .push("frame= ---- fps=978 q=-0.0 size=10600200kB time=00:02:10.86 bitrate=663552.0kbits/s speed=32.6x") 437 | .unwrap_err().reason, 438 | "frame is no number".to_string()); 439 | } 440 | 441 | /// More test cases by sampling: https://gist.github.com/jsturgis/3b19447b304616f18657 442 | static BULL_RUN: &str = r#"Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/WeAreGoingOnBullrun.mp4': 443 | Metadata: 444 | major_brand : mp42 445 | minor_version : 0 446 | compatible_brands: isomavc1mp42 447 | creation_time : 2010-06-28T18:32:12.000000Z 448 | Duration: 00:00:47.46, start: 0.000000, bitrate: 2222 kb/s 449 | Stream #0:0(und): Audio: aac (LC) (mp4a / 0x6134706D), 44100 Hz, stereo, fltp, 125 kb/s (default)"#; 450 | 451 | #[test] 452 | fn test_bull_run_audio_input() { 453 | for info in InfoParser::default().iter_on(BULL_RUN.lines()) { 454 | info.unwrap(); 455 | } 456 | } 457 | 458 | #[test] 459 | fn test_ffmpeg_lines() { 460 | use super::FFMpegLineIter; 461 | 462 | let mut lines = 463 | b"foo\nbar\n\rbaz\rbaf\r".iter().map(|b| std::io::Result::Ok(*b)).ffmpeg_lines(); 464 | 465 | assert_eq!(lines.next().unwrap().unwrap(), "foo"); 466 | assert_eq!(lines.next().unwrap().unwrap(), "bar"); 467 | assert_eq!(lines.next().unwrap().unwrap(), "baz"); 468 | assert_eq!(lines.next().unwrap().unwrap(), "baf"); 469 | assert!(lines.next().is_none()); 470 | 471 | assert_eq!(lines.state(), b"baf"); 472 | } 473 | } 474 | --------------------------------------------------------------------------------