├── .gitignore ├── rav1e-by-gop ├── src │ ├── remote │ │ ├── mod.rs │ │ ├── client.rs │ │ └── server.rs │ ├── compress.rs │ ├── muxer │ │ ├── ivf.rs │ │ └── mod.rs │ ├── bin │ │ ├── decode.rs │ │ ├── remote.rs │ │ ├── progress.rs │ │ ├── encode.rs │ │ ├── rav1e-by-gop.rs │ │ └── analyze.rs │ ├── lib.rs │ └── encode │ │ ├── mod.rs │ │ └── stats.rs └── Cargo.toml ├── rustfmt.toml ├── .editorconfig ├── Cargo.toml ├── workers.example.toml ├── CHANGELOG.md ├── LICENSE ├── rav1e-worker ├── Cargo.toml └── src │ ├── server │ ├── mod.rs │ ├── helpers.rs │ └── routes.rs │ ├── main.rs │ └── worker.rs ├── .github └── workflows │ ├── deploy.yml │ └── rav1e-by-gop.yml └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | **/*.rs.bk 3 | /.idea 4 | workers.toml 5 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/remote/mod.rs: -------------------------------------------------------------------------------- 1 | pub use client::*; 2 | pub use server::*; 3 | 4 | mod client; 5 | mod server; 6 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | unstable_features = true 2 | imports_layout = "HorizontalVertical" 3 | imports_granularity = "Crate" 4 | group_imports = "StdExternalCrate" 5 | format_strings = true -------------------------------------------------------------------------------- /.editorconfig: -------------------------------------------------------------------------------- 1 | root = true 2 | 3 | [*] 4 | end_of_line = lf 5 | insert_final_newline = true 6 | 7 | [*.rs] 8 | charset = utf-8 9 | indent_style = space 10 | indent_size = 4 11 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = ["rav1e-by-gop", "rav1e-worker"] 3 | 4 | [profile.dev] 5 | opt-level = 1 6 | 7 | [profile.release] 8 | lto = "thin" 9 | codegen-units = 1 10 | -------------------------------------------------------------------------------- /workers.example.toml: -------------------------------------------------------------------------------- 1 | # Each worker will begin with the `[[workers]]` key 2 | [[workers]] 3 | # `host` is required, for obvious reasons 4 | host = "192.168.1.101" 5 | # `port` is optional, and will use the default server port if not specified 6 | port = 9000 7 | # `password` is required; it should match the server password 8 | password = "Friendly-Influence-Warm-Notice-6" 9 | 10 | [[workers]] 11 | # IPv6 also works 12 | host = "2001:db8::8a2e:370:7334" 13 | password = "Reasonable-Speak-Bear-Hammer-9" 14 | 15 | [[workers]] 16 | # A domain name will work also 17 | host = "worker-1.mydomain.com" 18 | password = "Interfere-Customer-Cut-Landlord-2" 19 | # Recommended, but requires the server be configured 20 | # with a valid TLS certificate 21 | secure = true 22 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/compress.rs: -------------------------------------------------------------------------------- 1 | //! Utilities for losslessly compressing frames in memory. 2 | //! The goal here is to reduce memory usage, 3 | //! because uncompressed Y4M frames are quite large. 4 | 5 | use serde::{de::DeserializeOwned, Serialize}; 6 | use v_frame::{frame::Frame, pixel::Pixel}; 7 | 8 | pub fn compress_frame(frame: &Frame) -> Vec { 9 | let mut compressed_frame = Vec::new(); 10 | let mut encoder = zstd::Encoder::new(&mut compressed_frame, 0).unwrap(); 11 | bincode::serialize_into(&mut encoder, frame).unwrap(); 12 | encoder.finish().unwrap(); 13 | compressed_frame 14 | } 15 | 16 | pub fn decompress_frame(compressed_frame: &[u8]) -> Frame { 17 | let decoder = zstd::Decoder::new(compressed_frame).unwrap(); 18 | bincode::deserialize_from(decoder).unwrap() 19 | } 20 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## Version 0.2.0 (unreleased) 2 | - [Feature] Encoding is now done in one pass instead of two. 3 | This should somewhat speed up the encoding process 4 | and make it more streamlined overall. 5 | - [Breaking/Feature] The `--pass` option has been removed. 6 | It is no longer needed, and you can use your OS's standard 7 | method for piping in input instead, e.g. 8 | `ffmpeg -i myinput.mkv -f yuv4mpegpipe - | rav1e-by-gop - -o myencode.ivf` 9 | - [Breaking] The `--fast-fp` option has been removed. 10 | It is no longer useful. 11 | - [Feature] Add options for limiting number of threads 12 | based on available memory. 13 | No more needing to manually guess how many threads to use. 14 | Adjust via the `--memory` CLI option-- 15 | enabled with "light" setting by default. 16 | - [Breaking] Progress files from prior to this release will not work 17 | with this release. This means that in-progress encodes 18 | from a previous version cannot be resumed. 19 | - Improved scenechange detection algorithm. 20 | 21 | ## Version 0.1.0 22 | - Initial release 23 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/remote/client.rs: -------------------------------------------------------------------------------- 1 | use serde::{Deserialize, Serialize}; 2 | 3 | use crate::{EncodeOptions, VideoDetails}; 4 | 5 | /// The client sends this as a text-based "handshake" to discover 6 | /// the number of worker threads available in the server. 7 | pub const WORKER_QUERY_MESSAGE: &str = "handshake"; 8 | 9 | /// Indicates that the client is requesting a worker slot on the server. 10 | /// Intended to be a very small message to acquire a slot 11 | /// before loading frames into memory and sending them over the network. 12 | /// 13 | /// This is needed to be a separate struct due to needing to know the pixel depth 14 | /// before receiving an encoder message. 15 | #[derive(Debug, Clone, Serialize, Deserialize)] 16 | pub struct SlotRequestMessage { 17 | pub options: EncodeOptions, 18 | pub video_info: VideoDetails, 19 | pub client_version: semver::Version, 20 | } 21 | 22 | #[derive(Debug, Clone, Serialize, Deserialize)] 23 | pub struct PostSegmentMessage { 24 | pub keyframe_number: usize, 25 | pub segment_idx: usize, 26 | pub next_analysis_frame: usize, 27 | } 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Multimedia and Rust 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 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/remote/server.rs: -------------------------------------------------------------------------------- 1 | use std::sync::Arc; 2 | 3 | use chrono::{DateTime, Utc}; 4 | use serde::{Deserialize, Serialize}; 5 | use uuid::Uuid; 6 | 7 | use crate::{ProgressInfo, SegmentFrameData, SerializableProgressInfo}; 8 | 9 | #[derive(Debug, Clone, Copy, Serialize, Deserialize)] 10 | pub struct GetInfoResponse { 11 | pub worker_count: usize, 12 | } 13 | 14 | #[derive(Debug, Clone, Copy, Serialize, Deserialize)] 15 | pub struct PostEnqueueResponse { 16 | pub request_id: Uuid, 17 | } 18 | 19 | #[derive(Debug, Clone, Serialize, Deserialize)] 20 | pub struct GetProgressResponse { 21 | pub progress: SerializableProgressInfo, 22 | pub done: bool, 23 | } 24 | 25 | pub enum EncodeState { 26 | Enqueued, 27 | AwaitingInfo { 28 | time_ready: DateTime, 29 | }, 30 | AwaitingData { 31 | keyframe_number: usize, 32 | segment_idx: usize, 33 | next_analysis_frame: usize, 34 | time_ready: DateTime, 35 | }, 36 | Ready { 37 | keyframe_number: usize, 38 | segment_idx: usize, 39 | next_analysis_frame: usize, 40 | raw_frames: Arc, 41 | }, 42 | InProgress { 43 | progress: ProgressInfo, 44 | }, 45 | EncodingDone { 46 | progress: ProgressInfo, 47 | encoded_data: Vec, 48 | time_finished: DateTime, 49 | }, 50 | } 51 | -------------------------------------------------------------------------------- /rav1e-worker/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "rav1e-worker" 3 | version = "0.5.0" 4 | authors = ["Josh Holmer "] 5 | edition = "2018" 6 | license = "MIT" 7 | description = "A server to enable distributed encoding with rav1e-by-gop" 8 | 9 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 10 | [dependencies] 11 | anyhow = "1" 12 | bincode = "1.3" 13 | bcrypt = "0.9" 14 | byteorder = "1.4.2" 15 | bytes = "1.0.1" 16 | chrono = "0.4" 17 | clap = "2" 18 | crossbeam-channel = "0.5" 19 | env_logger = "0.8" 20 | http = "0.2" 21 | itertools = "0.10" 22 | lazy_static = "1.4.0" 23 | log = "0.4" 24 | num_cpus = "1" 25 | parking_lot = "0.11" 26 | rand = "0.8.3" 27 | rav1e = { version = "0.4", default-features = false, features = ["asm", "serialize"] } 28 | rav1e-by-gop = { path = "../rav1e-by-gop", default-features = false, features = ["remote"] } 29 | rayon = "1.5" 30 | semver = "0.11" 31 | serde = { version = "1", features = ["derive"] } 32 | serde_json = "1.0.64" 33 | tokio = { version = "1.2", features = ["rt-multi-thread", "macros", "time"] } 34 | uuid = { version = "0.8", features = ["v1", "serde"] } 35 | v_frame = { version = "0.2", features = ["serialize"] } 36 | warp = "0.3" 37 | zstd = "0.6" 38 | 39 | [target.'cfg(all(target_arch = "x86_64", target_os = "linux"))'.dependencies] 40 | jemallocator = { version = "0.3.2", features = ["background_threads"] } 41 | 42 | [features] 43 | default = ["tls"] 44 | tls = ["warp/tls"] 45 | -------------------------------------------------------------------------------- /rav1e-worker/src/server/mod.rs: -------------------------------------------------------------------------------- 1 | use std::{env, net::SocketAddrV4, path::PathBuf}; 2 | 3 | use bcrypt::{hash, DEFAULT_COST}; 4 | use lazy_static::lazy_static; 5 | use log::info; 6 | 7 | use crate::server::routes::get_routes; 8 | 9 | mod helpers; 10 | mod routes; 11 | 12 | lazy_static! { 13 | static ref HASHED_SERVER_PASSWORD: String = 14 | hash(env::var("SERVER_PASSWORD").unwrap(), DEFAULT_COST).unwrap(); 15 | static ref CLIENT_VERSION_REQUIRED: semver::VersionReq = { 16 | let server_version = semver::Version::parse(env!("CARGO_PKG_VERSION")).unwrap(); 17 | if server_version.major > 0 { 18 | semver::VersionReq::parse(&format!("^{}.0.0", server_version.major)).unwrap() 19 | } else { 20 | semver::VersionReq::parse(&format!("~0.{}.0", server_version.minor)).unwrap() 21 | } 22 | }; 23 | } 24 | 25 | pub async fn start_listener( 26 | server_ip: SocketAddrV4, 27 | temp_dir: Option, 28 | worker_threads: usize, 29 | ) { 30 | // This thread watches for new incoming connections, 31 | // both for the initial negotiation and for new slot requests 32 | tokio::spawn(async move { 33 | info!("Remote listener started on {}", server_ip); 34 | 35 | match (env::var("TLS_CERT_PATH"), env::var("TLS_KEY_PATH")) { 36 | (Ok(cert_path), Ok(key_path)) => { 37 | warp::serve(get_routes(temp_dir, worker_threads)) 38 | .tls() 39 | .cert_path(&cert_path) 40 | .key_path(&key_path) 41 | .run(server_ip) 42 | .await; 43 | } 44 | _ => { 45 | warp::serve(get_routes(temp_dir, worker_threads)) 46 | .run(server_ip) 47 | .await; 48 | } 49 | }; 50 | }); 51 | } 52 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/muxer/ivf.rs: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2001-2016, Alliance for Open Media. All rights reserved 2 | // Copyright (c) 2017-2019, The rav1e contributors. All rights reserved 3 | // 4 | // This source code is subject to the terms of the BSD 2 Clause License and 5 | // the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License 6 | // was not distributed with this source code in the LICENSE file, you can 7 | // obtain it at www.aomedia.org/license/software. If the Alliance for Open 8 | // Media Patent License 1.0 was not distributed with this source code in the 9 | // PATENTS file, you can obtain it at www.aomedia.org/license/patent. 10 | 11 | use std::{ 12 | fs::File, 13 | io, 14 | io::{sink, BufWriter, Sink, Write}, 15 | }; 16 | 17 | use anyhow::Result; 18 | use ivf::*; 19 | use rav1e::prelude::*; 20 | 21 | use super::Muxer; 22 | 23 | pub struct IvfMuxer { 24 | pub output: W, 25 | } 26 | 27 | impl Muxer for IvfMuxer { 28 | fn write_header( 29 | &mut self, 30 | width: usize, 31 | height: usize, 32 | framerate_num: usize, 33 | framerate_den: usize, 34 | ) { 35 | write_ivf_header( 36 | &mut self.output, 37 | width, 38 | height, 39 | framerate_num, 40 | framerate_den, 41 | ); 42 | } 43 | 44 | fn write_frame(&mut self, pts: u64, data: &[u8], _frame_type: FrameType) { 45 | write_ivf_frame(&mut self.output, pts, data); 46 | } 47 | 48 | fn flush(&mut self) -> io::Result<()> { 49 | self.output.flush() 50 | } 51 | } 52 | 53 | impl IvfMuxer { 54 | pub fn open(path: &str) -> Result>> { 55 | let ivf = IvfMuxer { 56 | output: BufWriter::new(File::create(path)?), 57 | }; 58 | Ok(ivf) 59 | } 60 | 61 | pub fn null() -> IvfMuxer { 62 | IvfMuxer { output: sink() } 63 | } 64 | 65 | pub fn in_memory() -> IvfMuxer> { 66 | IvfMuxer { output: Vec::new() } 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/muxer/mod.rs: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017-2019, The rav1e contributors. All rights reserved 2 | // 3 | // This source code is subject to the terms of the BSD 2 Clause License and 4 | // the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License 5 | // was not distributed with this source code in the LICENSE file, you can 6 | // obtain it at www.aomedia.org/license/software. If the Alliance for Open 7 | // Media Patent License 1.0 was not distributed with this source code in the 8 | // PATENTS file, you can obtain it at www.aomedia.org/license/patent. 9 | 10 | mod ivf; 11 | 12 | use std::{ 13 | ffi::OsStr, 14 | fs::File, 15 | io, 16 | io::{BufWriter, Sink}, 17 | }; 18 | 19 | use anyhow::Result; 20 | use rav1e::prelude::*; 21 | 22 | pub use self::ivf::IvfMuxer; 23 | use crate::Output; 24 | 25 | pub trait Muxer { 26 | fn write_header( 27 | &mut self, 28 | width: usize, 29 | height: usize, 30 | framerate_num: usize, 31 | framerate_den: usize, 32 | ); 33 | 34 | fn write_frame(&mut self, pts: u64, data: &[u8], frame_type: FrameType); 35 | 36 | fn flush(&mut self) -> io::Result<()>; 37 | } 38 | 39 | pub fn create_muxer(path: &Output) -> Result> { 40 | match path { 41 | Output::File(path) => { 42 | let ext = path 43 | .extension() 44 | .and_then(OsStr::to_str) 45 | .map(str::to_lowercase) 46 | .unwrap_or_else(|| "ivf".into()); 47 | 48 | match &ext[..] { 49 | "ivf" => Ok(Box::new(IvfMuxer::>::open( 50 | path.to_str().unwrap(), 51 | )?)), 52 | _e => { 53 | panic!( 54 | "{} is not a supported extension, please change to .ivf", 55 | ext 56 | ); 57 | } 58 | } 59 | } 60 | Output::Memory => Ok(Box::new(IvfMuxer::>::in_memory())), 61 | Output::Null => Ok(Box::new(IvfMuxer::::null())), 62 | } 63 | } 64 | -------------------------------------------------------------------------------- /.github/workflows/deploy.yml: -------------------------------------------------------------------------------- 1 | name: deploy 2 | 3 | on: 4 | push: 5 | tags: 6 | - 'v*.*.*' 7 | 8 | jobs: 9 | 10 | deploy-binaries: 11 | 12 | runs-on: windows-latest 13 | 14 | steps: 15 | - uses: actions/checkout@v2 16 | 17 | - name: Install nasm 18 | run: | 19 | $NASM_VERSION="2.15.04" 20 | $LINK="https://www.nasm.us/pub/nasm/releasebuilds/$NASM_VERSION/win64" 21 | curl --ssl-no-revoke -LO "$LINK/nasm-$NASM_VERSION-win64.zip" 22 | 7z e -y "nasm-$NASM_VERSION-win64.zip" -o"C:\nasm" 23 | echo "C:\nasm" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append 24 | 25 | - name: Install Rust 26 | uses: actions-rs/toolchain@v1 27 | with: 28 | profile: minimal 29 | toolchain: stable 30 | override: true 31 | 32 | - name: Set MSVC x86_64 linker path 33 | run: | 34 | $LinkGlob = "VC\Tools\MSVC\*\bin\Hostx64\x64" 35 | $env:PATH = "$env:PATH;${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer" 36 | $LinkPath = vswhere -latest -products * -find "$LinkGlob" | 37 | Select-Object -Last 1 38 | echo "$LinkPath" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append 39 | 40 | - name: Build 41 | run: | 42 | cargo build --release 43 | 44 | - name: Strip binaries 45 | run: | 46 | cd target/release 47 | strip rav1e-by-gop.exe rav1e-worker.exe 48 | 49 | - name: Handle release data and files 50 | shell: bash 51 | id: data 52 | run: | 53 | VERSION=$(head -n 1 CHANGELOG.md | tr -d "## Version ") 54 | echo "::set-output name=version::$VERSION" 55 | tail -n +2 CHANGELOG.md | sed -e '/^$/,$d' > CHANGELOG.txt 56 | 57 | - name: Create a release 58 | uses: softprops/action-gh-release@v1 59 | with: 60 | name: Version ${{ steps.data.outputs.version }} 61 | body_path: CHANGELOG.txt 62 | files: | 63 | target/release/rav1e-by-gop.exe 64 | target/release/rav1e-worker.exe 65 | env: 66 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 67 | -------------------------------------------------------------------------------- /rav1e-worker/src/server/helpers.rs: -------------------------------------------------------------------------------- 1 | use std::convert::Infallible; 2 | 3 | use bcrypt::verify; 4 | use serde::de::DeserializeOwned; 5 | use warp::{ 6 | http::StatusCode, 7 | reject::{MissingHeader, Reject}, 8 | reply::{Json, WithStatus}, 9 | Filter, 10 | Rejection, 11 | Reply, 12 | }; 13 | 14 | use crate::server::HASHED_SERVER_PASSWORD; 15 | 16 | pub fn json_body( 17 | ) -> impl Filter + Copy { 18 | warp::body::content_length_limit(1024 * 1024).and(warp::body::json()) 19 | } 20 | 21 | pub fn with_state( 22 | state: T, 23 | ) -> impl Filter + Clone { 24 | warp::any().map(move || state.clone()) 25 | } 26 | 27 | pub fn require_auth() -> impl Filter + Copy { 28 | warp::header("X-RAV1E-AUTH").and_then(move |password: String| async move { 29 | if verify(password, &HASHED_SERVER_PASSWORD).unwrap() { 30 | return Ok(()); 31 | } 32 | Err(warp::reject::custom(InvalidAuthorization)) 33 | }) 34 | } 35 | 36 | #[derive(Debug, Clone, Copy)] 37 | pub struct InvalidAuthorization; 38 | 39 | impl Reject for InvalidAuthorization {} 40 | 41 | #[derive(Debug, Clone, Copy)] 42 | pub struct ClientVersionMismatch; 43 | 44 | impl Reject for ClientVersionMismatch {} 45 | 46 | pub async fn handle_rejection_types(err: Rejection) -> Result { 47 | if err.find::().is_some() { 48 | return Ok(warp::reply::with_status( 49 | "Incorrect server password".to_string(), 50 | StatusCode::UNAUTHORIZED, 51 | )); 52 | } 53 | if let Some(err) = err.find::() { 54 | if err.name() == "X-RAV1E-AUTH" { 55 | return Ok(warp::reply::with_status( 56 | "Password header not provided".to_string(), 57 | StatusCode::UNAUTHORIZED, 58 | )); 59 | } 60 | } 61 | if err.find::().is_some() { 62 | return Ok(warp::reply::with_status( 63 | "Client/server version mismatch".to_string(), 64 | StatusCode::BAD_REQUEST, 65 | )); 66 | } 67 | 68 | Err(err) 69 | } 70 | 71 | pub fn map_error_to_500(_e: E) -> WithStatus { 72 | warp::reply::with_status(warp::reply::json(&()), StatusCode::INTERNAL_SERVER_ERROR) 73 | } 74 | 75 | #[macro_export] 76 | macro_rules! try_or_500 { 77 | ( $expr:expr ) => { 78 | match $expr { 79 | Ok(val) => val, 80 | Err(e) => { 81 | return Ok(crate::server::helpers::map_error_to_500(e)); 82 | } 83 | } 84 | }; 85 | } 86 | -------------------------------------------------------------------------------- /rav1e-by-gop/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "rav1e-by-gop" 3 | version = "0.5.0" 4 | authors = ["Josh Holmer "] 5 | edition = "2018" 6 | license = "MIT" 7 | description = "A tool to parallel encode GOP segments of a video with rav1e" 8 | autobins = false 9 | 10 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 11 | [dependencies] 12 | anyhow = "1.0" 13 | arrayvec = "0.5" 14 | av-scenechange = { version = "0.5", default-features = false, features = [ "serde" ], optional = true } 15 | bincode = "1.3" 16 | chrono = { version = "0.4.19", features = ["serde"] } 17 | clap = { version = "2", optional = true } 18 | console = { version = "0.14.0", optional = true } 19 | crossbeam-channel = "0.5" 20 | crossbeam-utils = { version = "0.8", optional = true } 21 | dialoguer = { version = "0.7", optional = true } 22 | env_logger = { version = "0.8.1", optional = true } 23 | http = { version = "0.2", optional = true } 24 | indicatif = { version = "0.15", optional = true } 25 | itertools = { version = "0.10", optional = true } 26 | ivf = "0.1" 27 | lazy_static = { version = "1.0", optional = true } 28 | log = "0.4" 29 | num_cpus = { version = "1", optional = true } 30 | parking_lot = { version = "0.11", optional = true } 31 | rav1e = { version = "0.4", default-features = false, features = ["asm", "serialize", "unstable"] } 32 | rayon = "1.5" 33 | reqwest = { version = "0.11", features = ["blocking", "json"], optional = true } 34 | rmp-serde = { version = "0.15", optional = true } 35 | semver = { version = "0.11", features = ["serde"], optional = true } 36 | serde = { version = "1", features = ["derive"] } 37 | systemstat = "0.1.7" 38 | thiserror = { version = "1", optional = true } 39 | threadpool = "1" 40 | toml = { version = "0.5", optional = true } 41 | url = { version = "2.2", optional = true } 42 | uuid = { version = "0.8", features = ["serde"], optional = true } 43 | v_frame = { version = "0.2", features = ["serialize"] } 44 | y4m = { version = "0.7", optional = true } 45 | zstd = "0.6" 46 | 47 | [target.'cfg(all(target_arch = "x86_64", target_os = "linux"))'.dependencies] 48 | jemallocator = { version = "0.3.2", features = ["background_threads"], optional = true } 49 | 50 | [[bin]] 51 | name = "rav1e-by-gop" 52 | required-features = ["binary"] 53 | bench = false 54 | 55 | [features] 56 | default = ["binary", "remote"] 57 | binary = [ 58 | "av-scenechange", 59 | "clap", 60 | "console", 61 | "crossbeam-utils", 62 | "dialoguer", 63 | "env_logger", 64 | "indicatif", 65 | "itertools", 66 | "jemallocator", 67 | "num_cpus", 68 | "rmp-serde", 69 | "thiserror", 70 | "y4m" 71 | ] 72 | remote = ["http", "lazy_static", "parking_lot", "reqwest", "semver", "toml", "url", "uuid"] 73 | 74 | -------------------------------------------------------------------------------- /.github/workflows/rav1e-by-gop.yml: -------------------------------------------------------------------------------- 1 | name: rav1e-by-gop 2 | 3 | on: 4 | pull_request: 5 | branches: 6 | - master 7 | push: 8 | branches: 9 | - master 10 | 11 | jobs: 12 | clippy-rustfmt: 13 | runs-on: ubuntu-latest 14 | steps: 15 | 16 | - uses: actions/checkout@v2 17 | 18 | - name: Install stable 19 | uses: actions-rs/toolchain@v1 20 | with: 21 | profile: minimal 22 | toolchain: stable 23 | override: true 24 | components: clippy, rustfmt 25 | 26 | - name: Install nasm 27 | env: 28 | LINK: http://debian-archive.trafficmanager.net/debian/pool/main/n/nasm 29 | NASM_VERSION: 2.15.05-1 30 | run: | 31 | curl -O "$LINK/nasm_${NASM_VERSION}_amd64.deb" 32 | sudo dpkg -i "nasm_${NASM_VERSION}_amd64.deb" 33 | 34 | - name: Run rustfmt 35 | uses: actions-rs/cargo@v1 36 | with: 37 | command: fmt 38 | args: -- --check --verbose 39 | 40 | - name: Lint 41 | uses: actions-rs/clippy-check@v1 42 | with: 43 | token: ${{ secrets.GITHUB_TOKEN }} 44 | args: --all-features --tests --benches 45 | name: lint 46 | 47 | 48 | build: 49 | 50 | strategy: 51 | matrix: 52 | platform: [ubuntu-latest, windows-latest] 53 | 54 | runs-on: ${{ matrix.platform }} 55 | 56 | steps: 57 | - uses: actions/checkout@v2 58 | 59 | - name: Install stable 60 | uses: actions-rs/toolchain@v1 61 | with: 62 | profile: minimal 63 | toolchain: stable 64 | override: true 65 | 66 | - name: Install nasm for Ubuntu 67 | if: matrix.platform == 'ubuntu-latest' 68 | env: 69 | LINK: http://debian-archive.trafficmanager.net/debian/pool/main/n/nasm 70 | NASM_VERSION: 2.15.05-1 71 | run: | 72 | curl -O "$LINK/nasm_${NASM_VERSION}_amd64.deb" 73 | sudo dpkg -i "nasm_${NASM_VERSION}_amd64.deb" 74 | 75 | - name: Install nasm for Windows 76 | if: matrix.platform == 'windows-latest' 77 | run: | 78 | $NASM_VERSION="2.15.04" 79 | $LINK="https://www.nasm.us/pub/nasm/releasebuilds/$NASM_VERSION/win64" 80 | curl --ssl-no-revoke -LO "$LINK/nasm-$NASM_VERSION-win64.zip" 81 | 7z e -y "nasm-$NASM_VERSION-win64.zip" -o"C:\nasm" 82 | echo "C:\nasm" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append 83 | 84 | - name: Set MSVC x86_64 linker path 85 | if: matrix.platform == 'windows-latest' 86 | run: | 87 | $LinkGlob = "VC\Tools\MSVC\*\bin\Hostx64\x64" 88 | $env:PATH = "$env:PATH;${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer" 89 | $LinkPath = vswhere -latest -products * -find "$LinkGlob" | 90 | Select-Object -Last 1 91 | echo "$LinkPath" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append 92 | 93 | - name: Build 94 | run: cargo build --all-features --tests --benches 95 | 96 | - name: Run tests 97 | run: cargo test --all-features 98 | 99 | - name: Generate docs 100 | run: cargo doc --all-features --no-deps 101 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rav1e-by-gop 2 | 3 | ## DEPRECATION NOTICE 4 | 5 | **DEPRECATED**: Please use rav1e with the [channel-api feature](https://github.com/xiph/rav1e#current-unstable-features), or [av1an](https://github.com/master-of-zen/Av1an/). 6 | 7 | ## Introduction 8 | 9 | `rav1e-by-gop` is a tool to multithread a rav1e encode 10 | by splitting it into GOPs (at keyframe points). 11 | This allows multithreading without the quality loss of tiles. 12 | 13 | Disclaimer: This tool attempts to be memory efficient 14 | by only opening one reader for the entire encoding step. 15 | This avoids tools such as Vapoursynth, 16 | which may be quite memory hungry, 17 | from using all the memory on your system. 18 | However, memory usage of this tool itself 19 | is still `O(n)` where n is the number of threads. 20 | 21 | By default, this tool will limit the number of threads 22 | based on the memory available in your system. 23 | You can adjust or disable this safeguard using the 24 | `--memory` CLI option. 25 | 26 | ## Basic Usage 27 | 28 | To encode a video from y4m to AV1 (in ivf container): 29 | 30 | `rav1e-by-gop input.y4m -o output.ivf` 31 | 32 | This will use the default encode settings specified by rav1e, 33 | and use threads equal to the number of logical CPU cores. 34 | 35 | ## Advanced Usage 36 | 37 | To configure rav1e settings (only some settings are configurable): 38 | 39 | `rav1e-by-gop input.y4m -s 3 --qp 50 --min-keyint 12 --keyint 360 -o output.ivf` 40 | 41 | To pipe in input: 42 | 43 | `vspipe --y4m input.vpy - | rav1e-by-gop - -o output.ivf` 44 | 45 | To limit the number of threads 46 | (this tool will never use more than the number of CPUs): 47 | 48 | `rav1e-by-gop input.y4m --threads 4 -o output.ivf` 49 | 50 | ## Distributed encoding 51 | 52 | By default, the rav1e-by-gop client can run on a single machine. 53 | However, it can be set up to take advantage of remote machines 54 | running the rav1e worker for distributed encoding. 55 | 56 | ### Server setup 57 | 58 | rav1e-worker runs a server that listens for connections from rav1e-by-gop, 59 | then uses worker threads to encode segments for the remote machines. 60 | 61 | By default, rav1e-worker will accept insecure connections, 62 | because unfortunately self-signed certs will not work. 63 | 64 | You can configure rav1e-worker to use secure TLS connections. 65 | If you have a signed TLS certificate you would like to use, 66 | or would like to get one for free from Let's Encrypt, 67 | you can use environment variables to tell rav1e-by-gop to use it. 68 | 69 | e.g. 70 | 71 | ``` 72 | TLS_CERT_PATH=/etc/letsencrypt/live/mydomain.example/cert.pem 73 | TLS_KEY_PATH=/etc/letsencrypt/live/mydomain.example/privkey.pem 74 | ``` 75 | 76 | Setting both of these will automatically enable TLS in the rav1e-worker server. 77 | 78 | You must also create a password for clients to use to connect to the server. 79 | This is mandatory. You should set this password to the `SERVER_PASSWORD` environment variable. 80 | It is highly recommended to use a secure password generator like [this one](https://correcthorsebatterystaple.net/). 81 | 82 | You can then run `rav1e-worker`, which will listen by default on port 13415, 83 | and spawn workers up to the number of logical CPU cores. 84 | You can configure which IP and port to listen on as well as how many workers 85 | to spawn via CLI args to rav1e-worker. 86 | 87 | ### Client setup 88 | 89 | You can tell rav1e-by-gop to connect to remote workers by using a `workers.toml` file. 90 | There is an example of this file in this repository. 91 | By default, the client will look in the current working directory for the file. 92 | You can override it by using the `--workers /path/to/workers.toml` CLI flag. 93 | 94 | rav1e-by-gop will by default use the local worker threads as well. 95 | You can disable local encoding entirely with the `--no-local` flag. 96 | 97 | ## Limitations 98 | 99 | - Only supports one-pass QP mode 100 | - One-pass bitrate mode will never be supported, 101 | because it's not suitable for the use cases where one would use this tool 102 | - Two-pass bitrate mode may be supported in the future 103 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/bin/decode.rs: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2001-2016, Alliance for Open Media. All rights reserved 2 | // Copyright (c) 2017-2019, The rav1e contributors. All rights reserved 3 | // 4 | // This source code is subject to the terms of the BSD 2 Clause License and 5 | // the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License 6 | // was not distributed with this source code in the LICENSE file, you can 7 | // obtain it at www.aomedia.org/license/software. If the Alliance for Open 8 | // Media Patent License 1.0 was not distributed with this source code in the 9 | // PATENTS file, you can obtain it at www.aomedia.org/license/patent. 10 | 11 | use std::io::{self, Read}; 12 | 13 | use rav1e::prelude::*; 14 | use rav1e_by_gop::VideoDetails; 15 | use thiserror::Error; 16 | use y4m::Decoder; 17 | 18 | #[derive(Debug, Error)] 19 | pub enum DecodeError { 20 | #[error("Reached end of file")] 21 | EndOfFile, 22 | #[error("Bad input")] 23 | BadInput, 24 | #[error("Unknown colorspace")] 25 | UnknownColorspace, 26 | #[error("Parse error")] 27 | ParseError, 28 | #[error("IO Error: {0}")] 29 | IoError(io::Error), 30 | #[error("Memory limit exceeded")] 31 | MemoryLimitExceeded, 32 | } 33 | 34 | pub(crate) fn get_video_details(dec: &Decoder) -> VideoDetails { 35 | let width = dec.get_width(); 36 | let height = dec.get_height(); 37 | let color_space = dec.get_colorspace(); 38 | let bit_depth = color_space.get_bit_depth(); 39 | let (chroma_sampling, chroma_sample_position) = map_y4m_color_space(color_space); 40 | let framerate = dec.get_framerate(); 41 | let time_base = Rational::new(framerate.den as u64, framerate.num as u64); 42 | 43 | VideoDetails { 44 | width, 45 | height, 46 | bit_depth, 47 | chroma_sampling, 48 | chroma_sample_position, 49 | time_base, 50 | } 51 | } 52 | 53 | pub(crate) fn read_raw_frame(dec: &mut Decoder) -> Result, DecodeError> { 54 | dec.read_frame().map_err(Into::into) 55 | } 56 | 57 | pub(crate) fn process_raw_frame( 58 | frame: &y4m::Frame, 59 | ctx: &Context, 60 | cfg: &VideoDetails, 61 | ) -> Frame { 62 | let mut f: Frame = ctx.new_frame(); 63 | let bytes = if cfg.bit_depth <= 8 { 1 } else { 2 }; 64 | 65 | let (chroma_width, _) = cfg 66 | .chroma_sampling 67 | .get_chroma_dimensions(cfg.width, cfg.height); 68 | 69 | f.planes[0].copy_from_raw_u8(frame.get_y_plane(), cfg.width * bytes, bytes); 70 | f.planes[1].copy_from_raw_u8(frame.get_u_plane(), chroma_width * bytes, bytes); 71 | f.planes[2].copy_from_raw_u8(frame.get_v_plane(), chroma_width * bytes, bytes); 72 | f 73 | } 74 | 75 | impl From for DecodeError { 76 | fn from(e: y4m::Error) -> DecodeError { 77 | match e { 78 | y4m::Error::EOF => DecodeError::EndOfFile, 79 | y4m::Error::BadInput => DecodeError::BadInput, 80 | y4m::Error::UnknownColorspace => DecodeError::UnknownColorspace, 81 | y4m::Error::ParseError(_) => DecodeError::ParseError, 82 | y4m::Error::IoError(e) => DecodeError::IoError(e), 83 | // Note that this error code has nothing to do with the system running out of memory, 84 | // it means the y4m decoder has exceeded its memory allocation limit. 85 | y4m::Error::OutOfMemory => DecodeError::MemoryLimitExceeded, 86 | } 87 | } 88 | } 89 | 90 | pub fn map_y4m_color_space( 91 | color_space: y4m::Colorspace, 92 | ) -> (v_frame::prelude::ChromaSampling, ChromaSamplePosition) { 93 | use v_frame::prelude::ChromaSampling::*; 94 | use y4m::Colorspace::*; 95 | use ChromaSamplePosition::*; 96 | 97 | match color_space { 98 | Cmono => (Cs400, Unknown), 99 | C420jpeg | C420paldv => (Cs420, Unknown), 100 | C420mpeg2 => (Cs420, Vertical), 101 | C420 | C420p10 | C420p12 => (Cs420, Colocated), 102 | C422 | C422p10 | C422p12 => (Cs422, Colocated), 103 | C444 | C444p10 | C444p12 => (Cs444, Colocated), 104 | } 105 | } 106 | -------------------------------------------------------------------------------- /rav1e-worker/src/main.rs: -------------------------------------------------------------------------------- 1 | use std::{collections::BTreeMap, env, net::SocketAddrV4, path::PathBuf, time::Duration}; 2 | 3 | use clap::{App, Arg}; 4 | use lazy_static::lazy_static; 5 | use log::{debug, log_enabled}; 6 | use parking_lot::RwLock; 7 | use rand::Rng; 8 | use rav1e_by_gop::{EncodeOptions, EncodeState, VideoDetails}; 9 | use server::*; 10 | use tokio::time::sleep; 11 | use uuid::{v1::Context, Uuid}; 12 | use worker::*; 13 | 14 | mod server; 15 | mod worker; 16 | 17 | #[cfg(all(target_arch = "x86_64", target_os = "linux"))] 18 | #[global_allocator] 19 | static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc; 20 | 21 | lazy_static! { 22 | pub static ref ENCODER_QUEUE: RwLock>> = 23 | RwLock::new(BTreeMap::new()); 24 | pub static ref UUID_CONTEXT: Context = Context::new(0); 25 | pub static ref UUID_NODE_ID: Box<[u8]> = { 26 | let mut id = Vec::with_capacity(6); 27 | let mut rng = rand::thread_rng(); 28 | for _ in 0..6 { 29 | id.push(rng.gen()); 30 | } 31 | id.into_boxed_slice() 32 | }; 33 | } 34 | 35 | pub struct EncodeItem { 36 | pub state: EncodeState, 37 | pub options: EncodeOptions, 38 | pub video_info: VideoDetails, 39 | } 40 | 41 | impl EncodeItem { 42 | fn new(options: EncodeOptions, video_info: VideoDetails) -> Self { 43 | EncodeItem { 44 | state: EncodeState::Enqueued, 45 | options, 46 | video_info, 47 | } 48 | } 49 | } 50 | 51 | impl std::fmt::Debug for EncodeItem { 52 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 53 | match self.state { 54 | EncodeState::Enqueued => f.write_str("Enqueued"), 55 | EncodeState::AwaitingInfo { .. } => f.write_str("Awaiting Segment Info"), 56 | EncodeState::AwaitingData { .. } => f.write_str("Awaiting Data"), 57 | EncodeState::Ready { ref raw_frames, .. } => f.write_fmt(format_args!( 58 | "Ready to encode {} frames", 59 | raw_frames.frame_count() 60 | )), 61 | EncodeState::InProgress { ref progress } => f.write_fmt(format_args!( 62 | "Encoding {} of {} frames", 63 | progress.frame_info.len(), 64 | progress.total_frames 65 | )), 66 | EncodeState::EncodingDone { 67 | ref encoded_data, .. 68 | } => f.write_fmt(format_args!("Done encoding {} bytes", encoded_data.len())), 69 | } 70 | } 71 | } 72 | 73 | #[tokio::main] 74 | async fn main() { 75 | env::var("SERVER_PASSWORD").expect("SERVER_PASSWORD env var MUST be set!"); 76 | 77 | if env::var("RUST_LOG").is_err() { 78 | env::set_var("RUST_LOG", "rav1e_worker=info"); 79 | } 80 | env_logger::init(); 81 | 82 | let matches = App::new("rav1e-worker") 83 | .arg( 84 | Arg::with_name("LISTEN_IP") 85 | .help("Select which IP to listen on") 86 | .long("ip") 87 | .visible_alias("host") 88 | .default_value("0.0.0.0") 89 | .takes_value(true), 90 | ) 91 | .arg( 92 | Arg::with_name("LISTEN_PORT") 93 | .help("Select which port to listen on") 94 | .long("port") 95 | .short("p") 96 | .default_value("13415") 97 | .takes_value(true), 98 | ) 99 | .arg( 100 | Arg::with_name("MAX_THREADS") 101 | .help( 102 | "Limit the number of threads that can be used for workers [default: num cpus]", 103 | ) 104 | .long("threads") 105 | .takes_value(true), 106 | ) 107 | .arg( 108 | Arg::with_name("TEMP_DIR") 109 | .help( 110 | "Store input segments in temp files in the specified directory; by default \ 111 | stores in memory", 112 | ) 113 | .long("temp-dir") 114 | .takes_value(true), 115 | ) 116 | .get_matches(); 117 | 118 | let server_ip = SocketAddrV4::new( 119 | matches.value_of("LISTEN_IP").unwrap().parse().unwrap(), 120 | matches.value_of("LISTEN_PORT").unwrap().parse().unwrap(), 121 | ); 122 | let mut threads = num_cpus::get(); 123 | if let Some(thread_setting) = matches 124 | .value_of("MAX_THREADS") 125 | .and_then(|val| val.parse().ok()) 126 | { 127 | threads = threads.min(thread_setting); 128 | } 129 | let temp_dir = if let Some(temp_dir) = matches.value_of("TEMP_DIR") { 130 | let dir = PathBuf::from(temp_dir); 131 | if !dir.is_dir() { 132 | panic!("Specified temp dir does not exist or is not a directory"); 133 | } 134 | if dir.metadata().unwrap().permissions().readonly() { 135 | panic!("Specified temp dir is not writeable"); 136 | } 137 | Some(dir) 138 | } else { 139 | None 140 | }; 141 | 142 | start_listener(server_ip, temp_dir, threads).await; 143 | start_workers(threads).await; 144 | 145 | loop { 146 | // Run the main thread forever until terminated 147 | if log_enabled!(log::Level::Debug) { 148 | let queue_handle = ENCODER_QUEUE.read(); 149 | let mut items = Vec::with_capacity(queue_handle.len()); 150 | for (key, item) in queue_handle.iter() { 151 | items.push((key, item.read())); 152 | } 153 | debug!("Items in queue: {:?}", items); 154 | } 155 | sleep(Duration::from_secs(5)).await; 156 | } 157 | } 158 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::cognitive_complexity)] 2 | 3 | pub mod compress; 4 | pub mod encode; 5 | pub mod muxer; 6 | #[cfg(feature = "remote")] 7 | pub mod remote; 8 | 9 | use std::{path::PathBuf, sync::Arc}; 10 | 11 | use crossbeam_channel::{Receiver, Sender}; 12 | use rav1e::prelude::*; 13 | use serde::{Deserialize, Serialize}; 14 | #[cfg(feature = "remote")] 15 | use url::Url; 16 | #[cfg(feature = "remote")] 17 | use uuid::Uuid; 18 | 19 | #[cfg(feature = "remote")] 20 | pub use self::remote::*; 21 | pub use self::{compress::*, encode::*, muxer::*}; 22 | 23 | #[derive(Debug, Clone, Copy, Serialize, Deserialize)] 24 | pub struct VideoDetails { 25 | pub width: usize, 26 | pub height: usize, 27 | pub bit_depth: usize, 28 | pub chroma_sampling: ChromaSampling, 29 | pub chroma_sample_position: ChromaSamplePosition, 30 | pub time_base: Rational, 31 | } 32 | 33 | impl Default for VideoDetails { 34 | fn default() -> Self { 35 | VideoDetails { 36 | width: 640, 37 | height: 480, 38 | bit_depth: 8, 39 | chroma_sampling: ChromaSampling::Cs420, 40 | chroma_sample_position: ChromaSamplePosition::Unknown, 41 | time_base: Rational { num: 30, den: 1 }, 42 | } 43 | } 44 | } 45 | 46 | pub enum Slot { 47 | Local(usize), 48 | #[cfg(feature = "remote")] 49 | Remote(Box), 50 | } 51 | 52 | impl Slot { 53 | pub fn is_remote(&self) -> bool { 54 | match self { 55 | Slot::Local(_) => false, 56 | #[cfg(feature = "remote")] 57 | Slot::Remote(_) => true, 58 | } 59 | } 60 | } 61 | 62 | pub struct SegmentData { 63 | pub segment_no: usize, 64 | pub slot: usize, 65 | pub next_analysis_frame: usize, 66 | pub start_frameno: usize, 67 | pub frame_data: SegmentFrameData, 68 | } 69 | 70 | pub enum SegmentFrameData { 71 | CompressedFrames(Vec>), 72 | Y4MFile { path: PathBuf, frame_count: usize }, 73 | } 74 | 75 | impl SegmentFrameData { 76 | pub fn frame_count(&self) -> usize { 77 | match self { 78 | SegmentFrameData::CompressedFrames(frames) => frames.len(), 79 | SegmentFrameData::Y4MFile { frame_count, .. } => *frame_count, 80 | } 81 | } 82 | } 83 | 84 | #[cfg(feature = "remote")] 85 | pub struct ActiveConnection { 86 | pub worker_uri: Url, 87 | pub worker_password: String, 88 | pub request_id: Uuid, 89 | pub slot_in_worker: usize, 90 | pub video_info: VideoDetails, 91 | pub encode_info: Option, 92 | pub worker_update_sender: WorkerUpdateSender, 93 | pub progress_sender: ProgressSender, 94 | } 95 | 96 | #[derive(Debug, Clone)] 97 | pub struct EncodeInfo { 98 | pub output_file: Output, 99 | pub frame_count: usize, 100 | pub next_analysis_frame: usize, 101 | pub segment_idx: usize, 102 | pub start_frameno: usize, 103 | } 104 | 105 | pub type WorkerUpdateSender = Sender; 106 | pub type WorkerUpdateReceiver = Receiver; 107 | pub type WorkerUpdateChannel = (WorkerUpdateSender, WorkerUpdateReceiver); 108 | 109 | #[derive(Debug, Clone, Copy)] 110 | pub struct WorkerStatusUpdate { 111 | pub status: Option, 112 | pub slot_delta: Option<(usize, bool)>, 113 | } 114 | 115 | #[derive(Debug, Clone, Copy, PartialEq)] 116 | pub enum SlotStatus { 117 | Empty, 118 | Requested, 119 | } 120 | 121 | #[allow(clippy::clippy::too_many_arguments)] 122 | pub fn build_config( 123 | speed: usize, 124 | qp: usize, 125 | max_bitrate: Option, 126 | tiles: usize, 127 | video_info: VideoDetails, 128 | pool: Arc, 129 | color_primaries: ColorPrimaries, 130 | transfer_characteristics: TransferCharacteristics, 131 | matrix_coefficients: MatrixCoefficients, 132 | ) -> Config { 133 | Config::new() 134 | .with_encoder_config(build_encoder_config( 135 | speed, 136 | qp, 137 | max_bitrate, 138 | tiles, 139 | video_info, 140 | color_primaries, 141 | transfer_characteristics, 142 | matrix_coefficients, 143 | )) 144 | .with_thread_pool(pool) 145 | } 146 | 147 | #[allow(clippy::clippy::too_many_arguments)] 148 | pub fn build_encoder_config( 149 | speed: usize, 150 | qp: usize, 151 | max_bitrate: Option, 152 | tiles: usize, 153 | video_info: VideoDetails, 154 | color_primaries: ColorPrimaries, 155 | transfer_characteristics: TransferCharacteristics, 156 | matrix_coefficients: MatrixCoefficients, 157 | ) -> EncoderConfig { 158 | let mut enc_config = EncoderConfig::with_speed_preset(speed); 159 | enc_config.width = video_info.width; 160 | enc_config.height = video_info.height; 161 | enc_config.bit_depth = video_info.bit_depth; 162 | enc_config.chroma_sampling = video_info.chroma_sampling; 163 | enc_config.chroma_sample_position = video_info.chroma_sample_position; 164 | enc_config.time_base = video_info.time_base; 165 | enc_config.tiles = tiles; 166 | enc_config.min_key_frame_interval = 0; 167 | enc_config.max_key_frame_interval = u16::max_value() as u64; 168 | enc_config.speed_settings.no_scene_detection = true; 169 | enc_config.color_description = if color_primaries == ColorPrimaries::Unspecified 170 | && transfer_characteristics == TransferCharacteristics::Unspecified 171 | && matrix_coefficients == MatrixCoefficients::Unspecified 172 | { 173 | // No need to set a color description with all parameters unspecified. 174 | None 175 | } else { 176 | Some(ColorDescription { 177 | color_primaries, 178 | transfer_characteristics, 179 | matrix_coefficients, 180 | }) 181 | }; 182 | 183 | if let Some(max_bitrate) = max_bitrate { 184 | enc_config.min_quantizer = qp as u8; 185 | enc_config.bitrate = max_bitrate; 186 | } else { 187 | enc_config.quantizer = qp; 188 | } 189 | 190 | enc_config 191 | } 192 | 193 | #[derive(Debug, Clone)] 194 | pub enum Output { 195 | File(PathBuf), 196 | Null, 197 | Memory, 198 | } 199 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/bin/remote.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | fs::File, 3 | io::{sink, Write}, 4 | thread::sleep, 5 | time::Duration, 6 | }; 7 | 8 | use anyhow::Result; 9 | use crossbeam_channel::unbounded; 10 | use log::{debug, error}; 11 | use rav1e_by_gop::{ 12 | ActiveConnection, 13 | GetInfoResponse, 14 | GetProgressResponse, 15 | Output, 16 | ProgressStatus, 17 | Slot, 18 | SlotStatus, 19 | WorkerStatusUpdate, 20 | WorkerUpdateChannel, 21 | }; 22 | use reqwest::StatusCode; 23 | use serde::de::DeserializeOwned; 24 | use url::Url; 25 | use v_frame::pixel::Pixel; 26 | 27 | use crate::{analyze::SlotReadySender, WorkerConfig, CLIENT}; 28 | 29 | pub(crate) struct RemoteWorkerInfo { 30 | pub uri: Url, 31 | pub password: String, 32 | // Like local slots, this tracks which workers are active or inactive 33 | pub workers: Box<[bool]>, 34 | pub slot_status: SlotStatus, 35 | pub update_channel: WorkerUpdateChannel, 36 | } 37 | 38 | pub(crate) fn discover_remote_worker(worker: &WorkerConfig) -> Result { 39 | // The http crate has a crap interface for setting the port (i.e. "no interface"), 40 | // so do it this kind of stupid way using two crates instead. 41 | let mut url = Url::parse(&if worker.secure { 42 | format!("https://{}", worker.host) 43 | } else { 44 | format!("http://{}", worker.host) 45 | })?; 46 | url.set_port(worker.port.or(Some(13415))).unwrap(); 47 | 48 | debug!("Initial query to remote worker {}", url); 49 | let resp: GetInfoResponse = CLIENT 50 | .get(&format!("{}{}", url, "info")) 51 | .header("X-RAV1E-AUTH", &worker.password) 52 | .send()? 53 | .json()?; 54 | 55 | Ok(RemoteWorkerInfo { 56 | uri: url, 57 | password: worker.password.clone(), 58 | workers: vec![false; resp.worker_count].into_boxed_slice(), 59 | slot_status: SlotStatus::Empty, 60 | update_channel: unbounded(), 61 | }) 62 | } 63 | 64 | pub(crate) fn wait_for_slot_allocation( 65 | connection: ActiveConnection, 66 | slot_ready_sender: SlotReadySender, 67 | ) { 68 | loop { 69 | if let Ok(response) = CLIENT 70 | .get(&format!( 71 | "{}{}/{}", 72 | &connection.worker_uri, "enqueue", connection.request_id 73 | )) 74 | .header("X-RAV1E-AUTH", &connection.worker_password) 75 | .send() 76 | .and_then(|res| res.error_for_status()) 77 | { 78 | if response.status() == StatusCode::ACCEPTED { 79 | debug!("Slot reserved, sending to encoder"); 80 | let update_sender = connection.worker_update_sender.clone(); 81 | let slot_in_worker = connection.slot_in_worker; 82 | if slot_ready_sender 83 | .send(Slot::Remote(Box::new(connection))) 84 | .is_ok() 85 | { 86 | debug!("Slot allocated and sent to encoder"); 87 | update_sender 88 | .send(WorkerStatusUpdate { 89 | status: None, 90 | slot_delta: Some((slot_in_worker, true)), 91 | }) 92 | .unwrap(); 93 | } else { 94 | debug!("Ignoring final slot"); 95 | } 96 | return; 97 | } 98 | } 99 | sleep(Duration::from_secs(5)); 100 | } 101 | } 102 | 103 | pub(crate) fn remote_encode_segment( 104 | connection: ActiveConnection, 105 | ) { 106 | loop { 107 | if let Ok(response) = CLIENT 108 | .get(&format!( 109 | "{}{}/{}", 110 | &connection.worker_uri, "segment", connection.request_id 111 | )) 112 | .header("X-RAV1E-AUTH", &connection.worker_password) 113 | .send() 114 | .and_then(|res| res.error_for_status()) 115 | .and_then(|res| res.json::()) 116 | { 117 | let _ = connection 118 | .progress_sender 119 | .send(ProgressStatus::Encoding(Box::new( 120 | (&response.progress).into(), 121 | ))); 122 | if response.done { 123 | loop { 124 | match CLIENT 125 | .get(&format!( 126 | "{}{}/{}", 127 | &connection.worker_uri, "segment_data", connection.request_id 128 | )) 129 | .header("X-RAV1E-AUTH", &connection.worker_password) 130 | .send() 131 | .and_then(|res| res.error_for_status()) 132 | { 133 | Ok(response) => { 134 | let mut output: Box = 135 | match connection.encode_info.unwrap().output_file { 136 | Output::File(filename) => { 137 | Box::new(File::create(filename).unwrap()) 138 | } 139 | Output::Null => Box::new(sink()), 140 | Output::Memory => unreachable!(), 141 | }; 142 | output.write_all(&response.bytes().unwrap()).unwrap(); 143 | output.flush().unwrap(); 144 | break; 145 | } 146 | Err(e) => { 147 | error!("Failed to get encode from remote server: {}", e); 148 | sleep(Duration::from_secs(5)); 149 | } 150 | }; 151 | } 152 | break; 153 | } 154 | } 155 | sleep(Duration::from_secs(2)); 156 | } 157 | 158 | connection 159 | .worker_update_sender 160 | .send(WorkerStatusUpdate { 161 | status: None, 162 | slot_delta: Some((connection.slot_in_worker, false)), 163 | }) 164 | .unwrap(); 165 | } 166 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/encode/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod stats; 2 | 3 | use std::{collections::BTreeSet, fs, fs::File, io::BufReader, path::PathBuf, sync::Arc}; 4 | 5 | use anyhow::Result; 6 | use crossbeam_channel::{Receiver, Sender}; 7 | use rav1e::prelude::*; 8 | use serde::{de::DeserializeOwned, Deserialize, Serialize}; 9 | use systemstat::data::ByteSize; 10 | use threadpool::ThreadPool; 11 | 12 | pub use self::stats::*; 13 | use super::VideoDetails; 14 | use crate::{ 15 | build_config, 16 | decompress_frame, 17 | muxer::create_muxer, 18 | Output, 19 | SegmentData, 20 | SegmentFrameData, 21 | }; 22 | 23 | #[derive(Debug, Clone, Copy, Serialize, Deserialize)] 24 | pub struct EncodeOptions { 25 | pub speed: usize, 26 | pub qp: usize, 27 | pub max_bitrate: Option, 28 | pub tiles: usize, 29 | pub color_primaries: ColorPrimaries, 30 | pub transfer_characteristics: TransferCharacteristics, 31 | pub matrix_coefficients: MatrixCoefficients, 32 | } 33 | 34 | pub fn encode_segment( 35 | opts: EncodeOptions, 36 | video_info: VideoDetails, 37 | data: SegmentData, 38 | thread_pool: &mut ThreadPool, 39 | rayon_pool: Arc, 40 | progress_sender: ProgressSender, 41 | segment_output_file: Output, 42 | ) -> Result<()> { 43 | let progress = ProgressInfo::new( 44 | Rational { 45 | num: video_info.time_base.den, 46 | den: video_info.time_base.num, 47 | }, 48 | match data.frame_data { 49 | SegmentFrameData::Y4MFile { frame_count, .. } => frame_count, 50 | SegmentFrameData::CompressedFrames(ref frames) => frames.len(), 51 | }, 52 | { 53 | let mut kf = BTreeSet::new(); 54 | kf.insert(data.start_frameno); 55 | kf 56 | }, 57 | data.segment_no + 1, 58 | data.next_analysis_frame, 59 | None, 60 | ); 61 | let _ = progress_sender.send(ProgressStatus::Encoding(Box::new(progress.clone()))); 62 | 63 | thread_pool.execute(move || { 64 | let source = Source { 65 | frame_data: match data.frame_data { 66 | SegmentFrameData::Y4MFile { path, frame_count } => SourceFrameData::Y4MFile { 67 | frame_count, 68 | video_info, 69 | input: { 70 | let file = File::open(&path).unwrap(); 71 | BufReader::new(file) 72 | }, 73 | path, 74 | }, 75 | SegmentFrameData::CompressedFrames(frames) => { 76 | SourceFrameData::CompressedFrames(frames) 77 | } 78 | }, 79 | sent_count: 0, 80 | }; 81 | if video_info.bit_depth > 8 { 82 | do_encode::( 83 | rayon_pool, 84 | opts, 85 | video_info, 86 | source, 87 | &segment_output_file, 88 | progress, 89 | progress_sender, 90 | ) 91 | .expect("Failed encoding segment"); 92 | } else { 93 | do_encode::( 94 | rayon_pool, 95 | opts, 96 | video_info, 97 | source, 98 | &segment_output_file, 99 | progress, 100 | progress_sender, 101 | ) 102 | .expect("Failed encoding segment"); 103 | } 104 | }); 105 | Ok(()) 106 | } 107 | 108 | fn do_encode( 109 | pool: Arc, 110 | opts: EncodeOptions, 111 | video_info: VideoDetails, 112 | mut source: Source, 113 | segment_output_file: &Output, 114 | mut progress: ProgressInfo, 115 | progress_sender: ProgressSender, 116 | ) -> Result { 117 | let cfg = build_config( 118 | opts.speed, 119 | opts.qp, 120 | opts.max_bitrate, 121 | opts.tiles, 122 | video_info, 123 | pool, 124 | opts.color_primaries, 125 | opts.transfer_characteristics, 126 | opts.matrix_coefficients, 127 | ); 128 | 129 | let mut ctx: Context = cfg.new_context()?; 130 | let _ = progress_sender.send(ProgressStatus::Encoding(Box::new(progress.clone()))); 131 | 132 | let mut output = create_muxer(&segment_output_file).expect("Failed to create segment output"); 133 | loop { 134 | match process_frame(&mut ctx, &mut source)? { 135 | ProcessFrameResult::Packet(packet) => { 136 | output.write_frame( 137 | packet.input_frameno as u64, 138 | packet.data.as_ref(), 139 | packet.frame_type, 140 | ); 141 | progress.add_packet(*packet); 142 | let _ = progress_sender.send(ProgressStatus::Encoding(Box::new(progress.clone()))); 143 | } 144 | ProcessFrameResult::NoPacket(_) => { 145 | // Next iteration 146 | } 147 | ProcessFrameResult::EndOfSegment => { 148 | output.flush().unwrap(); 149 | break; 150 | } 151 | }; 152 | } 153 | if let SourceFrameData::Y4MFile { path, .. } = source.frame_data { 154 | fs::remove_file(&path).unwrap(); 155 | } 156 | 157 | Ok(progress) 158 | } 159 | 160 | pub enum SourceFrameData { 161 | CompressedFrames(Vec>), 162 | Y4MFile { 163 | path: PathBuf, 164 | input: BufReader, 165 | frame_count: usize, 166 | video_info: VideoDetails, 167 | }, 168 | } 169 | 170 | pub struct Source { 171 | pub sent_count: usize, 172 | pub frame_data: SourceFrameData, 173 | } 174 | 175 | impl Source { 176 | fn read_frame(&mut self, ctx: &mut Context) { 177 | if self.sent_count == self.frame_count() { 178 | ctx.flush(); 179 | return; 180 | } 181 | 182 | match self.frame_data { 183 | SourceFrameData::CompressedFrames(ref mut frames) => { 184 | let _ = ctx.send_frame(Some(Arc::new(decompress_frame(&frames[self.sent_count])))); 185 | // Deallocate the compressed frame from memory, we no longer need it 186 | frames[self.sent_count] = Vec::new(); 187 | } 188 | SourceFrameData::Y4MFile { ref mut input, .. } => { 189 | let _ = ctx.send_frame(Some(Arc::new(bincode::deserialize_from(input).unwrap()))); 190 | } 191 | }; 192 | self.sent_count += 1; 193 | } 194 | 195 | pub fn frame_count(&self) -> usize { 196 | match self.frame_data { 197 | SourceFrameData::CompressedFrames(ref frames) => frames.len(), 198 | SourceFrameData::Y4MFile { frame_count, .. } => frame_count, 199 | } 200 | } 201 | } 202 | 203 | pub enum ProcessFrameResult { 204 | Packet(Box>), 205 | NoPacket(bool), 206 | EndOfSegment, 207 | } 208 | 209 | pub fn process_frame( 210 | ctx: &mut Context, 211 | source: &mut Source, 212 | ) -> Result> { 213 | let pkt_wrapped = ctx.receive_packet(); 214 | match pkt_wrapped { 215 | Ok(pkt) => Ok(ProcessFrameResult::Packet(Box::new(pkt))), 216 | Err(EncoderStatus::NeedMoreData) => { 217 | source.read_frame(ctx); 218 | Ok(ProcessFrameResult::NoPacket(false)) 219 | } 220 | Err(EncoderStatus::EnoughData) => { 221 | unreachable!(); 222 | } 223 | Err(EncoderStatus::LimitReached) => Ok(ProcessFrameResult::EndOfSegment), 224 | e @ Err(EncoderStatus::Failure) => { 225 | e?; 226 | unreachable!(); 227 | } 228 | Err(EncoderStatus::NotReady) => { 229 | unreachable!(); 230 | } 231 | Err(EncoderStatus::Encoded) => Ok(ProcessFrameResult::NoPacket(true)), 232 | } 233 | } 234 | 235 | pub type ProgressSender = Sender; 236 | pub type ProgressReceiver = Receiver; 237 | pub type ProgressChannel = (ProgressSender, ProgressReceiver); 238 | 239 | pub enum ProgressStatus { 240 | Idle, 241 | Loading, 242 | Compressing(usize), 243 | Sending(ByteSize), 244 | Encoding(Box), 245 | } 246 | -------------------------------------------------------------------------------- /rav1e-worker/src/worker.rs: -------------------------------------------------------------------------------- 1 | use std::{collections::BTreeSet, fs, fs::File, io::BufReader, sync::Arc, time::Duration}; 2 | 3 | use chrono::Utc; 4 | use log::{error, info}; 5 | use rav1e::prelude::*; 6 | use rav1e_by_gop::*; 7 | use serde::{de::DeserializeOwned, Serialize}; 8 | use tokio::time::sleep; 9 | use uuid::Uuid; 10 | use v_frame::pixel::Pixel; 11 | 12 | use crate::ENCODER_QUEUE; 13 | 14 | pub async fn start_workers(worker_threads: usize) { 15 | info!("Starting {} workers", worker_threads); 16 | let rayon_pool = Arc::new( 17 | rayon::ThreadPoolBuilder::new() 18 | .num_threads(worker_threads) 19 | .build() 20 | .unwrap(), 21 | ); 22 | 23 | // This thread watches the slot request queue 24 | // and allocates slots when they are available. 25 | tokio::spawn(async move { 26 | loop { 27 | sleep(Duration::from_secs(3)).await; 28 | 29 | let reader = ENCODER_QUEUE.read(); 30 | let mut in_progress_items = 0; 31 | for item in reader.values() { 32 | match item.read().state { 33 | EncodeState::Enqueued => (), 34 | _ => { 35 | in_progress_items += 1; 36 | } 37 | } 38 | } 39 | if in_progress_items >= worker_threads { 40 | continue; 41 | } 42 | 43 | for (&request_id, item) in reader.iter() { 44 | let mut item_handle = item.write(); 45 | match item_handle.state { 46 | EncodeState::Enqueued if in_progress_items < worker_threads => { 47 | info!("A slot is ready for request {}", request_id); 48 | item_handle.state = EncodeState::AwaitingInfo { 49 | time_ready: Utc::now(), 50 | }; 51 | in_progress_items += 1; 52 | } 53 | EncodeState::Ready { ref raw_frames, .. } => { 54 | info!("Beginning encode for request {}", request_id); 55 | let video_info = item_handle.video_info; 56 | let options = item_handle.options; 57 | let raw_frames = raw_frames.clone(); 58 | let pool_handle = rayon_pool.clone(); 59 | tokio::spawn(async move { 60 | if video_info.bit_depth <= 8 { 61 | encode_segment::( 62 | request_id, 63 | video_info, 64 | options, 65 | raw_frames, 66 | pool_handle, 67 | ) 68 | .await; 69 | } else { 70 | encode_segment::( 71 | request_id, 72 | video_info, 73 | options, 74 | raw_frames, 75 | pool_handle, 76 | ) 77 | .await; 78 | } 79 | }); 80 | } 81 | _ => (), 82 | } 83 | } 84 | } 85 | }); 86 | } 87 | 88 | pub async fn encode_segment( 89 | request_id: Uuid, 90 | video_info: VideoDetails, 91 | options: EncodeOptions, 92 | input: Arc, 93 | pool: Arc, 94 | ) { 95 | { 96 | let queue_handle = ENCODER_QUEUE.read(); 97 | let mut item_handle = queue_handle.get(&request_id).unwrap().write(); 98 | item_handle.state = EncodeState::InProgress { 99 | progress: ProgressInfo::new( 100 | Rational::from_reciprocal(item_handle.video_info.time_base), 101 | match input.as_ref() { 102 | SegmentFrameData::CompressedFrames(frames) => frames.len(), 103 | SegmentFrameData::Y4MFile { frame_count, .. } => *frame_count, 104 | }, 105 | { 106 | let mut keyframes = BTreeSet::new(); 107 | keyframes.insert(match item_handle.state { 108 | EncodeState::Ready { 109 | keyframe_number, .. 110 | } => keyframe_number, 111 | _ => unreachable!(), 112 | }); 113 | keyframes 114 | }, 115 | match item_handle.state { 116 | EncodeState::Ready { segment_idx, .. } => segment_idx, 117 | _ => unreachable!(), 118 | }, 119 | match item_handle.state { 120 | EncodeState::Ready { 121 | next_analysis_frame, 122 | .. 123 | } => next_analysis_frame, 124 | _ => unreachable!(), 125 | }, 126 | None, 127 | ), 128 | }; 129 | } 130 | 131 | let mut source = Source { 132 | frame_data: match Arc::try_unwrap(input) { 133 | Ok(SegmentFrameData::CompressedFrames(input)) => { 134 | SourceFrameData::CompressedFrames(input) 135 | } 136 | Ok(SegmentFrameData::Y4MFile { 137 | frame_count, 138 | ref path, 139 | }) => SourceFrameData::Y4MFile { 140 | frame_count, 141 | input: { 142 | let file = File::open(&path).unwrap(); 143 | BufReader::new(file) 144 | }, 145 | path: path.to_path_buf(), 146 | video_info, 147 | }, 148 | Err(_) => unreachable!("input should only have one reference at this point"), 149 | }, 150 | sent_count: 0, 151 | }; 152 | let cfg = build_config( 153 | options.speed, 154 | options.qp, 155 | options.max_bitrate, 156 | options.tiles, 157 | video_info, 158 | pool, 159 | options.color_primaries, 160 | options.transfer_characteristics, 161 | options.matrix_coefficients, 162 | ); 163 | 164 | let mut ctx: Context = match cfg.new_context() { 165 | Ok(ctx) => ctx, 166 | Err(e) => { 167 | error!( 168 | "Failed to create encode context for request {}: {}", 169 | request_id, e 170 | ); 171 | return; 172 | } 173 | }; 174 | let mut output = IvfMuxer::>::in_memory(); 175 | loop { 176 | match process_frame(&mut ctx, &mut source) { 177 | Ok(ProcessFrameResult::Packet(packet)) => { 178 | let queue_handle = ENCODER_QUEUE.read(); 179 | let mut item_handle = queue_handle.get(&request_id).unwrap().write(); 180 | if let EncodeState::InProgress { ref mut progress } = item_handle.state { 181 | output.write_frame( 182 | packet.input_frameno as u64, 183 | packet.data.as_ref(), 184 | packet.frame_type, 185 | ); 186 | progress.add_packet(*packet); 187 | } else { 188 | unreachable!(); 189 | } 190 | } 191 | Ok(ProcessFrameResult::NoPacket(_)) => { 192 | // Next iteration 193 | } 194 | Ok(ProcessFrameResult::EndOfSegment) => { 195 | break; 196 | } 197 | Err(e) => { 198 | error!("Encoding error for request {}: {}", request_id, e); 199 | if let SourceFrameData::Y4MFile { path, .. } = source.frame_data { 200 | let _ = fs::remove_file(path); 201 | }; 202 | return; 203 | } 204 | }; 205 | } 206 | 207 | { 208 | let queue_handle = ENCODER_QUEUE.read(); 209 | let mut item_handle = queue_handle.get(&request_id).unwrap().write(); 210 | item_handle.state = if let EncodeState::InProgress { ref progress } = item_handle.state { 211 | EncodeState::EncodingDone { 212 | progress: progress.clone(), 213 | encoded_data: output.output, 214 | time_finished: Utc::now(), 215 | } 216 | } else { 217 | unreachable!() 218 | }; 219 | } 220 | if let SourceFrameData::Y4MFile { path, .. } = source.frame_data { 221 | let _ = fs::remove_file(path); 222 | }; 223 | info!("Segment {} finished", request_id); 224 | } 225 | -------------------------------------------------------------------------------- /rav1e-worker/src/server/routes.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | fs::File, 3 | io::{BufWriter, Write}, 4 | path::PathBuf, 5 | sync::Arc, 6 | }; 7 | 8 | use byteorder::{LittleEndian, WriteBytesExt}; 9 | use bytes::Bytes; 10 | use chrono::Utc; 11 | use http::header::{HeaderValue, CONTENT_TYPE}; 12 | use parking_lot::RwLock; 13 | use rav1e_by_gop::{ 14 | decompress_frame, 15 | EncodeState, 16 | GetInfoResponse, 17 | GetProgressResponse, 18 | PostEnqueueResponse, 19 | PostSegmentMessage, 20 | SegmentFrameData, 21 | SerializableProgressInfo, 22 | SlotRequestMessage, 23 | }; 24 | use uuid::{v1::Timestamp, Uuid}; 25 | use warp::{http::StatusCode, reply::Response, Filter, Rejection, Reply}; 26 | 27 | use crate::{ 28 | server::{ 29 | helpers::{ 30 | handle_rejection_types, 31 | json_body, 32 | require_auth, 33 | with_state, 34 | ClientVersionMismatch, 35 | }, 36 | CLIENT_VERSION_REQUIRED, 37 | }, 38 | try_or_500, 39 | EncodeItem, 40 | ENCODER_QUEUE, 41 | UUID_CONTEXT, 42 | UUID_NODE_ID, 43 | }; 44 | 45 | pub fn get_routes( 46 | temp_dir: Option, 47 | worker_threads: usize, 48 | ) -> impl Filter + Clone { 49 | warp::path!("info") 50 | .and(warp::get()) 51 | .and(require_auth()) 52 | .and(with_state(worker_threads)) 53 | .and_then(get_info) 54 | .or(warp::path!("enqueue" / Uuid) 55 | .and(warp::get()) 56 | .and(require_auth()) 57 | .and_then(get_enqueue)) 58 | .or(warp::path!("enqueue") 59 | .and(warp::post()) 60 | .and(require_auth()) 61 | .and(json_body()) 62 | .and_then(post_enqueue)) 63 | .or(warp::path!("segment" / Uuid) 64 | .and(warp::post()) 65 | .and(require_auth()) 66 | .and(json_body()) 67 | .and_then(post_segment)) 68 | .or(warp::path!("segment_data" / Uuid) 69 | .and(warp::post()) 70 | .and(require_auth()) 71 | .and(warp::body::bytes()) 72 | .and(with_state(temp_dir)) 73 | .and_then(post_segment_data)) 74 | .or(warp::path!("segment" / Uuid) 75 | .and(warp::get()) 76 | .and(require_auth()) 77 | .and_then(get_segment)) 78 | .or(warp::path!("segment_data" / Uuid) 79 | .and(warp::get()) 80 | .and(require_auth()) 81 | .and_then(get_segment_data)) 82 | .recover(handle_rejection_types) 83 | } 84 | 85 | async fn get_info(_auth: (), worker_threads: usize) -> Result { 86 | Ok(warp::reply::with_status( 87 | warp::reply::json(&GetInfoResponse { 88 | worker_count: worker_threads, 89 | }), 90 | StatusCode::NOT_FOUND, 91 | )) 92 | } 93 | 94 | // this endpoint tells a client if their slot is ready 95 | async fn get_enqueue(request_id: Uuid, _auth: ()) -> Result { 96 | let reader = ENCODER_QUEUE.read(); 97 | let item = reader.get(&request_id); 98 | if let Some(item) = item { 99 | match item.read().state { 100 | EncodeState::Enqueued => Ok(warp::reply::with_status( 101 | warp::reply::json(&()), 102 | StatusCode::ACCEPTED, 103 | )), 104 | EncodeState::AwaitingInfo { .. } | EncodeState::AwaitingData { .. } => Ok( 105 | warp::reply::with_status(warp::reply::json(&()), StatusCode::OK), 106 | ), 107 | _ => Ok(warp::reply::with_status( 108 | warp::reply::json(&()), 109 | StatusCode::GONE, 110 | )), 111 | } 112 | } else { 113 | Ok(warp::reply::with_status( 114 | warp::reply::json(&()), 115 | StatusCode::NOT_FOUND, 116 | )) 117 | } 118 | } 119 | 120 | // client hits this endpoint to say that they want to request a slot 121 | // returns a JSON body with a request ID 122 | async fn post_enqueue(_auth: (), body: SlotRequestMessage) -> Result { 123 | if !CLIENT_VERSION_REQUIRED.matches(&body.client_version) { 124 | return Err(warp::reject::custom(ClientVersionMismatch)); 125 | } 126 | 127 | let ts = Timestamp::from_unix(&*UUID_CONTEXT, 1497624119, 1234); 128 | let request_id = try_or_500!(Uuid::new_v1(ts, &UUID_NODE_ID)); 129 | ENCODER_QUEUE.write().insert( 130 | request_id, 131 | RwLock::new(EncodeItem::new(body.options, body.video_info)), 132 | ); 133 | 134 | Ok(warp::reply::with_status( 135 | warp::reply::json(&PostEnqueueResponse { request_id }), 136 | StatusCode::ACCEPTED, 137 | )) 138 | } 139 | 140 | async fn post_segment( 141 | request_id: Uuid, 142 | _auth: (), 143 | body: PostSegmentMessage, 144 | ) -> Result { 145 | if let Some(item) = ENCODER_QUEUE.read().get(&request_id) { 146 | let mut item_handle = item.write(); 147 | match item_handle.state { 148 | EncodeState::AwaitingInfo { .. } => (), 149 | _ => { 150 | return Ok(warp::reply::with_status( 151 | warp::reply::json(&()), 152 | StatusCode::GONE, 153 | )); 154 | } 155 | }; 156 | item_handle.state = EncodeState::AwaitingData { 157 | keyframe_number: body.keyframe_number, 158 | segment_idx: body.segment_idx, 159 | next_analysis_frame: body.next_analysis_frame, 160 | time_ready: Utc::now(), 161 | }; 162 | return Ok(warp::reply::with_status( 163 | warp::reply::json(&()), 164 | StatusCode::OK, 165 | )); 166 | } 167 | 168 | Ok(warp::reply::with_status( 169 | warp::reply::json(&()), 170 | StatusCode::NOT_FOUND, 171 | )) 172 | } 173 | 174 | // client sends raw video data via this endpoint 175 | async fn post_segment_data( 176 | request_id: Uuid, 177 | _auth: (), 178 | body: Bytes, 179 | temp_dir: Option, 180 | ) -> Result { 181 | if let Some(item) = ENCODER_QUEUE.read().get(&request_id) { 182 | let mut item_handle = item.write(); 183 | let frame_data; 184 | let keyframe_number_outer; 185 | let next_analysis_frame_outer; 186 | let segment_idx_outer; 187 | if let EncodeState::AwaitingData { 188 | keyframe_number, 189 | next_analysis_frame, 190 | segment_idx, 191 | .. 192 | } = &mut item_handle.state 193 | { 194 | keyframe_number_outer = *keyframe_number; 195 | next_analysis_frame_outer = *next_analysis_frame; 196 | segment_idx_outer = *segment_idx; 197 | let compressed_frames: Vec> = try_or_500!(bincode::deserialize(&body)); 198 | frame_data = match temp_dir { 199 | Some(temp_dir) => { 200 | let frame_count = compressed_frames.len(); 201 | let mut temp_path = temp_dir; 202 | temp_path.push(request_id.to_hyphenated().to_string()); 203 | temp_path.set_extension("py4m"); 204 | let file = File::create(&temp_path).unwrap(); 205 | let mut writer = BufWriter::new(file); 206 | for frame in compressed_frames { 207 | let frame_data = if item_handle.video_info.bit_depth == 8 { 208 | bincode::serialize(&decompress_frame::(&frame)).unwrap() 209 | } else { 210 | bincode::serialize(&decompress_frame::(&frame)).unwrap() 211 | }; 212 | writer 213 | .write_u32::(frame_data.len() as u32) 214 | .unwrap(); 215 | writer.write_all(&frame_data).unwrap(); 216 | } 217 | writer.flush().unwrap(); 218 | SegmentFrameData::Y4MFile { 219 | path: temp_path, 220 | frame_count, 221 | } 222 | } 223 | None => SegmentFrameData::CompressedFrames(compressed_frames), 224 | } 225 | } else { 226 | return Ok(warp::reply::with_status( 227 | warp::reply::json(&()), 228 | StatusCode::NOT_FOUND, 229 | )); 230 | } 231 | item_handle.state = EncodeState::Ready { 232 | keyframe_number: keyframe_number_outer, 233 | next_analysis_frame: next_analysis_frame_outer, 234 | segment_idx: segment_idx_outer, 235 | raw_frames: Arc::new(frame_data), 236 | }; 237 | return Ok(warp::reply::with_status( 238 | warp::reply::json(&()), 239 | StatusCode::OK, 240 | )); 241 | } 242 | 243 | Ok(warp::reply::with_status( 244 | warp::reply::json(&()), 245 | StatusCode::NOT_FOUND, 246 | )) 247 | } 248 | 249 | // returns progress on currently encoded segment 250 | async fn get_segment(request_id: Uuid, _auth: ()) -> Result { 251 | if let Some(item) = ENCODER_QUEUE.read().get(&request_id) { 252 | let item_reader = item.read(); 253 | match item_reader.state { 254 | EncodeState::InProgress { ref progress, .. } => Ok(warp::reply::with_status( 255 | warp::reply::json(&GetProgressResponse { 256 | progress: SerializableProgressInfo::from(&*progress), 257 | done: false, 258 | }), 259 | StatusCode::OK, 260 | )), 261 | EncodeState::EncodingDone { ref progress, .. } => Ok(warp::reply::with_status( 262 | warp::reply::json(&GetProgressResponse { 263 | progress: SerializableProgressInfo::from(&*progress), 264 | done: true, 265 | }), 266 | StatusCode::OK, 267 | )), 268 | _ => Ok(warp::reply::with_status( 269 | warp::reply::json(&()), 270 | StatusCode::GONE, 271 | )), 272 | } 273 | } else { 274 | Ok(warp::reply::with_status( 275 | warp::reply::json(&()), 276 | StatusCode::NOT_FOUND, 277 | )) 278 | } 279 | } 280 | 281 | // if segment is ready, sends client the encoded video data 282 | async fn get_segment_data(request_id: Uuid, _auth: ()) -> Result { 283 | // Check first without mutating the state 284 | let mut can_send_data = false; 285 | if let Some(item) = ENCODER_QUEUE.read().get(&request_id) { 286 | if let EncodeState::EncodingDone { .. } = item.read().state { 287 | can_send_data = true; 288 | } 289 | } 290 | if !can_send_data { 291 | return Ok(warp::reply::with_status( 292 | Response::new(Vec::new().into()), 293 | StatusCode::NOT_FOUND, 294 | )); 295 | } 296 | 297 | // Now pop it from the queue and send it, freeing the resources simultaneously 298 | let item = { 299 | let mut queue_handle = ENCODER_QUEUE.write(); 300 | queue_handle.remove(&request_id) 301 | }; 302 | if let Some(item) = item { 303 | if let EncodeState::EncodingDone { encoded_data, .. } = item.into_inner().state { 304 | return Ok(warp::reply::with_status( 305 | { 306 | let mut response = Response::new(encoded_data.into()); 307 | response.headers_mut().insert( 308 | CONTENT_TYPE, 309 | HeaderValue::from_static("application/octet-stream"), 310 | ); 311 | response 312 | }, 313 | StatusCode::OK, 314 | )); 315 | } 316 | } 317 | 318 | unreachable!() 319 | } 320 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/bin/progress.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | cmp, 3 | fs::File, 4 | io::Write, 5 | path::{Path, PathBuf}, 6 | sync::{ 7 | atomic::{AtomicBool, Ordering}, 8 | Arc, 9 | }, 10 | thread, 11 | time::Duration, 12 | }; 13 | 14 | use clap::ArgMatches; 15 | use console::{style, StyledObject, Term}; 16 | use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; 17 | use log::{debug, error, trace}; 18 | #[cfg(feature = "remote")] 19 | use parking_lot::Mutex; 20 | use rav1e_by_gop::*; 21 | 22 | use crate::analyze::InputFinishedReceiver; 23 | #[cfg(feature = "remote")] 24 | use crate::remote::RemoteWorkerInfo; 25 | 26 | pub fn update_progress_file(output: &Output, progress: &ProgressInfo) { 27 | if let Output::File(output) = output { 28 | let data = rmp_serde::to_vec(&stats::SerializableProgressInfo::from(progress)) 29 | .expect("Failed to serialize data"); 30 | let mut progress_file = 31 | File::create(get_progress_filename(&output)).expect("Failed to open progress file"); 32 | progress_file 33 | .write_all(&data) 34 | .expect("Failed to write to progress file"); 35 | } 36 | } 37 | 38 | pub fn get_segment_input_filename(output: &Path, segment_idx: usize) -> PathBuf { 39 | // We're making a "pseudo y4m" file that is actually `Frame` data. 40 | output.with_extension(&format!("part{}.py4m", segment_idx)) 41 | } 42 | 43 | pub fn get_segment_output_filename(output: &Path, segment_idx: usize) -> PathBuf { 44 | output.with_extension(&format!("part{}.ivf", segment_idx)) 45 | } 46 | 47 | pub fn load_progress_file(outfile: &Output, matches: &ArgMatches) -> Option { 48 | let term_err = Term::stderr(); 49 | match outfile { 50 | Output::File(outfile) if get_progress_filename(outfile).is_file() => { 51 | let resume = if matches.is_present("FORCE_RESUME") { 52 | true 53 | } else if matches.is_present("FORCE_OVERWRITE") { 54 | false 55 | } else if term_err.features().is_attended() && matches.value_of("INPUT") != Some("-") { 56 | let resume; 57 | loop { 58 | let input = dialoguer::Input::::new() 59 | .with_prompt(&format!( 60 | "Found progress file for this encode. [{}]esume or [{}]verwrite?", 61 | style("R").cyan(), 62 | style("O").cyan() 63 | )) 64 | .interact() 65 | .unwrap(); 66 | match input.to_lowercase().as_str() { 67 | "r" | "resume" => { 68 | resume = true; 69 | break; 70 | } 71 | "o" | "overwrite" => { 72 | resume = false; 73 | break; 74 | } 75 | _ => { 76 | error!("Input not recognized"); 77 | } 78 | }; 79 | } 80 | resume 81 | } else { 82 | // Assume we want to resume if this is not a TTY 83 | // and no CLI option is given 84 | true 85 | }; 86 | if resume { 87 | let progress_file = File::open(get_progress_filename(outfile)) 88 | .expect("Failed to open progress file"); 89 | let progress_input: SerializableProgressInfo = rmp_serde::from_read(&progress_file) 90 | .expect("Progress file did not contain valid JSON"); 91 | Some(ProgressInfo::from(&progress_input)) 92 | } else { 93 | None 94 | } 95 | } 96 | _ => None, 97 | } 98 | } 99 | 100 | pub fn get_progress_filename(output: &Path) -> PathBuf { 101 | output.with_extension("progress.data") 102 | } 103 | 104 | pub(crate) fn watch_progress_receivers( 105 | receivers: Vec, 106 | slots: Arc>, 107 | #[cfg(feature = "remote")] remote_slots: Arc>>, 108 | output_file: Output, 109 | verbose: bool, 110 | mut overall_progress: ProgressInfo, 111 | input_finished_receiver: InputFinishedReceiver, 112 | display_progress: bool, 113 | max_frames: Option, 114 | ) { 115 | let slots_count = slots.len(); 116 | let segments_pb_holder = MultiProgress::new(); 117 | if !display_progress { 118 | segments_pb_holder.set_draw_target(ProgressDrawTarget::hidden()); 119 | } 120 | let main_pb = segments_pb_holder.add(ProgressBar::new_spinner()); 121 | main_pb.set_style(main_progress_style()); 122 | main_pb.set_prefix(&overall_prefix().to_string()); 123 | main_pb.set_message(&overall_progress.progress_overall()); 124 | if let Some(max_frames) = max_frames { 125 | main_pb.set_length(max_frames); 126 | } 127 | let segment_pbs = (0..receivers.len()) 128 | .map(|i| { 129 | let pb = segments_pb_holder.add(ProgressBar::new_spinner()); 130 | pb.set_style(progress_idle_style()); 131 | pb.set_prefix(&idle_prefix(i >= slots_count).to_string()); 132 | pb.set_position(0); 133 | pb.set_length(0); 134 | pb 135 | }) 136 | .collect::>(); 137 | thread::spawn(move || { 138 | segments_pb_holder.join_and_clear().unwrap(); 139 | }); 140 | 141 | let mut input_finished = false; 142 | loop { 143 | for (slot, rx) in receivers.iter().enumerate() { 144 | while let Ok(msg) = rx.try_recv() { 145 | if update_progress( 146 | msg, 147 | &mut overall_progress, 148 | slots.clone(), 149 | slots_count, 150 | slot, 151 | &segment_pbs[slot], 152 | ) { 153 | update_overall_progress(&output_file, &overall_progress, &main_pb); 154 | } 155 | } 156 | } 157 | 158 | if !input_finished && input_finished_receiver.is_full() { 159 | debug!("Set input finished"); 160 | input_finished = true; 161 | } 162 | 163 | let local_slots_done = slots.iter().all(|slot| !slot.load(Ordering::SeqCst)); 164 | #[cfg(feature = "remote")] 165 | let remote_slots_done = remote_slots 166 | .lock() 167 | .iter() 168 | .all(|slot| slot.workers.iter().all(|worker| !worker)); 169 | #[cfg(not(feature = "remote"))] 170 | let remote_slots_done = true; 171 | if input_finished && local_slots_done && remote_slots_done { 172 | debug!("Finishing progress"); 173 | // Done encoding 174 | main_pb.finish(); 175 | for pb in &segment_pbs { 176 | pb.finish_and_clear(); 177 | } 178 | debug!("Finished progress"); 179 | break; 180 | } 181 | thread::sleep(Duration::from_millis(100)); 182 | } 183 | 184 | thread::sleep(Duration::from_millis(500)); 185 | overall_progress.print_summary(verbose); 186 | debug!("Printed summary"); 187 | } 188 | 189 | fn update_progress( 190 | progress: ProgressStatus, 191 | overall_progress: &mut ProgressInfo, 192 | slots: Arc>, 193 | slots_count: usize, 194 | slot_idx: usize, 195 | pb: &ProgressBar, 196 | ) -> bool { 197 | let remote = slot_idx >= slots_count; 198 | 199 | match progress { 200 | ProgressStatus::Encoding(progress) => { 201 | if pb.length() == 0 { 202 | debug!("Updating progress--new segment starting"); 203 | pb.set_message(""); 204 | pb.set_style(progress_active_style()); 205 | pb.set_prefix(&segment_prefix(progress.segment_idx).to_string()); 206 | pb.reset_elapsed(); 207 | pb.set_position(0); 208 | pb.set_length(progress.total_frames as u64); 209 | if let Some(keyframe) = progress.keyframes.iter().next().copied() { 210 | overall_progress.keyframes.insert(keyframe); 211 | } 212 | overall_progress.next_analysis_frame = cmp::max( 213 | overall_progress.next_analysis_frame, 214 | progress.next_analysis_frame, 215 | ); 216 | true 217 | } else if progress.total_frames == progress.frame_info.len() { 218 | debug!("Updating progress--segment complete"); 219 | overall_progress 220 | .frame_info 221 | .extend_from_slice(&progress.frame_info); 222 | overall_progress.total_frames += progress.total_frames; 223 | overall_progress.encoded_size += progress.encoded_size; 224 | overall_progress 225 | .completed_segments 226 | .insert(progress.segment_idx); 227 | overall_progress.encoding_stats.0 += &progress.encoding_stats.0; 228 | overall_progress.encoding_stats.1 += &progress.encoding_stats.1; 229 | if !remote { 230 | slots[slot_idx].store(false, Ordering::SeqCst); 231 | } 232 | 233 | pb.set_message(""); 234 | pb.set_style(progress_idle_style()); 235 | pb.set_prefix(&idle_prefix(remote).to_string()); 236 | pb.reset_elapsed(); 237 | pb.set_position(0); 238 | pb.set_length(0); 239 | true 240 | } else { 241 | trace!("Updating progress--tick"); 242 | pb.set_position(progress.frame_info.len() as u64); 243 | pb.set_message(&progress.progress()); 244 | false 245 | } 246 | } 247 | ProgressStatus::Loading => { 248 | debug!("Updating progress--frames are loading"); 249 | pb.set_message("Loading frames..."); 250 | false 251 | } 252 | ProgressStatus::Compressing(frames) => { 253 | debug!("Updating progress--compressing {} frames", frames); 254 | pb.set_message(&format!("Compressing {} input frames...", frames)); 255 | false 256 | } 257 | ProgressStatus::Sending(size) => { 258 | debug!("Updating progress--sending {}", size); 259 | pb.set_message(&format!("Sending {} input...", size)); 260 | false 261 | } 262 | ProgressStatus::Idle => { 263 | debug!("Updating progress--idle"); 264 | pb.set_message(""); 265 | pb.set_style(progress_idle_style()); 266 | pb.set_prefix(&idle_prefix(remote).to_string()); 267 | pb.reset_elapsed(); 268 | pb.set_position(0); 269 | pb.set_length(0); 270 | false 271 | } 272 | } 273 | } 274 | 275 | fn update_overall_progress( 276 | output_file: &Output, 277 | overall_progress: &ProgressInfo, 278 | pb: &ProgressBar, 279 | ) { 280 | update_progress_file(output_file, &overall_progress); 281 | 282 | pb.set_message(&overall_progress.progress_overall()); 283 | } 284 | 285 | fn progress_idle_style() -> ProgressStyle { 286 | ProgressStyle::default_spinner().template("[{prefix}] {msg}") 287 | } 288 | 289 | fn main_progress_style() -> ProgressStyle { 290 | ProgressStyle::default_spinner().template("[{prefix}] [{elapsed_precise}] {wide_msg}") 291 | } 292 | 293 | fn progress_active_style() -> ProgressStyle { 294 | ProgressStyle::default_bar() 295 | .template( 296 | "[{prefix}] [{elapsed_precise}] {bar:24.blue/white.dim} {pos:>4}/{len:4} {wide_msg}", 297 | ) 298 | .progress_chars("##-") 299 | } 300 | 301 | fn overall_prefix() -> StyledObject<&'static str> { 302 | style("Overall").blue().bold() 303 | } 304 | 305 | fn idle_prefix(remote: bool) -> StyledObject<&'static str> { 306 | if remote { 307 | style("Remote").yellow().dim() 308 | } else { 309 | style("Idle").cyan().dim() 310 | } 311 | } 312 | 313 | fn segment_prefix(segment_idx: usize) -> StyledObject { 314 | style(format!("Seg. {:>4}", segment_idx)).cyan() 315 | } 316 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/bin/encode.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::BTreeSet, 3 | fs::{remove_file, File}, 4 | io::{BufReader, BufWriter, Read, Write}, 5 | path::Path, 6 | sync::{ 7 | atomic::{AtomicBool, AtomicUsize, Ordering}, 8 | Arc, 9 | }, 10 | thread, 11 | thread::sleep, 12 | time::Duration, 13 | }; 14 | 15 | #[cfg(feature = "remote")] 16 | use anyhow::bail; 17 | use anyhow::Result; 18 | use console::style; 19 | use crossbeam_channel::{bounded, unbounded, TryRecvError}; 20 | use crossbeam_utils::thread::{scope, Scope}; 21 | #[cfg(feature = "remote")] 22 | use log::error; 23 | use log::{debug, info}; 24 | #[cfg(feature = "remote")] 25 | use parking_lot::Mutex; 26 | use rav1e::prelude::*; 27 | use rav1e_by_gop::*; 28 | use serde::{de::DeserializeOwned, Serialize}; 29 | use threadpool::ThreadPool; 30 | use y4m::Decoder; 31 | 32 | use super::VideoDetails; 33 | #[cfg(feature = "remote")] 34 | use crate::analyze::{InputFinishedReceiver, RemoteAnalyzerChannel, RemoteAnalyzerReceiver}; 35 | #[cfg(feature = "remote")] 36 | use crate::remote::{discover_remote_worker, remote_encode_segment, RemoteWorkerInfo}; 37 | use crate::{ 38 | analyze::{run_first_pass, AnalyzerChannel, AnalyzerReceiver, InputFinishedChannel}, 39 | decide_thread_count, 40 | decode::*, 41 | muxer::create_muxer, 42 | progress::*, 43 | CliOptions, 44 | Output, 45 | }; 46 | 47 | pub fn perform_encode( 48 | keyframes: BTreeSet, 49 | next_analysis_frame: usize, 50 | opts: &CliOptions, 51 | progress: Option, 52 | ) -> Result<()> { 53 | let reader = opts.input.as_reader()?; 54 | let dec = y4m::decode(reader).expect("input is not a y4m file"); 55 | let video_info = get_video_details(&dec); 56 | scope(|s| { 57 | if video_info.bit_depth == 8 { 58 | perform_encode_inner::( 59 | s, 60 | keyframes, 61 | next_analysis_frame, 62 | opts, 63 | progress, 64 | dec, 65 | video_info, 66 | ) 67 | } else { 68 | perform_encode_inner::( 69 | s, 70 | keyframes, 71 | next_analysis_frame, 72 | opts, 73 | progress, 74 | dec, 75 | video_info, 76 | ) 77 | } 78 | }) 79 | .unwrap() 80 | } 81 | 82 | pub fn perform_encode_inner< 83 | T: Pixel + Serialize + DeserializeOwned + Default, 84 | R: 'static + Read + Send, 85 | >( 86 | s: &Scope, 87 | keyframes: BTreeSet, 88 | next_analysis_frame: usize, 89 | opts: &CliOptions, 90 | progress: Option, 91 | dec: Decoder, 92 | video_info: VideoDetails, 93 | ) -> Result<()> { 94 | info!( 95 | "Using {} decoder: {}p @ {} fps, {}, {}", 96 | style("y4m").cyan(), 97 | style(format!("{}x{}", video_info.width, video_info.height)).cyan(), 98 | style(format!( 99 | "{}/{}", 100 | video_info.time_base.den, video_info.time_base.num 101 | )) 102 | .cyan(), 103 | style(video_info.chroma_sampling).cyan(), 104 | style(format!("{}-bit", video_info.bit_depth)).cyan() 105 | ); 106 | 107 | #[cfg(feature = "remote")] 108 | let remote_workers = opts 109 | .workers 110 | .iter() 111 | .map(|worker| discover_remote_worker(worker)) 112 | .filter_map(|worker| { 113 | worker 114 | .map_err(|e| { 115 | error!("Failed to negotiate with remote worker: {}", e); 116 | e 117 | }) 118 | .ok() 119 | }) 120 | .collect::>(); 121 | #[cfg(feature = "remote")] 122 | let remote_worker_slot_count = remote_workers 123 | .iter() 124 | .map(|worker| worker.workers.len()) 125 | .sum::(); 126 | #[cfg(not(feature = "remote"))] 127 | let remote_worker_slot_count = 0; 128 | 129 | let num_threads = 130 | decide_thread_count(opts, &video_info, remote_worker_slot_count > 0, opts.tiles); 131 | 132 | #[cfg(feature = "remote")] 133 | if !remote_workers.is_empty() { 134 | info!( 135 | "Discovered {} remote workers with up to {} slots", 136 | remote_workers.len(), 137 | remote_worker_slot_count 138 | ) 139 | } else if num_threads == 0 { 140 | bail!("Cannot disable local threads without having remote workers available!"); 141 | } 142 | 143 | let mut thread_pool = if num_threads > 0 { 144 | Some(ThreadPool::new(num_threads)) 145 | } else { 146 | None 147 | }; 148 | let rayon_pool = Arc::new( 149 | rayon::ThreadPoolBuilder::new() 150 | .num_threads(num_threads) 151 | .build() 152 | .unwrap(), 153 | ); 154 | 155 | let overall_progress = if let Some(progress) = progress { 156 | progress 157 | } else { 158 | let progress = ProgressInfo::new( 159 | Rational { 160 | num: video_info.time_base.den, 161 | den: video_info.time_base.num, 162 | }, 163 | 0, 164 | keyframes, 165 | 0, 166 | next_analysis_frame, 167 | opts.max_frames, 168 | ); 169 | 170 | // Do an initial write of the progress file, 171 | // so we don't need to redo keyframe search. 172 | update_progress_file(&opts.output, &progress); 173 | 174 | progress 175 | }; 176 | 177 | let num_local_slots = if num_threads == 0 { 178 | 0 179 | } else { 180 | let local_workers = opts 181 | .local_workers 182 | .unwrap_or_else(|| (num_threads / opts.tiles).max(1)); 183 | info!( 184 | "Using {} encoder threads ({} local workers)", 185 | style(num_threads).cyan(), 186 | style(local_workers).cyan() 187 | ); 188 | local_workers 189 | }; 190 | 191 | let analyzer_channel: AnalyzerChannel = unbounded(); 192 | #[cfg(feature = "remote")] 193 | let remote_analyzer_channel: RemoteAnalyzerChannel = unbounded(); 194 | let progress_channels: Vec = (0..(num_local_slots + remote_worker_slot_count)) 195 | .map(|_| unbounded()) 196 | .collect(); 197 | let input_finished_channel: InputFinishedChannel = bounded(1); 198 | let slots = Arc::new({ 199 | let mut slots = Vec::with_capacity(num_local_slots); 200 | for _ in 0..num_local_slots { 201 | slots.push(AtomicBool::new(false)); 202 | } 203 | slots 204 | }); 205 | #[cfg(feature = "remote")] 206 | let remote_slots = Arc::new(Mutex::new(remote_workers)); 207 | 208 | let output_file = opts.output.clone(); 209 | let verbose = opts.verbose; 210 | let start_frameno = overall_progress.next_analysis_frame; 211 | let known_keyframes = overall_progress.keyframes.clone(); 212 | let skipped_segments = overall_progress.completed_segments.clone(); 213 | let display_progress = opts.display_progress; 214 | 215 | let receivers = progress_channels 216 | .iter() 217 | .map(|(_, rx)| rx.clone()) 218 | .collect::>(); 219 | let slots_ref = slots.clone(); 220 | #[cfg(feature = "remote")] 221 | let remote_slots_ref = remote_slots.clone(); 222 | let input_finished_receiver = input_finished_channel.1.clone(); 223 | let max_frames = opts.max_frames; 224 | let num_segments = Arc::new(AtomicUsize::new(0)); 225 | s.spawn(move |_| { 226 | watch_progress_receivers( 227 | receivers, 228 | slots_ref, 229 | #[cfg(feature = "remote")] 230 | remote_slots_ref, 231 | output_file, 232 | verbose, 233 | overall_progress, 234 | input_finished_receiver, 235 | display_progress, 236 | max_frames, 237 | ); 238 | }); 239 | 240 | let opts_ref = opts.clone(); 241 | let analyzer_sender = analyzer_channel.0.clone(); 242 | #[cfg(feature = "remote")] 243 | let remote_analyzer_sender = remote_analyzer_channel.0.clone(); 244 | #[cfg(feature = "remote")] 245 | let remote_slots_ref = remote_slots.clone(); 246 | let input_finished_sender = input_finished_channel.0.clone(); 247 | #[cfg(feature = "remote")] 248 | let input_finished_receiver = input_finished_channel.1.clone(); 249 | #[cfg(not(feature = "remote"))] 250 | let input_finished_receiver = input_finished_channel.1; 251 | let progress_senders = progress_channels 252 | .iter() 253 | .map(|(tx, _)| tx.clone()) 254 | .collect::>(); 255 | #[cfg(feature = "remote")] 256 | let remote_progress_senders = progress_senders 257 | .iter() 258 | .skip(num_local_slots) 259 | .cloned() 260 | .collect::>(); 261 | let rayon_handle = rayon_pool.clone(); 262 | let num_segments_handle = num_segments.clone(); 263 | s.spawn(move |s| { 264 | #[cfg(feature = "remote")] 265 | let color_primaries = opts_ref.color_primaries; 266 | #[cfg(feature = "remote")] 267 | let transfer_characteristics = opts_ref.transfer_characteristics; 268 | #[cfg(feature = "remote")] 269 | let matrix_coefficients = opts_ref.matrix_coefficients; 270 | run_first_pass::( 271 | dec, 272 | opts_ref, 273 | analyzer_sender, 274 | #[cfg(feature = "remote")] 275 | remote_analyzer_sender, 276 | progress_senders, 277 | #[cfg(feature = "remote")] 278 | remote_progress_senders, 279 | input_finished_sender, 280 | input_finished_receiver, 281 | slots, 282 | #[cfg(feature = "remote")] 283 | remote_slots_ref, 284 | rayon_handle, 285 | start_frameno, 286 | known_keyframes, 287 | skipped_segments, 288 | video_info, 289 | s, 290 | num_segments_handle, 291 | #[cfg(feature = "remote")] 292 | color_primaries, 293 | #[cfg(feature = "remote")] 294 | transfer_characteristics, 295 | #[cfg(feature = "remote")] 296 | matrix_coefficients, 297 | ); 298 | }); 299 | 300 | #[cfg(feature = "remote")] 301 | for worker in remote_slots.lock().iter() { 302 | let receiver = worker.update_channel.1.clone(); 303 | let remote_slots_ref = remote_slots.clone(); 304 | let input_finished_receiver = input_finished_channel.1.clone(); 305 | s.spawn(move |_| watch_worker_updates(remote_slots_ref, receiver, input_finished_receiver)); 306 | } 307 | 308 | // Write only the ivf header 309 | create_muxer(&match &opts.output { 310 | Output::File(output_file) => Output::File(get_segment_output_filename(&output_file, 0)), 311 | x => x.clone(), 312 | }) 313 | .map(|mut output| { 314 | output.write_header( 315 | video_info.width, 316 | video_info.height, 317 | video_info.time_base.den as usize, 318 | video_info.time_base.num as usize, 319 | ); 320 | }) 321 | .expect("Failed to create segment output"); 322 | 323 | #[cfg(feature = "remote")] 324 | let remote_listener = { 325 | let input_finished_receiver = input_finished_channel.1; 326 | let remote_analyzer_receiver = remote_analyzer_channel.1; 327 | let remote_slots_ref = remote_slots; 328 | s.spawn(move |s| { 329 | listen_for_remote_workers( 330 | s, 331 | remote_analyzer_receiver, 332 | input_finished_receiver, 333 | remote_slots_ref, 334 | ) 335 | }) 336 | }; 337 | 338 | if num_threads > 0 { 339 | let _ = listen_for_local_workers::( 340 | EncodeOptions::from(opts), 341 | &opts.output, 342 | thread_pool.as_mut().unwrap(), 343 | rayon_pool, 344 | video_info, 345 | &progress_channels, 346 | analyzer_channel.1, 347 | ); 348 | thread_pool.unwrap().join(); 349 | } 350 | #[cfg(feature = "remote")] 351 | let _ = remote_listener.join(); 352 | 353 | if let Output::File(output) = &opts.output { 354 | mux_output_files(&output, num_segments.load(Ordering::SeqCst))?; 355 | }; 356 | 357 | Ok(()) 358 | } 359 | 360 | #[cfg(feature = "remote")] 361 | fn watch_worker_updates( 362 | remote_slots: Arc>>, 363 | update_receiver: WorkerUpdateReceiver, 364 | input_finished_receiver: InputFinishedReceiver, 365 | ) { 366 | let mut done = false; 367 | loop { 368 | sleep(Duration::from_millis(100)); 369 | 370 | if !done && input_finished_receiver.is_full() { 371 | debug!("Worker update thread knows input is done"); 372 | done = true; 373 | let remote_slots_ref = remote_slots.lock(); 374 | if remote_slots_ref 375 | .iter() 376 | .find(|worker| worker.update_channel.1.same_channel(&update_receiver)) 377 | .unwrap() 378 | .workers 379 | .iter() 380 | .filter(|worker| **worker) 381 | .count() 382 | == 0 383 | { 384 | debug!("Exiting worker update thread"); 385 | return; 386 | } 387 | } 388 | 389 | if let Ok(message) = update_receiver.try_recv() { 390 | debug!("Updating worker status: {:?}", message); 391 | let mut remote_slots_ref = remote_slots.lock(); 392 | let worker = remote_slots_ref 393 | .iter_mut() 394 | .find(|worker| worker.update_channel.1.same_channel(&update_receiver)) 395 | .unwrap(); 396 | if let Some(status) = message.status { 397 | worker.slot_status = status; 398 | } 399 | if let Some((slot, new_state)) = message.slot_delta { 400 | worker.workers[slot] = new_state; 401 | if done 402 | && !new_state 403 | && worker.workers.iter().filter(|worker| **worker).count() == 0 404 | { 405 | debug!("Exiting worker update thread after receiving message"); 406 | return; 407 | } 408 | } 409 | debug!( 410 | "Worker now at {:?}, {}/{} workers", 411 | worker.slot_status, 412 | worker.workers.iter().filter(|worker| **worker).count(), 413 | worker.workers.len() 414 | ); 415 | } 416 | } 417 | } 418 | 419 | fn listen_for_local_workers( 420 | encode_opts: EncodeOptions, 421 | output_file: &Output, 422 | thread_pool: &mut ThreadPool, 423 | rayon_pool: Arc, 424 | video_info: VideoDetails, 425 | progress_channels: &[ProgressChannel], 426 | analyzer_receiver: AnalyzerReceiver, 427 | ) -> Result<()> { 428 | loop { 429 | match analyzer_receiver.try_recv() { 430 | Ok(Some(data)) => { 431 | let slot = data.slot; 432 | let segment_idx = data.segment_no + 1; 433 | encode_segment( 434 | encode_opts, 435 | video_info, 436 | data, 437 | thread_pool, 438 | rayon_pool.clone(), 439 | progress_channels[slot].0.clone(), 440 | match output_file { 441 | Output::File(output_file) => { 442 | Output::File(get_segment_output_filename(output_file, segment_idx)) 443 | } 444 | x => x.clone(), 445 | }, 446 | )?; 447 | } 448 | Ok(None) => { 449 | // No more input frames, finish. 450 | break; 451 | } 452 | Err(TryRecvError::Empty) => { 453 | sleep(Duration::from_millis(1000)); 454 | } 455 | Err(e) => { 456 | debug!("{}", e); 457 | break; 458 | } 459 | } 460 | } 461 | Ok(()) 462 | } 463 | 464 | #[cfg(feature = "remote")] 465 | fn listen_for_remote_workers( 466 | scope: &Scope, 467 | remote_analyzer_receiver: RemoteAnalyzerReceiver, 468 | input_finished_receiver: InputFinishedReceiver, 469 | remote_slots: Arc>>, 470 | ) { 471 | let mut threads = Vec::new(); 472 | loop { 473 | if let Ok(message) = remote_analyzer_receiver.try_recv() { 474 | threads.push(scope.spawn(move |_| { 475 | if message.video_info.bit_depth <= 8 { 476 | remote_encode_segment::(message) 477 | } else { 478 | remote_encode_segment::(message) 479 | } 480 | })); 481 | } 482 | if input_finished_receiver.is_full() 483 | && remote_slots 484 | .lock() 485 | .iter() 486 | .map(|slot| slot.workers.iter()) 487 | .flatten() 488 | .filter(|worker| **worker) 489 | .count() 490 | == 0 491 | { 492 | break; 493 | } 494 | sleep(Duration::from_millis(100)); 495 | } 496 | for thread in threads { 497 | thread.join().unwrap(); 498 | } 499 | } 500 | 501 | fn mux_output_files(out_filename: &Path, num_segments: usize) -> Result<()> { 502 | let mut out = BufWriter::new(File::create(out_filename)?); 503 | let segments = 504 | (0..=num_segments).map(|seg_idx| get_segment_output_filename(out_filename, seg_idx)); 505 | let mut files = segments.clone(); 506 | let header = files.next().unwrap(); 507 | std::io::copy(&mut File::open(header)?, &mut out)?; 508 | 509 | let mut pts = 0; 510 | for seg_filename in files { 511 | let mut in_seg = BufReader::new(File::open(seg_filename)?); 512 | loop { 513 | match ivf::read_packet(&mut in_seg) { 514 | Ok(pkt) => { 515 | ivf::write_ivf_frame(&mut out, pts, &pkt.data); 516 | pts += 1; 517 | } 518 | Err(err) => match err.kind() { 519 | std::io::ErrorKind::UnexpectedEof => break, 520 | _ => return Err(err.into()), 521 | }, 522 | } 523 | } 524 | } 525 | out.flush()?; 526 | 527 | // Allow the progress indicator thread 528 | // enough time to output the end-of-encode stats 529 | thread::sleep(Duration::from_secs(3)); 530 | 531 | for segment in segments { 532 | let _ = remove_file(segment); 533 | } 534 | let _ = remove_file(get_progress_filename(out_filename)); 535 | 536 | Ok(()) 537 | } 538 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/bin/rav1e-by-gop.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::cognitive_complexity)] 2 | #![allow(clippy::too_many_arguments)] 3 | 4 | mod analyze; 5 | mod decode; 6 | mod encode; 7 | mod progress; 8 | #[cfg(feature = "remote")] 9 | mod remote; 10 | 11 | #[cfg(feature = "remote")] 12 | use std::time::Duration; 13 | use std::{ 14 | cmp, 15 | collections::BTreeSet, 16 | env, 17 | fmt, 18 | fs::File, 19 | io::{stdin, Read, Write}, 20 | path::PathBuf, 21 | str::FromStr, 22 | }; 23 | 24 | use anyhow::{ensure, Result}; 25 | use clap::{App, Arg, ArgMatches}; 26 | use console::{style, Term}; 27 | #[cfg(feature = "remote")] 28 | use lazy_static::lazy_static; 29 | use log::info; 30 | use rav1e::prelude::*; 31 | use rav1e_by_gop::*; 32 | #[cfg(feature = "remote")] 33 | use serde::Deserialize; 34 | use systemstat::{ByteSize, Platform, System}; 35 | 36 | use self::{encode::*, progress::*}; 37 | 38 | #[cfg(all(target_arch = "x86_64", target_os = "linux"))] 39 | #[global_allocator] 40 | static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc; 41 | 42 | #[cfg(feature = "remote")] 43 | lazy_static! { 44 | static ref CLIENT: reqwest::blocking::Client = reqwest::blocking::ClientBuilder::new() 45 | .timeout(None) 46 | .connect_timeout(Duration::from_secs(10)) 47 | .build() 48 | .unwrap(); 49 | } 50 | 51 | fn main() -> Result<()> { 52 | if env::var("RUST_LOG").is_err() { 53 | env::set_var("RUST_LOG", "rav1e_by_gop=info"); 54 | } 55 | env_logger::builder() 56 | .format(|buf, record| writeln!(buf, "{}", record.args())) 57 | .init(); 58 | 59 | // Thanks, borrow checker 60 | let mem_usage_default = MemoryUsage::default().to_string(); 61 | 62 | #[allow(unused_mut)] 63 | let mut app = App::new("rav1e-by-gop") 64 | .arg( 65 | Arg::with_name("INPUT") 66 | .help("Sets the input file or command to use") 67 | .required(true) 68 | .index(1), 69 | ) 70 | .arg( 71 | Arg::with_name("OUTPUT") 72 | .help("IVF video output") 73 | .short("o") 74 | .long("output") 75 | .takes_value(true), 76 | ) 77 | .arg( 78 | Arg::with_name("SPEED") 79 | .help("rav1e speed level (0-10), smaller values are slower") 80 | .default_value("6") 81 | .short("s") 82 | .long("speed") 83 | .takes_value(true), 84 | ) 85 | .arg( 86 | Arg::with_name("QP") 87 | .help("Quantizer (0-255), smaller values are higher quality") 88 | .default_value("100") 89 | .short("q") 90 | .long("quantizer") 91 | .takes_value(true), 92 | ) 93 | .arg( 94 | Arg::with_name("MAX_BITRATE") 95 | .help("Max local bitrate (kbps)") 96 | .long("max-bitrate") 97 | .alias("vbv-maxrate") 98 | .takes_value(true), 99 | ) 100 | .arg( 101 | Arg::with_name("MIN_KEYINT") 102 | .help("Minimum distance between two keyframes") 103 | .default_value("12") 104 | .short("i") 105 | .long("min-keyint") 106 | .takes_value(true), 107 | ) 108 | .arg( 109 | Arg::with_name("MAX_KEYINT") 110 | .help("Maximum distance between two keyframes") 111 | .default_value("240") 112 | .short("I") 113 | .long("keyint") 114 | .takes_value(true), 115 | ) 116 | .arg( 117 | Arg::with_name("MAX_THREADS") 118 | .help("Limit the maximum number of threads that can be used") 119 | .long("threads") 120 | .takes_value(true), 121 | ) 122 | .arg( 123 | Arg::with_name("MEMORY_LIMIT") 124 | .help("Limit the number of threads based on the amount of memory on the system") 125 | .long("memory") 126 | .takes_value(true) 127 | .possible_values(&["light", "heavy", "unlimited"]) 128 | .default_value(&mem_usage_default), 129 | ) 130 | .arg( 131 | Arg::with_name("FRAMES") 132 | .help("Limit the number of frames to encode") 133 | .long("limit") 134 | .alias("frames") 135 | .takes_value(true), 136 | ) 137 | .arg( 138 | Arg::with_name("FORCE_RESUME") 139 | .help("Resume any in-progress encodes without being prompted") 140 | .long("resume") 141 | .conflicts_with("FORCE_OVERWRITE"), 142 | ) 143 | .arg( 144 | Arg::with_name("FORCE_OVERWRITE") 145 | .help("Overwrite any in-progress encodes without being prompted") 146 | .long("overwrite") 147 | .conflicts_with("FORCE_RESUME"), 148 | ) 149 | .arg( 150 | Arg::with_name("VERBOSE") 151 | .help("Print more stats at end of encode") 152 | .long("verbose") 153 | .short("v"), 154 | ) 155 | .arg( 156 | Arg::with_name("NO_PROGRESS") 157 | .help( 158 | "Hide the progress bars. Mostly useful for debugging. Automatically set if \ 159 | not running from a TTY.", 160 | ) 161 | .long("no-progress") 162 | .hidden(true), 163 | ) 164 | .arg( 165 | Arg::with_name("TILES") 166 | .help("Sets the number of tiles to use") 167 | .long("tiles") 168 | .short("t") 169 | .default_value("1"), 170 | ) 171 | .arg( 172 | Arg::with_name("INPUT_TEMP_FILES") 173 | .help( 174 | "Write y4m input segments to temporary files instead of keeping them in \ 175 | memory. Reduces memory usage but increases disk usage. Should enable more \ 176 | segments to run simultaneously.", 177 | ) 178 | .long("tmp-input"), 179 | ) 180 | .arg( 181 | Arg::with_name("LOCAL_WORKERS") 182 | .help("Limit the maximum number of local workers that can be used") 183 | .long("local-workers") 184 | .takes_value(true), 185 | ) 186 | .arg( 187 | Arg::with_name("COLOR_PRIMARIES") 188 | .help("Color primaries used to describe color parameters") 189 | .long("primaries") 190 | .possible_values(&ColorPrimaries::variants()) 191 | .default_value("unspecified") 192 | .case_insensitive(true), 193 | ) 194 | .arg( 195 | Arg::with_name("TRANSFER_CHARACTERISTICS") 196 | .help("Transfer characteristics used to describe color parameters") 197 | .long("transfer") 198 | .possible_values(&TransferCharacteristics::variants()) 199 | .default_value("unspecified") 200 | .case_insensitive(true), 201 | ) 202 | .arg( 203 | Arg::with_name("MATRIX_COEFFICIENTS") 204 | .help("Matrix coefficients used to describe color parameters") 205 | .long("matrix") 206 | .possible_values(&MatrixCoefficients::variants()) 207 | .default_value("unspecified") 208 | .case_insensitive(true), 209 | ); 210 | #[cfg(feature = "remote")] 211 | { 212 | app = app 213 | .arg( 214 | Arg::with_name("WORKERS") 215 | .help("A path to the TOML file defining remote encoding workers") 216 | .long("workers") 217 | .default_value("workers.toml") 218 | .takes_value(true), 219 | ) 220 | .arg( 221 | Arg::with_name("NO_LOCAL") 222 | .help("Disable local encoding threads--requires distributed workers") 223 | .long("no-local"), 224 | ); 225 | } 226 | let matches = app.get_matches(); 227 | let opts = CliOptions::from(&matches); 228 | ensure!( 229 | opts.max_keyint >= opts.min_keyint, 230 | "Max keyint must be greater than or equal to min keyint" 231 | ); 232 | ensure!(opts.qp <= 255, "QP must be between 0-255"); 233 | 234 | let progress = load_progress_file(&opts.output, &matches); 235 | let (keyframes, next_analysis_frame) = if let Some(ref progress) = progress { 236 | info!("{} encode", style("Resuming").yellow()); 237 | (progress.keyframes.clone(), progress.next_analysis_frame) 238 | } else { 239 | info!("{} encode", style("Starting").yellow()); 240 | (BTreeSet::new(), 0) 241 | }; 242 | 243 | info!( 244 | "Encoding using input from `{}`, speed {}, quantizer {}, keyint {}-{}", 245 | opts.input, opts.speed, opts.qp, opts.min_keyint, opts.max_keyint 246 | ); 247 | perform_encode(keyframes, next_analysis_frame, &opts, progress).expect("Failed encoding"); 248 | 249 | info!("{}", style("Finished!").yellow()); 250 | 251 | Ok(()) 252 | } 253 | 254 | #[derive(Debug, Clone)] 255 | pub struct CliOptions { 256 | input: Input, 257 | output: Output, 258 | speed: usize, 259 | qp: usize, 260 | max_bitrate: Option, 261 | min_keyint: u64, 262 | max_keyint: u64, 263 | max_threads: Option, 264 | local_workers: Option, 265 | tiles: usize, 266 | max_frames: Option, 267 | verbose: bool, 268 | memory_usage: MemoryUsage, 269 | display_progress: bool, 270 | #[cfg(feature = "remote")] 271 | workers: Vec, 272 | #[cfg(feature = "remote")] 273 | use_local: bool, 274 | temp_input: bool, 275 | color_primaries: ColorPrimaries, 276 | transfer_characteristics: TransferCharacteristics, 277 | matrix_coefficients: MatrixCoefficients, 278 | } 279 | 280 | impl From<&ArgMatches<'_>> for CliOptions { 281 | fn from(matches: &ArgMatches) -> Self { 282 | let input = matches.value_of("INPUT").unwrap(); 283 | let output = PathBuf::from(matches.value_of("OUTPUT").unwrap()); 284 | let output = if ["/dev/null", "nul", "null"] 285 | .contains(&output.as_os_str().to_string_lossy().to_lowercase().as_str()) 286 | { 287 | Output::Null 288 | } else { 289 | Output::File(output) 290 | }; 291 | let temp_input = match output { 292 | Output::File(_) => matches.is_present("INPUT_TEMP_FILES"), 293 | Output::Null => { 294 | info!("Null output unsupported with temp input files, disabling temp input"); 295 | false 296 | } 297 | _ => unreachable!(), 298 | }; 299 | CliOptions { 300 | input: if input == "-" { 301 | Input::Stdin 302 | } else { 303 | Input::File(PathBuf::from(input)) 304 | }, 305 | output, 306 | speed: matches.value_of("SPEED").unwrap().parse().unwrap(), 307 | qp: matches.value_of("QP").unwrap().parse().unwrap(), 308 | max_bitrate: matches 309 | .value_of("MAX_BITRATE") 310 | .map(|val| val.parse::().unwrap() * 1000), 311 | min_keyint: matches.value_of("MIN_KEYINT").unwrap().parse().unwrap(), 312 | max_keyint: matches.value_of("MAX_KEYINT").unwrap().parse().unwrap(), 313 | max_threads: matches 314 | .value_of("MAX_THREADS") 315 | .map(|threads| threads.parse().unwrap()), 316 | tiles: matches.value_of("TILES").unwrap().parse().unwrap(), 317 | local_workers: matches 318 | .value_of("LOCAL_WORKERS") 319 | .map(|workers| workers.parse().unwrap()), 320 | max_frames: matches 321 | .value_of("FRAMES") 322 | .map(|frames| frames.parse().unwrap()), 323 | verbose: matches.is_present("VERBOSE"), 324 | memory_usage: matches 325 | .value_of("MEMORY_LIMIT") 326 | .map(|val| MemoryUsage::from_str(val).expect("Invalid option for memory limit")) 327 | .unwrap_or_default(), 328 | display_progress: Term::stderr().features().is_attended() 329 | && !matches.is_present("NO_PROGRESS"), 330 | #[cfg(feature = "remote")] 331 | workers: { 332 | let workers_file_path = matches.value_of("WORKERS").unwrap(); 333 | match File::open(workers_file_path) { 334 | Ok(mut file) => { 335 | let mut contents = String::new(); 336 | file.read_to_string(&mut contents).unwrap(); 337 | match toml::from_str::(&contents) { 338 | Ok(res) => { 339 | info!("Loaded remote workers file"); 340 | res.workers 341 | } 342 | Err(e) => { 343 | panic!("Malformed remote workers file: {}", e); 344 | } 345 | } 346 | } 347 | Err(_) => { 348 | info!("Could not open workers file; using local encoding only"); 349 | Vec::new() 350 | } 351 | } 352 | }, 353 | #[cfg(feature = "remote")] 354 | use_local: !matches.is_present("NO_LOCAL"), 355 | temp_input, 356 | color_primaries: matches 357 | .value_of("COLOR_PRIMARIES") 358 | .unwrap() 359 | .parse() 360 | .unwrap_or_default(), 361 | transfer_characteristics: matches 362 | .value_of("TRANSFER_CHARACTERISTICS") 363 | .unwrap() 364 | .parse() 365 | .unwrap_or_default(), 366 | matrix_coefficients: matches 367 | .value_of("MATRIX_COEFFICIENTS") 368 | .unwrap() 369 | .parse() 370 | .unwrap_or_default(), 371 | } 372 | } 373 | } 374 | 375 | impl From<&CliOptions> for EncodeOptions { 376 | fn from(other: &CliOptions) -> Self { 377 | EncodeOptions { 378 | speed: other.speed, 379 | qp: other.qp, 380 | max_bitrate: other.max_bitrate, 381 | tiles: other.tiles, 382 | color_primaries: other.color_primaries, 383 | transfer_characteristics: other.transfer_characteristics, 384 | matrix_coefficients: other.matrix_coefficients, 385 | } 386 | } 387 | } 388 | 389 | #[derive(Debug, Clone)] 390 | pub enum Input { 391 | File(PathBuf), 392 | Stdin, 393 | } 394 | 395 | impl Input { 396 | pub fn as_reader(&self) -> Result> { 397 | Ok(match self { 398 | Input::File(filename) => Box::new(File::open(filename)?), 399 | Input::Stdin => Box::new(stdin()), 400 | }) 401 | } 402 | } 403 | 404 | impl fmt::Display for Input { 405 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 406 | match self { 407 | Input::File(path) => write!(f, "{}", path.to_string_lossy().as_ref()), 408 | Input::Stdin => write!(f, "stdin"), 409 | } 410 | } 411 | } 412 | 413 | #[derive(Debug, Clone, Copy)] 414 | pub enum MemoryUsage { 415 | Light, 416 | Heavy, 417 | Unlimited, 418 | } 419 | 420 | impl Default for MemoryUsage { 421 | fn default() -> Self { 422 | MemoryUsage::Light 423 | } 424 | } 425 | 426 | impl FromStr for MemoryUsage { 427 | type Err = (); 428 | 429 | fn from_str(s: &str) -> Result { 430 | match s.to_lowercase().as_str() { 431 | "light" => Ok(MemoryUsage::Light), 432 | "heavy" => Ok(MemoryUsage::Heavy), 433 | "unlimited" => Ok(MemoryUsage::Unlimited), 434 | _ => Err(()), 435 | } 436 | } 437 | } 438 | 439 | impl fmt::Display for MemoryUsage { 440 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 441 | write!( 442 | f, 443 | "{}", 444 | match self { 445 | MemoryUsage::Light => "light", 446 | MemoryUsage::Heavy => "heavy", 447 | MemoryUsage::Unlimited => "unlimited", 448 | } 449 | ) 450 | } 451 | } 452 | 453 | fn decide_thread_count( 454 | opts: &CliOptions, 455 | video_info: &VideoDetails, 456 | has_remote_workers: bool, 457 | tiles: usize, 458 | ) -> usize { 459 | #[cfg(feature = "remote")] 460 | if !opts.use_local { 461 | return 0; 462 | } 463 | 464 | let mut num_threads = 0; 465 | 466 | // Limit based on available memory. 467 | let sys = System::new(); 468 | let sys_memory = sys.memory(); 469 | if let Ok(sys_memory) = sys_memory { 470 | let bytes_per_frame = bytes_per_frame(video_info); 471 | let rdo_lookahead_frames = 40; 472 | // The RDO Lookahead will have extra data loaded, 473 | // and will be uncompressed in memory. 474 | // The remainder of frames will have only basic YUV data loaded, 475 | // and will be compressed in memory. 476 | let bytes_per_segment = if opts.temp_input { 477 | bytes_per_frame * rdo_lookahead_frames * 9 478 | } else if opts.max_keyint <= rdo_lookahead_frames { 479 | bytes_per_frame * opts.max_keyint * 6 480 | } else { 481 | (bytes_per_frame * rdo_lookahead_frames * 9) 482 | + bytes_per_frame * (opts.max_keyint - rdo_lookahead_frames) * 6 / 10 483 | }; 484 | let total = sys_memory.total.as_u64(); 485 | match opts.memory_usage { 486 | MemoryUsage::Light => { 487 | // Uses 50% of memory, minimum 2GB, 488 | // minimum unreserved memory of 2GB, 489 | // maximum unreserved memory of 12GB 490 | let unreserved = cmp::min( 491 | ByteSize::gb(12).as_u64(), 492 | cmp::max(ByteSize::gb(2).as_u64(), sys_memory.total.as_u64() / 2), 493 | ); 494 | let limit = cmp::max(ByteSize::gb(2).as_u64(), total.saturating_sub(unreserved)); 495 | num_threads = cmp::max( 496 | 1, 497 | // Having remote workers means input must continue 498 | // to be read and analyzed while local encodes 499 | // are happening. i.e. it's like having one extra thread. 500 | (limit / bytes_per_segment) as usize - if has_remote_workers { 1 } else { 0 }, 501 | ) * tiles; 502 | } 503 | MemoryUsage::Heavy => { 504 | // Uses 80% of memory, minimum 2GB, 505 | // minimum unreserved memory of 1GB, 506 | // maximum unreserved memory of 6GB 507 | let unreserved = cmp::min( 508 | ByteSize::gb(6).as_u64(), 509 | cmp::max(ByteSize::gb(1).as_u64(), sys_memory.total.as_u64() * 4 / 5), 510 | ); 511 | let limit = cmp::max(ByteSize::gb(2).as_u64(), total.saturating_sub(unreserved)); 512 | num_threads = cmp::max( 513 | 1, 514 | (limit / bytes_per_segment) as usize - if has_remote_workers { 1 } else { 0 }, 515 | ) * tiles; 516 | } 517 | MemoryUsage::Unlimited => { 518 | num_threads = num_cpus::get(); 519 | } 520 | } 521 | } 522 | 523 | // Limit to the number of logical CPUs. 524 | num_threads = cmp::min(num_cpus::get(), num_threads); 525 | if let Some(max_threads) = opts.max_threads { 526 | num_threads = cmp::min(num_threads, max_threads); 527 | } 528 | 529 | num_threads 530 | } 531 | 532 | fn bytes_per_frame(video_info: &VideoDetails) -> u64 { 533 | let bytes_per_plane = 534 | video_info.width * video_info.height * if video_info.bit_depth > 8 { 2 } else { 1 }; 535 | (match video_info.chroma_sampling { 536 | ChromaSampling::Cs420 => bytes_per_plane * 3 / 2, 537 | ChromaSampling::Cs422 => bytes_per_plane * 2, 538 | ChromaSampling::Cs444 => bytes_per_plane * 3, 539 | ChromaSampling::Cs400 => bytes_per_plane, 540 | }) as u64 541 | } 542 | 543 | #[cfg(feature = "remote")] 544 | #[derive(Debug, Clone, Deserialize)] 545 | pub struct WorkerFile { 546 | pub workers: Vec, 547 | } 548 | 549 | #[cfg(feature = "remote")] 550 | #[derive(Debug, Clone, Deserialize)] 551 | pub struct WorkerConfig { 552 | pub host: String, 553 | #[serde(default)] 554 | pub port: Option, 555 | pub password: String, 556 | #[serde(default)] 557 | pub secure: bool, 558 | } 559 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/bin/analyze.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::{BTreeMap, BTreeSet}, 3 | fs::File, 4 | io::{BufWriter, Read, Write}, 5 | sync::{ 6 | atomic::{AtomicBool, AtomicUsize, Ordering}, 7 | Arc, 8 | }, 9 | thread::sleep, 10 | time::Duration, 11 | }; 12 | 13 | use av_scenechange::{new_detector, DetectionOptions}; 14 | use crossbeam_channel::{unbounded, Receiver, Sender}; 15 | use crossbeam_utils::thread::Scope; 16 | use itertools::Itertools; 17 | use log::{debug, error}; 18 | #[cfg(feature = "remote")] 19 | use parking_lot::Mutex; 20 | use rav1e::prelude::*; 21 | use rav1e_by_gop::*; 22 | use serde::{de::DeserializeOwned, Serialize}; 23 | #[cfg(feature = "remote")] 24 | use systemstat::ByteSize; 25 | use v_frame::{frame::Frame, pixel::Pixel}; 26 | use y4m::Decoder; 27 | 28 | #[cfg(feature = "remote")] 29 | use crate::progress::get_segment_output_filename; 30 | #[cfg(feature = "remote")] 31 | use crate::remote::{wait_for_slot_allocation, RemoteWorkerInfo}; 32 | #[cfg(feature = "remote")] 33 | use crate::CLIENT; 34 | use crate::{ 35 | compress_frame, 36 | decode::{get_video_details, process_raw_frame, read_raw_frame, DecodeError}, 37 | progress::get_segment_input_filename, 38 | CliOptions, 39 | }; 40 | 41 | pub(crate) type AnalyzerSender = Sender>; 42 | pub(crate) type AnalyzerReceiver = Receiver>; 43 | pub(crate) type AnalyzerChannel = (AnalyzerSender, AnalyzerReceiver); 44 | 45 | #[cfg(feature = "remote")] 46 | pub(crate) type RemoteAnalyzerSender = Sender; 47 | #[cfg(feature = "remote")] 48 | pub(crate) type RemoteAnalyzerReceiver = Receiver; 49 | #[cfg(feature = "remote")] 50 | pub(crate) type RemoteAnalyzerChannel = (RemoteAnalyzerSender, RemoteAnalyzerReceiver); 51 | 52 | pub(crate) type InputFinishedSender = Sender<()>; 53 | pub(crate) type InputFinishedReceiver = Receiver<()>; 54 | pub(crate) type InputFinishedChannel = (InputFinishedSender, InputFinishedReceiver); 55 | 56 | pub(crate) type SlotReadySender = Sender; 57 | pub(crate) type SlotReadyReceiver = Receiver; 58 | pub(crate) type SlotReadyChannel = (SlotReadySender, SlotReadyReceiver); 59 | 60 | #[cfg_attr(not(feature = "remote"), allow(unused_variables))] 61 | pub(crate) fn run_first_pass< 62 | T: Pixel + Serialize + DeserializeOwned + Default, 63 | R: 'static + Read + Send, 64 | >( 65 | mut dec: Decoder, 66 | opts: CliOptions, 67 | sender: AnalyzerSender, 68 | #[cfg(feature = "remote")] remote_sender: RemoteAnalyzerSender, 69 | progress_senders: Vec, 70 | #[cfg(feature = "remote")] remote_progress_senders: Vec, 71 | input_finished_sender: InputFinishedSender, 72 | input_finished_receiver: InputFinishedReceiver, 73 | slot_pool: Arc>, 74 | #[cfg(feature = "remote")] remote_pool: Arc>>, 75 | rayon_pool: Arc, 76 | next_frameno: usize, 77 | known_keyframes: BTreeSet, 78 | skipped_segments: BTreeSet, 79 | video_info: VideoDetails, 80 | scope: &Scope, 81 | num_segments: Arc, 82 | #[cfg(feature = "remote")] color_primaries: ColorPrimaries, 83 | #[cfg(feature = "remote")] transfer_characteristics: TransferCharacteristics, 84 | #[cfg(feature = "remote")] matrix_coefficients: MatrixCoefficients, 85 | ) { 86 | let sc_opts = DetectionOptions { 87 | fast_analysis: opts.speed >= 10, 88 | ignore_flashes: true, 89 | lookahead_distance: 5, 90 | min_scenecut_distance: Some(opts.min_keyint as usize), 91 | max_scenecut_distance: Some(opts.max_keyint as usize), 92 | }; 93 | let cfg = get_video_details(&dec); 94 | let slot_ready_channel: SlotReadyChannel = unbounded(); 95 | 96 | // Wait for an open slot before loading more frames, 97 | // to reduce memory usage. 98 | let slot_ready_sender = slot_ready_channel.0.clone(); 99 | let pool_handle = slot_pool.clone(); 100 | let speed = opts.speed; 101 | let qp = opts.qp; 102 | let max_bitrate = opts.max_bitrate; 103 | let tiles = opts.tiles; 104 | scope.spawn(move |s| { 105 | slot_checker_loop::( 106 | pool_handle, 107 | #[cfg(feature = "remote")] 108 | remote_pool, 109 | slot_ready_sender, 110 | #[cfg(feature = "remote")] 111 | remote_progress_senders, 112 | input_finished_receiver, 113 | #[cfg(feature = "remote")] 114 | s, 115 | #[cfg(feature = "remote")] 116 | video_info, 117 | #[cfg(feature = "remote")] 118 | speed, 119 | #[cfg(feature = "remote")] 120 | qp, 121 | #[cfg(feature = "remote")] 122 | max_bitrate, 123 | #[cfg(feature = "remote")] 124 | tiles, 125 | #[cfg(feature = "remote")] 126 | color_primaries, 127 | #[cfg(feature = "remote")] 128 | transfer_characteristics, 129 | #[cfg(feature = "remote")] 130 | matrix_coefficients, 131 | ); 132 | }); 133 | 134 | let pool_handle = slot_pool.clone(); 135 | let slot_ready_listener = slot_ready_channel.1.clone(); 136 | let lookahead_distance = sc_opts.lookahead_distance; 137 | let frame_limit = opts 138 | .max_frames 139 | .map(|f| f as usize) 140 | .unwrap_or(usize::max_value()); 141 | scope 142 | .spawn(move |scope| { 143 | let mut detector = new_detector(&mut dec, sc_opts); 144 | let mut analysis_frameno = next_frameno; 145 | let mut lookahead_frameno = 0; 146 | let mut segment_no = 0; 147 | let mut start_frameno; 148 | // The first keyframe will always be 0, so get the second keyframe. 149 | let mut next_known_keyframe = known_keyframes.iter().nth(1).copied(); 150 | let mut keyframes: BTreeSet = known_keyframes.clone(); 151 | keyframes.insert(0); 152 | let mut lookahead_queue: BTreeMap>> = BTreeMap::new(); 153 | 154 | let enc_cfg = build_config( 155 | opts.speed, 156 | opts.qp, 157 | opts.max_bitrate, 158 | opts.tiles, 159 | video_info, 160 | rayon_pool, 161 | opts.color_primaries, 162 | opts.transfer_characteristics, 163 | opts.matrix_coefficients, 164 | ); 165 | let ctx: Context = enc_cfg.new_context::().unwrap(); 166 | 167 | while let Ok(message) = slot_ready_listener.recv() { 168 | debug!("Received slot ready message"); 169 | match message { 170 | Slot::Local(slot) => &progress_senders[slot], 171 | #[cfg(feature = "remote")] 172 | Slot::Remote(ref conn) => &conn.progress_sender, 173 | } 174 | .send(ProgressStatus::Loading) 175 | .unwrap(); 176 | 177 | let mut processed_frames: Vec> = Vec::new(); 178 | let mut frame_count = 0; 179 | loop { 180 | if let Some(next_keyframe) = next_known_keyframe { 181 | start_frameno = lookahead_frameno; 182 | next_known_keyframe = known_keyframes 183 | .iter() 184 | .copied() 185 | .find(|kf| *kf > next_keyframe); 186 | 187 | // Quickly seek ahead if this is a skipped segment 188 | if skipped_segments.contains(&(segment_no + 1)) { 189 | while lookahead_frameno < next_keyframe { 190 | match read_raw_frame(&mut dec) { 191 | Ok(_) => { 192 | lookahead_frameno += 1; 193 | } 194 | Err(DecodeError::EndOfFile) => { 195 | break; 196 | } 197 | Err(e) => { 198 | error!("Decode error: {}", e); 199 | return; 200 | } 201 | } 202 | } 203 | segment_no += 1; 204 | num_segments.fetch_add(1, Ordering::SeqCst); 205 | continue; 206 | } else { 207 | processed_frames = 208 | Vec::with_capacity(next_keyframe - lookahead_frameno); 209 | frame_count = 0; 210 | 211 | if opts.temp_input && !message.is_remote() { 212 | let file = File::create(get_segment_input_filename( 213 | match &opts.output { 214 | Output::File(output) => &output, 215 | Output::Null => unimplemented!( 216 | "Temp file input not supported with /dev/null output" 217 | ), 218 | _ => unreachable!(), 219 | }, 220 | segment_no + 1, 221 | )) 222 | .unwrap(); 223 | let mut writer = BufWriter::new(file); 224 | while lookahead_frameno < next_keyframe { 225 | match read_raw_frame(&mut dec) { 226 | Ok(frame) => { 227 | frame_count += 1; 228 | bincode::serialize_into( 229 | &mut writer, 230 | &process_raw_frame(&frame, &ctx, &cfg), 231 | ) 232 | .unwrap(); 233 | writer.flush().unwrap(); 234 | lookahead_frameno += 1; 235 | } 236 | Err(DecodeError::EndOfFile) => { 237 | break; 238 | } 239 | Err(e) => { 240 | error!("Decode error: {}", e); 241 | return; 242 | } 243 | } 244 | } 245 | } else { 246 | while lookahead_frameno < next_keyframe { 247 | match read_raw_frame(&mut dec) { 248 | Ok(frame) => { 249 | processed_frames.push(compress_frame::( 250 | &process_raw_frame(&frame, &ctx, &cfg), 251 | )); 252 | lookahead_frameno += 1; 253 | } 254 | Err(DecodeError::EndOfFile) => { 255 | break; 256 | } 257 | Err(e) => { 258 | error!("Decode error: {}", e); 259 | return; 260 | } 261 | } 262 | } 263 | } 264 | } 265 | } else { 266 | start_frameno = keyframes.iter().copied().last().unwrap(); 267 | loop { 268 | // Load frames until the lookahead queue is filled 269 | while analysis_frameno + lookahead_distance > lookahead_frameno 270 | && lookahead_frameno < frame_limit 271 | { 272 | match read_raw_frame(&mut dec) { 273 | Ok(frame) => { 274 | lookahead_queue.insert( 275 | lookahead_frameno, 276 | Arc::new(process_raw_frame(&frame, &ctx, &cfg)), 277 | ); 278 | lookahead_frameno += 1; 279 | } 280 | Err(DecodeError::EndOfFile) => { 281 | break; 282 | } 283 | Err(e) => { 284 | error!("Decode error: {}", e); 285 | return; 286 | } 287 | }; 288 | } 289 | 290 | // Analyze the current frame for a scenechange 291 | if analysis_frameno != keyframes.iter().last().copied().unwrap() { 292 | // The frame_queue should start at whatever the previous frame was 293 | let frame_set = lookahead_queue 294 | .iter() 295 | .skip_while(|(frameno, _)| **frameno < analysis_frameno - 1) 296 | .take(lookahead_distance + 1) 297 | .map(|(_, frame)| frame) 298 | .cloned() 299 | .collect::>(); 300 | if frame_set.len() >= 2 { 301 | if detector.analyze_next_frame( 302 | &frame_set, 303 | analysis_frameno as u64, 304 | *keyframes.iter().last().unwrap() as u64, 305 | ) { 306 | keyframes.insert(analysis_frameno); 307 | } 308 | } else { 309 | // End of encode 310 | keyframes.insert(*lookahead_queue.iter().last().unwrap().0 + 1); 311 | break; 312 | } 313 | 314 | analysis_frameno += 1; 315 | if keyframes.iter().last().copied().unwrap() == analysis_frameno - 1 316 | { 317 | // Keyframe found 318 | break; 319 | } 320 | } else if analysis_frameno < lookahead_frameno { 321 | analysis_frameno += 1; 322 | } else { 323 | debug!("End of encode"); 324 | sender.send(None).unwrap(); 325 | input_finished_sender.send(()).unwrap(); 326 | match message { 327 | Slot::Local(slot) => { 328 | pool_handle[slot].store(false, Ordering::SeqCst); 329 | progress_senders[slot].send(ProgressStatus::Idle).unwrap(); 330 | } 331 | #[cfg(feature = "remote")] 332 | Slot::Remote(connection) => { 333 | connection 334 | .progress_sender 335 | .send(ProgressStatus::Idle) 336 | .unwrap(); 337 | connection 338 | .worker_update_sender 339 | .send(WorkerStatusUpdate { 340 | status: Some(SlotStatus::Empty), 341 | slot_delta: Some(( 342 | connection.slot_in_worker, 343 | false, 344 | )), 345 | }) 346 | .unwrap(); 347 | } 348 | }; 349 | return; 350 | } 351 | } 352 | 353 | // The frames comprising the segment are known 354 | let interval: (usize, usize) = keyframes 355 | .iter() 356 | .rev() 357 | .take(2) 358 | .rev() 359 | .copied() 360 | .collect_tuple() 361 | .unwrap(); 362 | let interval_len = interval.1 - interval.0; 363 | if opts.temp_input && !message.is_remote() { 364 | let file = File::create(get_segment_input_filename( 365 | match &opts.output { 366 | Output::File(output) => &output, 367 | Output::Null => unimplemented!( 368 | "Temp file input not supported with /dev/null output" 369 | ), 370 | _ => unimplemented!(), 371 | }, 372 | segment_no + 1, 373 | )) 374 | .unwrap(); 375 | let mut writer = BufWriter::new(file); 376 | frame_count = 0; 377 | for frameno in (interval.0)..(interval.1) { 378 | frame_count += 1; 379 | bincode::serialize_into( 380 | &mut writer, 381 | &Arc::try_unwrap(lookahead_queue.remove(&frameno).unwrap()) 382 | .unwrap(), 383 | ) 384 | .unwrap(); 385 | writer.flush().unwrap(); 386 | } 387 | } else { 388 | processed_frames = Vec::with_capacity(interval_len); 389 | for frameno in (interval.0)..(interval.1) { 390 | processed_frames.push(compress_frame( 391 | &lookahead_queue.remove(&frameno).unwrap(), 392 | )); 393 | } 394 | } 395 | } 396 | break; 397 | } 398 | 399 | match message { 400 | Slot::Local(slot) => { 401 | debug!("Encoding with local slot"); 402 | sender 403 | .send(Some(SegmentData { 404 | segment_no, 405 | slot, 406 | next_analysis_frame: analysis_frameno - 1, 407 | start_frameno, 408 | frame_data: if opts.temp_input { 409 | SegmentFrameData::Y4MFile { 410 | frame_count, 411 | path: get_segment_input_filename( 412 | match &opts.output { 413 | Output::File(output) => &output, 414 | Output::Null => unimplemented!( 415 | "Temp file input not supported with /dev/null \ 416 | output" 417 | ), 418 | _ => unimplemented!(), 419 | }, 420 | segment_no + 1, 421 | ), 422 | } 423 | } else { 424 | SegmentFrameData::CompressedFrames(processed_frames) 425 | }, 426 | })) 427 | .unwrap(); 428 | } 429 | #[cfg(feature = "remote")] 430 | Slot::Remote(mut connection) => { 431 | debug!("Encoding with remote slot"); 432 | let output = opts.output.clone(); 433 | let remote_sender = remote_sender.clone(); 434 | scope.spawn(move |_| { 435 | let frame_count = processed_frames.len(); 436 | connection 437 | .progress_sender 438 | .send(ProgressStatus::Sending(ByteSize( 439 | processed_frames 440 | .iter() 441 | .map(|frame| frame.len() as u64) 442 | .sum(), 443 | ))) 444 | .unwrap(); 445 | while CLIENT 446 | .post(&format!( 447 | "{}{}/{}", 448 | &connection.worker_uri, "segment", connection.request_id 449 | )) 450 | .header("X-RAV1E-AUTH", &connection.worker_password) 451 | .json(&PostSegmentMessage { 452 | keyframe_number: start_frameno, 453 | segment_idx: segment_no + 1, 454 | next_analysis_frame: analysis_frameno - 1, 455 | }) 456 | .send() 457 | .and_then(|res| res.error_for_status()) 458 | .is_err() 459 | { 460 | sleep(Duration::from_secs(5)); 461 | } 462 | while CLIENT 463 | .post(&format!( 464 | "{}{}/{}", 465 | &connection.worker_uri, "segment_data", connection.request_id 466 | )) 467 | .header("X-RAV1E-AUTH", &connection.worker_password) 468 | .body(bincode::serialize(&processed_frames).unwrap()) 469 | .send() 470 | .and_then(|res| res.error_for_status()) 471 | .is_err() 472 | { 473 | sleep(Duration::from_secs(5)); 474 | } 475 | connection 476 | .worker_update_sender 477 | .send(WorkerStatusUpdate { 478 | status: Some(SlotStatus::Empty), 479 | slot_delta: None, 480 | }) 481 | .unwrap(); 482 | connection.encode_info = Some(EncodeInfo { 483 | output_file: match output { 484 | Output::File(output) => Output::File( 485 | get_segment_output_filename(&output, segment_no + 1), 486 | ), 487 | x => x, 488 | }, 489 | frame_count, 490 | next_analysis_frame: analysis_frameno - 1, 491 | segment_idx: segment_no + 1, 492 | start_frameno, 493 | }); 494 | remote_sender.send(*connection).unwrap(); 495 | }); 496 | } 497 | } 498 | segment_no += 1; 499 | num_segments.fetch_add(1, Ordering::SeqCst); 500 | } 501 | }) 502 | .join() 503 | .unwrap(); 504 | 505 | // Close any extra ready sockets 506 | while let Ok(message) = slot_ready_channel.1.try_recv() { 507 | match message { 508 | Slot::Local(slot) => { 509 | slot_pool[slot].store(false, Ordering::SeqCst); 510 | } 511 | #[cfg(feature = "remote")] 512 | Slot::Remote(connection) => { 513 | connection 514 | .worker_update_sender 515 | .send(WorkerStatusUpdate { 516 | status: Some(SlotStatus::Empty), 517 | slot_delta: Some((connection.slot_in_worker, false)), 518 | }) 519 | .unwrap(); 520 | } 521 | }; 522 | } 523 | } 524 | 525 | fn slot_checker_loop( 526 | pool: Arc>, 527 | #[cfg(feature = "remote")] remote_pool: Arc>>, 528 | slot_ready_sender: SlotReadySender, 529 | #[cfg(feature = "remote")] remote_progress_senders: Vec, 530 | input_finished_receiver: InputFinishedReceiver, 531 | #[cfg(feature = "remote")] scope: &Scope, 532 | #[cfg(feature = "remote")] video_info: VideoDetails, 533 | #[cfg(feature = "remote")] speed: usize, 534 | #[cfg(feature = "remote")] qp: usize, 535 | #[cfg(feature = "remote")] max_bitrate: Option, 536 | #[cfg(feature = "remote")] tiles: usize, 537 | #[cfg(feature = "remote")] color_primaries: ColorPrimaries, 538 | #[cfg(feature = "remote")] transfer_characteristics: TransferCharacteristics, 539 | #[cfg(feature = "remote")] matrix_coefficients: MatrixCoefficients, 540 | ) { 541 | loop { 542 | if input_finished_receiver.is_full() { 543 | debug!("Exiting slot checker loop"); 544 | return; 545 | } 546 | 547 | sleep(Duration::from_millis(500)); 548 | 549 | if let Some(slot) = pool.iter().position(|slot| !slot.load(Ordering::SeqCst)) { 550 | if slot_ready_sender.send(Slot::Local(slot)).is_err() { 551 | debug!("Exiting slot checker loop"); 552 | return; 553 | }; 554 | pool[slot].store(true, Ordering::SeqCst); 555 | continue; 556 | } 557 | 558 | #[cfg(feature = "remote")] 559 | { 560 | let mut worker_start_idx = 0; 561 | for worker in remote_pool.lock().iter_mut() { 562 | if worker.workers.iter().all(|worker| *worker) { 563 | worker_start_idx += worker.workers.len(); 564 | continue; 565 | } 566 | if let SlotStatus::Empty = worker.slot_status { 567 | debug!("Empty connection--requesting new slot"); 568 | match CLIENT 569 | .post(&format!("{}{}", &worker.uri, "enqueue")) 570 | .header("X-RAV1E-AUTH", &worker.password) 571 | .json(&SlotRequestMessage { 572 | options: EncodeOptions { 573 | speed, 574 | qp, 575 | max_bitrate, 576 | tiles, 577 | color_primaries, 578 | transfer_characteristics, 579 | matrix_coefficients, 580 | }, 581 | video_info, 582 | client_version: semver::Version::parse(env!("CARGO_PKG_VERSION")) 583 | .unwrap(), 584 | }) 585 | .send() 586 | .and_then(|res| res.error_for_status()) 587 | .and_then(|res| res.json::()) 588 | { 589 | Ok(response) => { 590 | worker.slot_status = SlotStatus::Requested; 591 | let slot = worker.workers.iter().position(|worker| !worker).unwrap(); 592 | let connection = ActiveConnection { 593 | worker_uri: worker.uri.clone(), 594 | worker_password: worker.password.clone(), 595 | video_info, 596 | request_id: response.request_id, 597 | encode_info: None, 598 | worker_update_sender: worker.update_channel.0.clone(), 599 | slot_in_worker: slot, 600 | progress_sender: remote_progress_senders[worker_start_idx + slot] 601 | .clone(), 602 | }; 603 | let slot_ready_sender = slot_ready_sender.clone(); 604 | scope.spawn(move |_| { 605 | wait_for_slot_allocation::(connection, slot_ready_sender) 606 | }); 607 | } 608 | Err(e) => { 609 | error!("Failed to contact remote server: {}", e); 610 | } 611 | } 612 | } 613 | worker_start_idx += worker.workers.len(); 614 | } 615 | } 616 | } 617 | } 618 | -------------------------------------------------------------------------------- /rav1e-by-gop/src/encode/stats.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::BTreeSet, 3 | time::{Duration, Instant}, 4 | }; 5 | 6 | use arrayvec::ArrayVec; 7 | #[cfg(feature = "binary")] 8 | use console::style; 9 | #[cfg(feature = "binary")] 10 | use log::info; 11 | use rav1e::{data::EncoderStats, prelude::*}; 12 | use serde::{Deserialize, Deserializer, Serialize, Serializer}; 13 | 14 | #[derive(Debug, Clone)] 15 | pub struct ProgressInfo { 16 | // Frame rate of the video 17 | pub frame_rate: Rational, 18 | // The length of the whole video, in frames. `None` if not known. 19 | pub total_frames: usize, 20 | // The time the encode was started 21 | pub time_started: Instant, 22 | // List of frames encoded so far 23 | pub frame_info: Vec, 24 | // Summarized verbose encoding stats, split into I and P frames 25 | pub encoding_stats: (EncoderStats, EncoderStats), 26 | // Video size so far in bytes. 27 | // 28 | // This value will be updated in the CLI very frequently, so we cache the previous value 29 | // to reduce the overall complexity. 30 | pub encoded_size: usize, 31 | // The below are used for resume functionality 32 | pub keyframes: BTreeSet, 33 | pub completed_segments: BTreeSet, 34 | pub segment_idx: usize, 35 | pub next_analysis_frame: usize, 36 | pub frame_limit: Option, 37 | } 38 | 39 | impl ProgressInfo { 40 | pub fn new( 41 | frame_rate: Rational, 42 | total_frames: usize, 43 | keyframes: BTreeSet, 44 | segment_idx: usize, 45 | next_analysis_frame: usize, 46 | frame_limit: Option, 47 | ) -> Self { 48 | Self { 49 | frame_rate, 50 | total_frames, 51 | time_started: Instant::now(), 52 | frame_info: Vec::with_capacity(total_frames), 53 | encoded_size: 0, 54 | keyframes, 55 | completed_segments: BTreeSet::new(), 56 | segment_idx, 57 | encoding_stats: (EncoderStats::default(), EncoderStats::default()), 58 | next_analysis_frame, 59 | frame_limit, 60 | } 61 | } 62 | 63 | pub fn add_packet(&mut self, packet: Packet) { 64 | self.encoded_size += packet.data.len(); 65 | match packet.frame_type { 66 | FrameType::KEY => self.encoding_stats.0 += &packet.enc_stats, 67 | _ => self.encoding_stats.1 += &packet.enc_stats, 68 | }; 69 | self.frame_info.push(packet.into()); 70 | } 71 | 72 | #[cfg(feature = "binary")] 73 | fn frames_encoded(&self) -> usize { 74 | self.frame_info.len() 75 | } 76 | 77 | #[cfg(feature = "binary")] 78 | fn encoding_fps(&self) -> f64 { 79 | self.frame_info.len() as f64 / self.elapsed_time() 80 | } 81 | 82 | #[cfg(feature = "binary")] 83 | fn video_fps(&self) -> f64 { 84 | self.frame_rate.num as f64 / self.frame_rate.den as f64 85 | } 86 | 87 | #[cfg(feature = "binary")] 88 | // Returns the bitrate of the frames so far, in bits/second 89 | fn bitrate(&self) -> usize { 90 | let bits = self.encoded_size * 8; 91 | let seconds = self.frame_info.len() as f64 / self.video_fps(); 92 | (bits as f64 / seconds) as usize 93 | } 94 | 95 | #[cfg(feature = "binary")] 96 | // Estimates the final filesize in bytes, if the number of frames is known 97 | fn estimated_size(&self) -> usize { 98 | self.encoded_size * self.total_frames / self.frames_encoded() 99 | } 100 | 101 | #[cfg(feature = "binary")] 102 | // Estimates the remaining encoding time in seconds 103 | fn estimated_time(&self, total_frames: usize) -> u64 { 104 | ((total_frames - self.frames_encoded()) as f64 / self.encoding_fps()) as u64 105 | } 106 | 107 | pub fn elapsed_time(&self) -> f64 { 108 | let duration = Instant::now().duration_since(self.time_started); 109 | duration.as_secs() as f64 + duration.subsec_millis() as f64 / 1000f64 110 | } 111 | 112 | #[cfg(feature = "binary")] 113 | // Number of frames of given type which appear in the video 114 | fn get_frame_type_count(&self, frame_type: FrameType) -> usize { 115 | self.frame_info 116 | .iter() 117 | .filter(|frame| frame.frame_type == frame_type) 118 | .count() 119 | } 120 | 121 | #[cfg(feature = "binary")] 122 | fn get_frame_type_avg_size(&self, frame_type: FrameType) -> usize { 123 | let count = self.get_frame_type_count(frame_type); 124 | if count == 0 { 125 | return 0; 126 | } 127 | self.frame_info 128 | .iter() 129 | .filter(|frame| frame.frame_type == frame_type) 130 | .map(|frame| frame.size) 131 | .sum::() 132 | / count 133 | } 134 | 135 | #[cfg(feature = "binary")] 136 | fn get_frame_type_avg_qp(&self, frame_type: FrameType) -> f32 { 137 | let count = self.get_frame_type_count(frame_type); 138 | if count == 0 { 139 | return 0.; 140 | } 141 | self.frame_info 142 | .iter() 143 | .filter(|frame| frame.frame_type == frame_type) 144 | .map(|frame| frame.qp as f32) 145 | .sum::() 146 | / count as f32 147 | } 148 | 149 | #[cfg(feature = "binary")] 150 | fn get_block_count_by_frame_type(&self, frame_type: FrameType) -> usize { 151 | match frame_type { 152 | FrameType::KEY => self 153 | .encoding_stats 154 | .0 155 | .block_size_counts 156 | .iter() 157 | .sum::(), 158 | FrameType::INTER => self 159 | .encoding_stats 160 | .1 161 | .block_size_counts 162 | .iter() 163 | .sum::(), 164 | _ => unreachable!(), 165 | } 166 | } 167 | 168 | #[cfg(feature = "binary")] 169 | fn get_tx_count_by_frame_type(&self, frame_type: FrameType) -> usize { 170 | match frame_type { 171 | FrameType::KEY => self.encoding_stats.0.tx_type_counts.iter().sum::(), 172 | FrameType::INTER => self.encoding_stats.1.tx_type_counts.iter().sum::(), 173 | _ => unreachable!(), 174 | } 175 | } 176 | 177 | #[cfg(feature = "binary")] 178 | fn get_bsize_pct_by_frame_type(&self, bsize: BlockSize, frame_type: FrameType) -> f32 { 179 | let count = self.get_block_count_by_frame_type(frame_type); 180 | if count == 0 { 181 | return 0.; 182 | } 183 | (match frame_type { 184 | FrameType::KEY => self.encoding_stats.0.block_size_counts[bsize as usize], 185 | FrameType::INTER => self.encoding_stats.1.block_size_counts[bsize as usize], 186 | _ => unreachable!(), 187 | }) as f32 188 | / count as f32 189 | * 100. 190 | } 191 | 192 | #[cfg(feature = "binary")] 193 | fn get_skip_pct_by_frame_type(&self, frame_type: FrameType) -> f32 { 194 | let count = self.get_block_count_by_frame_type(frame_type); 195 | if count == 0 { 196 | return 0.; 197 | } 198 | (match frame_type { 199 | FrameType::KEY => self.encoding_stats.0.skip_block_count, 200 | FrameType::INTER => self.encoding_stats.1.skip_block_count, 201 | _ => unreachable!(), 202 | }) as f32 203 | / count as f32 204 | * 100. 205 | } 206 | 207 | #[cfg(feature = "binary")] 208 | fn get_txtype_pct_by_frame_type(&self, tx_type: TxType, frame_type: FrameType) -> f32 { 209 | let count = self.get_tx_count_by_frame_type(frame_type); 210 | if count == 0 { 211 | return 0.; 212 | } 213 | (match frame_type { 214 | FrameType::KEY => self.encoding_stats.0.tx_type_counts[tx_type as usize], 215 | FrameType::INTER => self.encoding_stats.1.tx_type_counts[tx_type as usize], 216 | _ => unreachable!(), 217 | }) as f32 218 | / count as f32 219 | * 100. 220 | } 221 | 222 | #[cfg(feature = "binary")] 223 | fn get_luma_pred_count_by_frame_type(&self, frame_type: FrameType) -> usize { 224 | match frame_type { 225 | FrameType::KEY => self 226 | .encoding_stats 227 | .0 228 | .luma_pred_mode_counts 229 | .iter() 230 | .sum::(), 231 | FrameType::INTER => self 232 | .encoding_stats 233 | .1 234 | .luma_pred_mode_counts 235 | .iter() 236 | .sum::(), 237 | _ => unreachable!(), 238 | } 239 | } 240 | 241 | #[cfg(feature = "binary")] 242 | fn get_chroma_pred_count_by_frame_type(&self, frame_type: FrameType) -> usize { 243 | match frame_type { 244 | FrameType::KEY => self 245 | .encoding_stats 246 | .0 247 | .chroma_pred_mode_counts 248 | .iter() 249 | .sum::(), 250 | FrameType::INTER => self 251 | .encoding_stats 252 | .1 253 | .chroma_pred_mode_counts 254 | .iter() 255 | .sum::(), 256 | _ => unreachable!(), 257 | } 258 | } 259 | 260 | #[cfg(feature = "binary")] 261 | fn get_luma_pred_mode_pct_by_frame_type( 262 | &self, 263 | pred_mode: PredictionMode, 264 | frame_type: FrameType, 265 | ) -> f32 { 266 | let count = self.get_luma_pred_count_by_frame_type(frame_type); 267 | if count == 0 { 268 | return 0.; 269 | } 270 | (match frame_type { 271 | FrameType::KEY => self.encoding_stats.0.luma_pred_mode_counts[pred_mode as usize], 272 | FrameType::INTER => self.encoding_stats.1.luma_pred_mode_counts[pred_mode as usize], 273 | _ => unreachable!(), 274 | }) as f32 275 | / count as f32 276 | * 100. 277 | } 278 | 279 | #[cfg(feature = "binary")] 280 | fn get_chroma_pred_mode_pct_by_frame_type( 281 | &self, 282 | pred_mode: PredictionMode, 283 | frame_type: FrameType, 284 | ) -> f32 { 285 | let count = self.get_chroma_pred_count_by_frame_type(frame_type); 286 | if count == 0 { 287 | return 0.; 288 | } 289 | (match frame_type { 290 | FrameType::KEY => self.encoding_stats.0.chroma_pred_mode_counts[pred_mode as usize], 291 | FrameType::INTER => self.encoding_stats.1.chroma_pred_mode_counts[pred_mode as usize], 292 | _ => unreachable!(), 293 | }) as f32 294 | / count as f32 295 | * 100. 296 | } 297 | 298 | #[cfg(feature = "binary")] 299 | pub fn print_summary(&self, verbose: bool) { 300 | info!("{}", self.end_of_encode_progress()); 301 | 302 | info!(""); 303 | 304 | info!("{}", style("Summary by Frame Type").yellow()); 305 | info!( 306 | "{:10} | {:>6} | {:>9} | {:>6}", 307 | style("Frame Type").blue(), 308 | style("Count").blue(), 309 | style("Avg Size").blue(), 310 | style("Avg QP").blue(), 311 | ); 312 | self.print_frame_type_summary(FrameType::KEY); 313 | self.print_frame_type_summary(FrameType::INTER); 314 | self.print_frame_type_summary(FrameType::INTRA_ONLY); 315 | self.print_frame_type_summary(FrameType::SWITCH); 316 | 317 | info!(""); 318 | 319 | if verbose { 320 | info!("{}", style("Block Type Usage").yellow()); 321 | self.print_block_type_summary(); 322 | info!(""); 323 | 324 | info!("{}", style("Transform Type Usage").yellow()); 325 | self.print_transform_type_summary(); 326 | info!(""); 327 | 328 | info!("{}", style("Prediction Mode Usage").yellow()); 329 | self.print_prediction_modes_summary(); 330 | info!(""); 331 | } 332 | } 333 | 334 | #[cfg(feature = "binary")] 335 | fn print_frame_type_summary(&self, frame_type: FrameType) { 336 | let count = self.get_frame_type_count(frame_type); 337 | let size = self.get_frame_type_avg_size(frame_type); 338 | let avg_qp = self.get_frame_type_avg_qp(frame_type); 339 | info!( 340 | "{:10} | {:>6} | {:>9} | {:>6.2}", 341 | style(frame_type.to_string().replace(" frame", "")).blue(), 342 | style(count).cyan(), 343 | style(format!("{} B", size)).cyan(), 344 | style(avg_qp).cyan(), 345 | ); 346 | } 347 | 348 | #[cfg(feature = "binary")] 349 | pub fn progress(&self) -> String { 350 | format!( 351 | "{:.2} fps, {:.1} Kb/s, ETA {}", 352 | self.encoding_fps(), 353 | self.bitrate() as f64 / 1000f64, 354 | secs_to_human_time(self.estimated_time(self.total_frames), false) 355 | ) 356 | } 357 | 358 | #[cfg(feature = "binary")] 359 | pub fn progress_overall(&self) -> String { 360 | if self.frames_encoded() == 0 { 361 | if let Some(max_frames) = self.frame_limit { 362 | format!( 363 | "Input Frame {}/{}, Output Frame {}/{}", 364 | self.next_analysis_frame + 1, 365 | max_frames, 366 | self.frames_encoded(), 367 | max_frames, 368 | ) 369 | } else { 370 | format!( 371 | "Input Frame {}, Output Frame {}", 372 | self.next_analysis_frame + 1, 373 | self.frames_encoded(), 374 | ) 375 | } 376 | } else if let Some(max_frames) = self.frame_limit { 377 | format!( 378 | "Input Frame {}/{}, Output Frame {}/{}, {:.2} fps, {:.1} Kb/s, ETA {}", 379 | self.next_analysis_frame + 1, 380 | max_frames, 381 | self.frames_encoded(), 382 | max_frames, 383 | self.encoding_fps(), 384 | self.bitrate() as f64 / 1000f64, 385 | secs_to_human_time(self.estimated_time(max_frames as usize), false) 386 | ) 387 | } else { 388 | format!( 389 | "Input Frame {}, Output Frame {}, {:.2} fps, {:.1} Kb/s", 390 | self.next_analysis_frame + 1, 391 | self.frames_encoded(), 392 | self.encoding_fps(), 393 | self.bitrate() as f64 / 1000f64, 394 | ) 395 | } 396 | } 397 | 398 | #[cfg(feature = "binary")] 399 | fn end_of_encode_progress(&self) -> String { 400 | format!( 401 | "Encoded {} in {}, {:.3} fps, {:.2} Kb/s, size: {:.2} MB", 402 | style(format!("{} frames", self.total_frames)).yellow(), 403 | style(secs_to_human_time(self.elapsed_time() as u64, true)).cyan(), 404 | self.encoding_fps(), 405 | self.bitrate() as f64 / 1000f64, 406 | self.estimated_size() as f64 / (1024 * 1024) as f64, 407 | ) 408 | } 409 | 410 | #[cfg(feature = "binary")] 411 | fn print_block_type_summary(&self) { 412 | self.print_block_type_summary_for_frame_type(FrameType::KEY, 'I'); 413 | self.print_block_type_summary_for_frame_type(FrameType::INTER, 'P'); 414 | } 415 | 416 | #[cfg(feature = "binary")] 417 | fn print_block_type_summary_for_frame_type(&self, frame_type: FrameType, type_label: char) { 418 | info!( 419 | "{:8} {:>6} {:>6} {:>6} {:>6} {:>6} {:>6}", 420 | style(format!("{} Frames", type_label)).yellow(), 421 | style("x128").blue(), 422 | style("x64").blue(), 423 | style("x32").blue(), 424 | style("x16").blue(), 425 | style("x8").blue(), 426 | style("x4").blue() 427 | ); 428 | info!( 429 | "{:>8} {:>5.1}% {:>5.1}% {}", 430 | style("128x").blue(), 431 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_128X128, frame_type)).cyan(), 432 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_128X64, frame_type)).cyan(), 433 | if frame_type == FrameType::INTER { 434 | format!( 435 | "{}: {:>5.1}%", 436 | style("skip").blue(), 437 | style(self.get_skip_pct_by_frame_type(frame_type)).cyan() 438 | ) 439 | } else { 440 | String::new() 441 | } 442 | ); 443 | info!( 444 | "{:>8} {:>5.1}% {:>5.1}% {:>5.1}% {:>5.1}%", 445 | style("64x").blue(), 446 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_64X128, frame_type)).cyan(), 447 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_64X64, frame_type)).cyan(), 448 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_64X32, frame_type)).cyan(), 449 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_64X16, frame_type)).cyan(), 450 | ); 451 | info!( 452 | "{:>8} {:>5.1}% {:>5.1}% {:>5.1}% {:>5.1}%", 453 | style("32x").blue(), 454 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_32X64, frame_type)).cyan(), 455 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_32X32, frame_type)).cyan(), 456 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_32X16, frame_type)).cyan(), 457 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_32X8, frame_type)).cyan(), 458 | ); 459 | info!( 460 | "{:>8} {:>5.1}% {:>5.1}% {:>5.1}% {:>5.1}% {:>5.1}%", 461 | style("16x").blue(), 462 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_16X64, frame_type)).cyan(), 463 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_16X32, frame_type)).cyan(), 464 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_16X16, frame_type)).cyan(), 465 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_16X8, frame_type)).cyan(), 466 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_16X4, frame_type)).cyan(), 467 | ); 468 | info!( 469 | "{:>8} {:>5.1}% {:>5.1}% {:>5.1}% {:>5.1}%", 470 | style("8x").blue(), 471 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_8X32, frame_type)).cyan(), 472 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_8X16, frame_type)).cyan(), 473 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_8X8, frame_type)).cyan(), 474 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_8X4, frame_type)).cyan(), 475 | ); 476 | info!( 477 | "{:>8} {:>5.1}% {:>5.1}% {:>5.1}%", 478 | style("4x").blue(), 479 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_4X16, frame_type)).cyan(), 480 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_4X8, frame_type)).cyan(), 481 | style(self.get_bsize_pct_by_frame_type(BlockSize::BLOCK_4X4, frame_type)).cyan(), 482 | ); 483 | } 484 | 485 | #[cfg(feature = "binary")] 486 | fn print_transform_type_summary(&self) { 487 | self.print_transform_type_summary_by_frame_type(FrameType::KEY, 'I'); 488 | self.print_transform_type_summary_by_frame_type(FrameType::INTER, 'P'); 489 | } 490 | 491 | #[cfg(feature = "binary")] 492 | fn print_transform_type_summary_by_frame_type(&self, frame_type: FrameType, type_label: char) { 493 | info!("{:8}", style(format!("{} Frames", type_label)).yellow()); 494 | info!( 495 | "{:9} {:>5.1}%", 496 | style("DCT_DCT").blue(), 497 | style(self.get_txtype_pct_by_frame_type(TxType::DCT_DCT, frame_type)).cyan() 498 | ); 499 | info!( 500 | "{:9} {:>5.1}%", 501 | style("ADST_DCT").blue(), 502 | style(self.get_txtype_pct_by_frame_type(TxType::ADST_DCT, frame_type)).cyan() 503 | ); 504 | info!( 505 | "{:9} {:>5.1}%", 506 | style("DCT_ADST").blue(), 507 | style(self.get_txtype_pct_by_frame_type(TxType::DCT_ADST, frame_type)).cyan() 508 | ); 509 | info!( 510 | "{:9} {:>5.1}%", 511 | style("ADST_ADST").blue(), 512 | style(self.get_txtype_pct_by_frame_type(TxType::ADST_ADST, frame_type)).cyan() 513 | ); 514 | info!( 515 | "{:9} {:>5.1}%", 516 | style("IDTX").blue(), 517 | style(self.get_txtype_pct_by_frame_type(TxType::IDTX, frame_type)).cyan() 518 | ); 519 | info!( 520 | "{:9} {:>5.1}%", 521 | style("V_DCT").blue(), 522 | style(self.get_txtype_pct_by_frame_type(TxType::V_DCT, frame_type)).cyan() 523 | ); 524 | info!( 525 | "{:9} {:>5.1}%", 526 | style("H_DCT").blue(), 527 | style(self.get_txtype_pct_by_frame_type(TxType::H_DCT, frame_type)).cyan() 528 | ); 529 | } 530 | 531 | #[cfg(feature = "binary")] 532 | fn print_prediction_modes_summary(&self) { 533 | self.print_luma_prediction_mode_summary_by_frame_type(FrameType::KEY, 'I'); 534 | self.print_chroma_prediction_mode_summary_by_frame_type(FrameType::KEY, 'I'); 535 | self.print_luma_prediction_mode_summary_by_frame_type(FrameType::INTER, 'P'); 536 | self.print_chroma_prediction_mode_summary_by_frame_type(FrameType::INTER, 'P'); 537 | } 538 | 539 | #[cfg(feature = "binary")] 540 | fn print_luma_prediction_mode_summary_by_frame_type( 541 | &self, 542 | frame_type: FrameType, 543 | type_label: char, 544 | ) { 545 | info!( 546 | "{}", 547 | style(format!("{} Frame Luma Modes", type_label)).yellow() 548 | ); 549 | if frame_type == FrameType::KEY { 550 | info!( 551 | "{:8} {:>5.1}%", 552 | style("DC").blue(), 553 | style( 554 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::DC_PRED, frame_type) 555 | ) 556 | .cyan() 557 | ); 558 | 559 | info!( 560 | "{:8} {:>5.1}%", 561 | style("Vert").blue(), 562 | style( 563 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::V_PRED, frame_type) 564 | ) 565 | .cyan() 566 | ); 567 | info!( 568 | "{:8} {:>5.1}%", 569 | style("Horiz").blue(), 570 | style( 571 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::H_PRED, frame_type) 572 | ) 573 | .cyan() 574 | ); 575 | info!( 576 | "{:8} {:>5.1}%", 577 | style("Paeth").blue(), 578 | style( 579 | self.get_luma_pred_mode_pct_by_frame_type( 580 | PredictionMode::PAETH_PRED, 581 | frame_type 582 | ) 583 | ) 584 | .cyan() 585 | ); 586 | info!( 587 | "{:8} {:>5.1}%", 588 | style("Smooth").blue(), 589 | style( 590 | self.get_luma_pred_mode_pct_by_frame_type( 591 | PredictionMode::SMOOTH_PRED, 592 | frame_type 593 | ) 594 | ) 595 | .cyan() 596 | ); 597 | info!( 598 | "{:8} {:>5.1}%", 599 | style("Smooth V").blue(), 600 | style(self.get_luma_pred_mode_pct_by_frame_type( 601 | PredictionMode::SMOOTH_V_PRED, 602 | frame_type 603 | )) 604 | .cyan() 605 | ); 606 | info!( 607 | "{:8} {:>5.1}%", 608 | style("Smooth H").blue(), 609 | style(self.get_luma_pred_mode_pct_by_frame_type( 610 | PredictionMode::SMOOTH_H_PRED, 611 | frame_type 612 | )) 613 | .cyan() 614 | ); 615 | info!( 616 | "{:8} {:>5.1}%", 617 | style("45-Deg").blue(), 618 | style( 619 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::D45_PRED, frame_type) 620 | ) 621 | .cyan() 622 | ); 623 | info!( 624 | "{:8} {:>5.1}%", 625 | style("63-Deg").blue(), 626 | style( 627 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::D67_PRED, frame_type) 628 | ) 629 | .cyan() 630 | ); 631 | info!( 632 | "{:8} {:>5.1}%", 633 | style("117-Deg").blue(), 634 | style( 635 | self.get_luma_pred_mode_pct_by_frame_type( 636 | PredictionMode::D113_PRED, 637 | frame_type 638 | ) 639 | ) 640 | .cyan() 641 | ); 642 | info!( 643 | "{:8} {:>5.1}%", 644 | style("135-Deg").blue(), 645 | style( 646 | self.get_luma_pred_mode_pct_by_frame_type( 647 | PredictionMode::D135_PRED, 648 | frame_type 649 | ) 650 | ) 651 | .cyan() 652 | ); 653 | info!( 654 | "{:8} {:>5.1}%", 655 | style("153-Deg").blue(), 656 | style( 657 | self.get_luma_pred_mode_pct_by_frame_type( 658 | PredictionMode::D157_PRED, 659 | frame_type 660 | ) 661 | ) 662 | .cyan() 663 | ); 664 | info!( 665 | "{:8} {:>5.1}%", 666 | style("207-Deg").blue(), 667 | style( 668 | self.get_luma_pred_mode_pct_by_frame_type( 669 | PredictionMode::D203_PRED, 670 | frame_type 671 | ) 672 | ) 673 | .cyan() 674 | ); 675 | } else if frame_type == FrameType::INTER { 676 | info!( 677 | "{:15} {:>5.1}%", 678 | style("Nearest").blue(), 679 | style( 680 | self.get_luma_pred_mode_pct_by_frame_type( 681 | PredictionMode::NEARESTMV, 682 | frame_type 683 | ) 684 | ) 685 | .cyan() 686 | ); 687 | info!( 688 | "{:15} {:>5.1}%", 689 | style("Near-0").blue(), 690 | style( 691 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::NEAR0MV, frame_type) 692 | ) 693 | .cyan() 694 | ); 695 | info!( 696 | "{:15} {:>5.1}%", 697 | style("Near-1").blue(), 698 | style( 699 | self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::NEAR1MV, frame_type) 700 | ) 701 | .cyan() 702 | ); 703 | info!( 704 | "{:15} {:>5.1}%", 705 | style("Near-Near-0").blue(), 706 | style(self.get_luma_pred_mode_pct_by_frame_type( 707 | PredictionMode::NEAR_NEAR0MV, 708 | frame_type 709 | )) 710 | .cyan() 711 | ); 712 | info!( 713 | "{:15} {:>5.1}%", 714 | style("Near-Near-1").blue(), 715 | style(self.get_luma_pred_mode_pct_by_frame_type( 716 | PredictionMode::NEAR_NEAR1MV, 717 | frame_type 718 | )) 719 | .cyan() 720 | ); 721 | info!( 722 | "{:15} {:>5.1}%", 723 | style("Near-Near-2").blue(), 724 | style(self.get_luma_pred_mode_pct_by_frame_type( 725 | PredictionMode::NEAR_NEAR2MV, 726 | frame_type 727 | )) 728 | .cyan() 729 | ); 730 | info!( 731 | "{:15} {:>5.1}%", 732 | style("New").blue(), 733 | style(self.get_luma_pred_mode_pct_by_frame_type(PredictionMode::NEWMV, frame_type)) 734 | .cyan() 735 | ); 736 | info!( 737 | "{:15} {:>5.1}%", 738 | style("New-New").blue(), 739 | style( 740 | self.get_luma_pred_mode_pct_by_frame_type( 741 | PredictionMode::NEW_NEWMV, 742 | frame_type 743 | ) 744 | ) 745 | .cyan() 746 | ); 747 | info!( 748 | "{:15} {:>5.1}%", 749 | style("Nearest-Nearest").blue(), 750 | style(self.get_luma_pred_mode_pct_by_frame_type( 751 | PredictionMode::NEAREST_NEARESTMV, 752 | frame_type 753 | )) 754 | .cyan() 755 | ); 756 | info!( 757 | "{:15} {:>5.1}%", 758 | style("Global-Global").blue(), 759 | style(self.get_luma_pred_mode_pct_by_frame_type( 760 | PredictionMode::GLOBAL_GLOBALMV, 761 | frame_type 762 | )) 763 | .cyan() 764 | ); 765 | } 766 | } 767 | 768 | #[cfg(feature = "binary")] 769 | fn print_chroma_prediction_mode_summary_by_frame_type( 770 | &self, 771 | frame_type: FrameType, 772 | type_label: char, 773 | ) { 774 | info!( 775 | "{}", 776 | style(format!("{} Frame Chroma Modes", type_label)).yellow() 777 | ); 778 | if frame_type == FrameType::KEY { 779 | info!( 780 | "{:8} {:>5.1}%", 781 | style("DC").blue(), 782 | style( 783 | self.get_chroma_pred_mode_pct_by_frame_type( 784 | PredictionMode::DC_PRED, 785 | frame_type 786 | ) 787 | ) 788 | .cyan() 789 | ); 790 | 791 | info!( 792 | "{:8} {:>5.1}%", 793 | style("Vert").blue(), 794 | style( 795 | self.get_chroma_pred_mode_pct_by_frame_type(PredictionMode::V_PRED, frame_type) 796 | ) 797 | .cyan() 798 | ); 799 | info!( 800 | "{:8} {:>5.1}%", 801 | style("Horiz").blue(), 802 | style( 803 | self.get_chroma_pred_mode_pct_by_frame_type(PredictionMode::H_PRED, frame_type) 804 | ) 805 | .cyan() 806 | ); 807 | info!( 808 | "{:8} {:>5.1}%", 809 | style("Paeth").blue(), 810 | style(self.get_chroma_pred_mode_pct_by_frame_type( 811 | PredictionMode::PAETH_PRED, 812 | frame_type 813 | )) 814 | .cyan() 815 | ); 816 | info!( 817 | "{:8} {:>5.1}%", 818 | style("Smooth").blue(), 819 | style(self.get_chroma_pred_mode_pct_by_frame_type( 820 | PredictionMode::SMOOTH_PRED, 821 | frame_type 822 | )) 823 | .cyan() 824 | ); 825 | info!( 826 | "{:8} {:>5.1}%", 827 | style("Smooth V").blue(), 828 | style(self.get_chroma_pred_mode_pct_by_frame_type( 829 | PredictionMode::SMOOTH_V_PRED, 830 | frame_type 831 | )) 832 | .cyan() 833 | ); 834 | info!( 835 | "{:8} {:>5.1}%", 836 | style("Smooth H").blue(), 837 | style(self.get_chroma_pred_mode_pct_by_frame_type( 838 | PredictionMode::SMOOTH_H_PRED, 839 | frame_type 840 | )) 841 | .cyan() 842 | ); 843 | info!( 844 | "{:8} {:>5.1}%", 845 | style("45-Deg").blue(), 846 | style( 847 | self.get_chroma_pred_mode_pct_by_frame_type( 848 | PredictionMode::D45_PRED, 849 | frame_type 850 | ) 851 | ) 852 | .cyan() 853 | ); 854 | info!( 855 | "{:8} {:>5.1}%", 856 | style("63-Deg").blue(), 857 | style( 858 | self.get_chroma_pred_mode_pct_by_frame_type( 859 | PredictionMode::D67_PRED, 860 | frame_type 861 | ) 862 | ) 863 | .cyan() 864 | ); 865 | info!( 866 | "{:8} {:>5.1}%", 867 | style("117-Deg").blue(), 868 | style( 869 | self.get_chroma_pred_mode_pct_by_frame_type( 870 | PredictionMode::D113_PRED, 871 | frame_type 872 | ) 873 | ) 874 | .cyan() 875 | ); 876 | info!( 877 | "{:8} {:>5.1}%", 878 | style("135-Deg").blue(), 879 | style( 880 | self.get_chroma_pred_mode_pct_by_frame_type( 881 | PredictionMode::D135_PRED, 882 | frame_type 883 | ) 884 | ) 885 | .cyan() 886 | ); 887 | info!( 888 | "{:8} {:>5.1}%", 889 | style("153-Deg").blue(), 890 | style( 891 | self.get_chroma_pred_mode_pct_by_frame_type( 892 | PredictionMode::D157_PRED, 893 | frame_type 894 | ) 895 | ) 896 | .cyan() 897 | ); 898 | info!( 899 | "{:8} {:>5.1}%", 900 | style("207-Deg").blue(), 901 | style( 902 | self.get_chroma_pred_mode_pct_by_frame_type( 903 | PredictionMode::D203_PRED, 904 | frame_type 905 | ) 906 | ) 907 | .cyan() 908 | ); 909 | info!( 910 | "{:8} {:>5.1}%", 911 | style("UV CFL").blue(), 912 | style(self.get_chroma_pred_mode_pct_by_frame_type( 913 | PredictionMode::UV_CFL_PRED, 914 | frame_type 915 | )) 916 | .cyan() 917 | ); 918 | } else if frame_type == FrameType::INTER { 919 | info!( 920 | "{:15} {:>5.1}%", 921 | style("Nearest").blue(), 922 | style( 923 | self.get_chroma_pred_mode_pct_by_frame_type( 924 | PredictionMode::NEARESTMV, 925 | frame_type 926 | ) 927 | ) 928 | .cyan() 929 | ); 930 | info!( 931 | "{:15} {:>5.1}%", 932 | style("Near-0").blue(), 933 | style( 934 | self.get_chroma_pred_mode_pct_by_frame_type( 935 | PredictionMode::NEAR0MV, 936 | frame_type 937 | ) 938 | ) 939 | .cyan() 940 | ); 941 | info!( 942 | "{:15} {:>5.1}%", 943 | style("Near-1").blue(), 944 | style( 945 | self.get_chroma_pred_mode_pct_by_frame_type( 946 | PredictionMode::NEAR1MV, 947 | frame_type 948 | ) 949 | ) 950 | .cyan() 951 | ); 952 | info!( 953 | "{:15} {:>5.1}%", 954 | style("Near-Near-0").blue(), 955 | style(self.get_chroma_pred_mode_pct_by_frame_type( 956 | PredictionMode::NEAR_NEAR0MV, 957 | frame_type 958 | )) 959 | .cyan() 960 | ); 961 | info!( 962 | "{:15} {:>5.1}%", 963 | style("Near-Near-1").blue(), 964 | style(self.get_chroma_pred_mode_pct_by_frame_type( 965 | PredictionMode::NEAR_NEAR1MV, 966 | frame_type 967 | )) 968 | .cyan() 969 | ); 970 | info!( 971 | "{:15} {:>5.1}%", 972 | style("Near-Near-2").blue(), 973 | style(self.get_chroma_pred_mode_pct_by_frame_type( 974 | PredictionMode::NEAR_NEAR2MV, 975 | frame_type 976 | )) 977 | .cyan() 978 | ); 979 | info!( 980 | "{:15} {:>5.1}%", 981 | style("New").blue(), 982 | style( 983 | self.get_chroma_pred_mode_pct_by_frame_type(PredictionMode::NEWMV, frame_type) 984 | ) 985 | .cyan() 986 | ); 987 | info!( 988 | "{:15} {:>5.1}%", 989 | style("New-New").blue(), 990 | style( 991 | self.get_chroma_pred_mode_pct_by_frame_type( 992 | PredictionMode::NEW_NEWMV, 993 | frame_type 994 | ) 995 | ) 996 | .cyan() 997 | ); 998 | info!( 999 | "{:15} {:>5.1}%", 1000 | style("Nearest-Nearest").blue(), 1001 | style(self.get_chroma_pred_mode_pct_by_frame_type( 1002 | PredictionMode::NEAREST_NEARESTMV, 1003 | frame_type 1004 | )) 1005 | .cyan() 1006 | ); 1007 | info!( 1008 | "{:15} {:>5.1}%", 1009 | style("Global-Global").blue(), 1010 | style(self.get_chroma_pred_mode_pct_by_frame_type( 1011 | PredictionMode::GLOBAL_GLOBALMV, 1012 | frame_type 1013 | )) 1014 | .cyan() 1015 | ); 1016 | } 1017 | } 1018 | } 1019 | 1020 | #[cfg(feature = "binary")] 1021 | fn secs_to_human_time(mut secs: u64, always_show_hours: bool) -> String { 1022 | let mut mins = secs / 60; 1023 | secs %= 60; 1024 | let hours = mins / 60; 1025 | mins %= 60; 1026 | if hours > 0 || always_show_hours { 1027 | format!("{:02}:{:02}:{:02}", hours, mins, secs) 1028 | } else { 1029 | format!("{:02}:{:02}", mins, secs) 1030 | } 1031 | } 1032 | 1033 | #[derive(Debug, Clone, Serialize, Deserialize)] 1034 | pub struct SerializableProgressInfo { 1035 | pub frame_rate: (u64, u64), 1036 | pub frame_info: Vec, 1037 | pub encoded_size: usize, 1038 | pub keyframes: BTreeSet, 1039 | pub completed_segments: BTreeSet, 1040 | pub next_analysis_frame: usize, 1041 | // Wall encoding time elapsed so far, in seconds 1042 | #[serde(default)] 1043 | pub elapsed_time: u64, 1044 | #[serde(default)] 1045 | pub encoding_stats: (SerializableEncoderStats, SerializableEncoderStats), 1046 | pub total_frames: usize, 1047 | #[serde(default)] 1048 | pub frame_limit: Option, 1049 | #[serde(default)] 1050 | pub segment_idx: usize, 1051 | } 1052 | 1053 | impl From<&ProgressInfo> for SerializableProgressInfo { 1054 | fn from(other: &ProgressInfo) -> Self { 1055 | SerializableProgressInfo { 1056 | frame_rate: (other.frame_rate.num, other.frame_rate.den), 1057 | frame_info: other 1058 | .frame_info 1059 | .iter() 1060 | .map(SerializableFrameSummary::from) 1061 | .collect(), 1062 | encoded_size: other.encoded_size, 1063 | keyframes: other.keyframes.clone(), 1064 | next_analysis_frame: other.next_analysis_frame, 1065 | completed_segments: other.completed_segments.clone(), 1066 | elapsed_time: other.elapsed_time() as u64, 1067 | encoding_stats: ( 1068 | (&other.encoding_stats.0).into(), 1069 | (&other.encoding_stats.1).into(), 1070 | ), 1071 | total_frames: other.total_frames, 1072 | frame_limit: other.frame_limit, 1073 | segment_idx: other.segment_idx, 1074 | } 1075 | } 1076 | } 1077 | 1078 | impl From<&SerializableProgressInfo> for ProgressInfo { 1079 | fn from(other: &SerializableProgressInfo) -> Self { 1080 | ProgressInfo { 1081 | frame_rate: Rational::new(other.frame_rate.0, other.frame_rate.1), 1082 | frame_info: other.frame_info.iter().map(FrameSummary::from).collect(), 1083 | encoded_size: other.encoded_size, 1084 | keyframes: other.keyframes.clone(), 1085 | next_analysis_frame: other.next_analysis_frame, 1086 | completed_segments: other.completed_segments.clone(), 1087 | time_started: Instant::now() - Duration::from_secs(other.elapsed_time), 1088 | segment_idx: other.segment_idx, 1089 | encoding_stats: ( 1090 | (&other.encoding_stats.0).into(), 1091 | (&other.encoding_stats.1).into(), 1092 | ), 1093 | total_frames: other.total_frames, 1094 | frame_limit: other.frame_limit, 1095 | } 1096 | } 1097 | } 1098 | 1099 | #[derive(Debug, Clone, Copy)] 1100 | pub struct FrameSummary { 1101 | /// Frame size in bytes 1102 | pub size: usize, 1103 | pub frame_type: FrameType, 1104 | /// QP selected for the frame. 1105 | pub qp: u8, 1106 | } 1107 | 1108 | impl From> for FrameSummary { 1109 | fn from(packet: Packet) -> Self { 1110 | Self { 1111 | size: packet.data.len(), 1112 | frame_type: packet.frame_type, 1113 | qp: packet.qp, 1114 | } 1115 | } 1116 | } 1117 | 1118 | #[derive(Debug, Clone, Copy, Serialize, Deserialize)] 1119 | pub struct SerializableFrameSummary { 1120 | pub size: usize, 1121 | pub frame_type: u8, 1122 | pub qp: u8, 1123 | } 1124 | 1125 | impl From<&FrameSummary> for SerializableFrameSummary { 1126 | fn from(summary: &FrameSummary) -> Self { 1127 | SerializableFrameSummary { 1128 | size: summary.size, 1129 | frame_type: summary.frame_type as u8, 1130 | qp: summary.qp, 1131 | } 1132 | } 1133 | } 1134 | 1135 | impl From<&SerializableFrameSummary> for FrameSummary { 1136 | fn from(summary: &SerializableFrameSummary) -> Self { 1137 | FrameSummary { 1138 | size: summary.size, 1139 | frame_type: match summary.frame_type { 1140 | 0 => FrameType::KEY, 1141 | 1 => FrameType::INTER, 1142 | 2 => FrameType::INTRA_ONLY, 1143 | 3 => FrameType::SWITCH, 1144 | _ => unreachable!(), 1145 | }, 1146 | qp: summary.qp, 1147 | } 1148 | } 1149 | } 1150 | 1151 | impl Serialize for FrameSummary { 1152 | fn serialize(&self, serializer: S) -> Result 1153 | where 1154 | S: Serializer, 1155 | { 1156 | let ser = SerializableFrameSummary::from(self); 1157 | ser.serialize(serializer) 1158 | } 1159 | } 1160 | 1161 | impl<'de> Deserialize<'de> for FrameSummary { 1162 | fn deserialize(deserializer: D) -> Result 1163 | where 1164 | D: Deserializer<'de>, 1165 | { 1166 | let de = SerializableFrameSummary::deserialize(deserializer)?; 1167 | Ok(FrameSummary::from(&de)) 1168 | } 1169 | } 1170 | 1171 | pub const TX_TYPES: usize = 16; 1172 | pub const PREDICTION_MODES: usize = 34; 1173 | 1174 | #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] 1175 | pub struct SerializableEncoderStats { 1176 | /// Stores count of pixels belonging to each block size in this frame 1177 | pub block_size_counts: [usize; BlockSize::BLOCK_SIZES_ALL], 1178 | /// Stores count of pixels belonging to skip blocks in this frame 1179 | pub skip_block_count: usize, 1180 | /// Stores count of pixels belonging to each transform type in this frame 1181 | pub tx_type_counts: [usize; TX_TYPES], 1182 | /// Stores count of pixels belonging to each luma prediction mode in this frame 1183 | pub luma_pred_mode_counts: ArrayVec<[usize; PREDICTION_MODES]>, 1184 | /// Stores count of pixels belonging to each chroma prediction mode in this frame 1185 | pub chroma_pred_mode_counts: ArrayVec<[usize; PREDICTION_MODES]>, 1186 | } 1187 | 1188 | impl From<&EncoderStats> for SerializableEncoderStats { 1189 | fn from(stats: &EncoderStats) -> Self { 1190 | SerializableEncoderStats { 1191 | block_size_counts: stats.block_size_counts, 1192 | skip_block_count: stats.skip_block_count, 1193 | tx_type_counts: stats.tx_type_counts, 1194 | luma_pred_mode_counts: stats.luma_pred_mode_counts.clone(), 1195 | chroma_pred_mode_counts: stats.chroma_pred_mode_counts.clone(), 1196 | } 1197 | } 1198 | } 1199 | 1200 | impl From<&SerializableEncoderStats> for EncoderStats { 1201 | fn from(stats: &SerializableEncoderStats) -> Self { 1202 | EncoderStats { 1203 | block_size_counts: stats.block_size_counts, 1204 | skip_block_count: stats.skip_block_count, 1205 | tx_type_counts: stats.tx_type_counts, 1206 | luma_pred_mode_counts: stats.luma_pred_mode_counts.clone(), 1207 | chroma_pred_mode_counts: stats.chroma_pred_mode_counts.clone(), 1208 | } 1209 | } 1210 | } 1211 | --------------------------------------------------------------------------------