From 2a89943f45031278b9f032c3be16d3a21a67ccd5 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Thu, 7 Mar 2024 03:36:42 -0500
Subject: [PATCH] Mocap job: read video file dimensions

---
 Cargo.lock                                    |  1 +
 crates/lib/files/videos/src/get_mp4_info.rs   | 27 +++++++-
 crates/service/job/inference_job/Cargo.toml   |  5 +-
 .../job_types/mocap/mocap_net/process_job.rs  | 66 ++++++++++++++-----
 .../upload_video_media_file_handler.rs        |  4 +-
 5 files changed, 78 insertions(+), 25 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 70f3cb9a12..d6f1270264 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3241,6 +3241,7 @@ dependencies = [
  "tokens",
  "tokio",
  "tts_common",
+ "videos",
  "walkdir",
 ]
 
diff --git a/crates/lib/files/videos/src/get_mp4_info.rs b/crates/lib/files/videos/src/get_mp4_info.rs
index 2a0327bf6c..11bd2f60fc 100644
--- a/crates/lib/files/videos/src/get_mp4_info.rs
+++ b/crates/lib/files/videos/src/get_mp4_info.rs
@@ -6,30 +6,47 @@ use mp4::TrackType;
 
 use errors::AnyhowResult;
 
+#[derive(Debug, Clone)]
 pub struct Mp4Info {
   /// Framerate of the longest video track.
   pub framerate: f64,
+
   /// Duration of the entire mp4.
   pub duration_millis: u128,
+
+  /// Width of the longest video track.
+  pub width: u16,
+
+  /// Height of the longest video track.
+  pub height: u16,
 }
 
-pub fn get_mp4_info_for_bytes(file_bytes: &[u8], file_size: u64) -> AnyhowResult<Mp4Info> {
+pub fn get_mp4_info_for_bytes(file_bytes: &[u8]) -> AnyhowResult<Mp4Info> {
+  get_mp4_info_for_bytes_and_len(file_bytes, file_bytes.len())
+}
+
+pub fn get_mp4_info_for_bytes_and_len(file_bytes: &[u8], file_size: usize) -> AnyhowResult<Mp4Info> {
   let reader = BufReader::new(Cursor::new(file_bytes));
-  get_mp4_info(reader, file_size)
+  get_mp4_info(reader, file_size as u64)
 }
 
 pub fn get_mp4_info<T: Seek + Read>(reader: T, file_size: u64) -> AnyhowResult<Mp4Info> {
   let mp4 = mp4::Mp4Reader::read_header(reader, file_size)?;
 
-  let mut framerate = 0.0;
   let mut longest_duration = Duration::from_secs(0);
 
+  let mut framerate = 0.0;
+  let mut width = 0;
+  let mut height = 0;
+
   for track in mp4.tracks().values() {
     match track.track_type() {
       Ok(TrackType::Video) => {
         if track.duration() > longest_duration {
           longest_duration = track.duration();
           framerate = track.frame_rate();
+          width = track.width();
+          height = track.height();
         }
       }
       Err(err) => {
@@ -43,6 +60,8 @@ pub fn get_mp4_info<T: Seek + Read>(reader: T, file_size: u64) -> AnyhowResult<M
   Ok(Mp4Info {
     framerate,
     duration_millis: mp4.duration().as_millis(),
+    width,
+    height,
   })
 }
 
@@ -76,5 +95,7 @@ pub mod tests {
 
     assert_eq!(info.framerate, 30.0);
     assert_eq!(info.duration_millis, 15168);
+    assert_eq!(info.width, 640);
+    assert_eq!(info.height, 480);
   }
 }
diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml
index 1c20087956..69ce218530 100644
--- a/crates/service/job/inference_job/Cargo.toml
+++ b/crates/service/job/inference_job/Cargo.toml
@@ -19,8 +19,6 @@ actix_helpers = { path = "../../../lib/actix_helpers" }
 bootstrap = { path = "../../../lib/bootstrap" }
 buckets = { path = "../../../_types/buckets" }
 cloud_storage = { path = "../../../lib/api_clients/cloud_storage" }
-google_drive_common = { path = "../../../lib/api_clients/google_drive_common" }
-thumbnail_generator = { path = "../../../lib/api_clients/thumbnail_generator"}
 collections = { path = "../../../lib/collections" }
 composite_identifiers = { path = "../../../_types/composite_identifiers" }
 concurrency = { path = "../../../lib/concurrency" }
@@ -29,6 +27,7 @@ easyenv = { path = "../../../lib/easyenv" }
 enums = { path = "../../../_types/enums" }
 errors = { path = "../../../lib/errors" }
 filesys = { path = "../../../lib/files/filesys" }
+google_drive_common = { path = "../../../lib/api_clients/google_drive_common" }
 hashing = { path = "../../../lib/files/hashing" }
 images = { path = "../../../lib/files/images" }
 jobs_common = { path = "../../../lib/jobs_common" }
@@ -39,8 +38,10 @@ mimetypes = { path = "../../../lib/files/mimetypes" }
 mysql_queries = { path = "../../../lib/mysql_queries" }
 redis_common = { path = "../../../lib/redis_common" }
 subprocess_common = { path = "../../../lib/subprocess_common" }
+thumbnail_generator = { path = "../../../lib/api_clients/thumbnail_generator"}
 tokens = { path = "../../../_types/tokens" }
 tts_common = { path = "../../../lib/tts_common" }
+videos = { path = "../../../lib/files/videos" }
 
 # External
 anyhow = "1.0.75"
diff --git a/crates/service/job/inference_job/src/job/job_types/mocap/mocap_net/process_job.rs b/crates/service/job/inference_job/src/job/job_types/mocap/mocap_net/process_job.rs
index aa20750272..9fe398efef 100644
--- a/crates/service/job/inference_job/src/job/job_types/mocap/mocap_net/process_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/mocap/mocap_net/process_job.rs
@@ -1,4 +1,6 @@
-use std::fs::read_to_string;
+use std::fs::{File, read_to_string};
+use std::io::{BufReader, Read};
+use std::path::PathBuf;
 use std::time::Instant;
 
 use anyhow::anyhow;
@@ -14,6 +16,7 @@ use hashing::sha256::sha256_hash_file::sha256_hash_file;
 use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs::Mc;
 use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
 use mysql_queries::queries::media_files::create::insert_media_file_from_mocapnet::{insert_media_file_from_mocapnet, InsertArgs};
+use videos::get_mp4_info::{get_mp4_info, get_mp4_info_for_bytes, Mp4Info};
 
 use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity};
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
@@ -53,6 +56,23 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result<JobSuccessR
 
     let job_args = validate_job(job)?;
 
+    // ==================== JOB ARGS ==================== //
+
+    let maybe_args = job.maybe_inference_args
+        .as_ref()
+        .map(|args| args.args.as_ref())
+        .flatten();
+
+    let poly_args = match maybe_args {
+        None => return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))),
+        Some(args) => args,
+    };
+
+    let mc_args = match poly_args {
+        Mc(args) => args,
+        _ => return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))),
+    };
+
     // ==================== TEMP DIR ==================== //
 
     let work_temp_dir = format!("temp_mocapnet_inference_{}", job.id.0);
@@ -86,6 +106,22 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result<JobSuccessR
 
     info!("Used video file: {:?}", usable_video_path);
 
+    // ==================== EXTRACT FILE METADATA (IF POSSIBLE) ==================== //
+
+    // TODO(bt): better check for video file type
+    let video_read_result = try_get_video_info(&video_path.filesystem_path);
+
+    let mut maybe_width = mc_args.maybe_size1.clone();
+    let mut maybe_height = mc_args.maybe_size2.clone();
+
+    // NB: Fail open. In the event this fails, we want to continue processing.
+    if let Some(video_info) = video_read_result {
+        info!("Read video info from file: {:?}", video_info);
+
+        maybe_width = Some(video_info.width as i32);
+        maybe_height = Some(video_info.height as i32);
+    }
+
     // ==================== SETUP FOR INFERENCE ==================== //
 
     info!("Ready for MocapNET inference...");
@@ -102,21 +138,6 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result<JobSuccessR
 
     // TODO: Limit output length for non-premium (???)
 
-    let maybe_args = job.maybe_inference_args
-        .as_ref()
-        .map(|args| args.args.as_ref())
-        .flatten();
-
-    let poly_args = match maybe_args {
-        None => return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))),
-        Some(args) => args,
-    };
-
-    let mc_args = match poly_args {
-        Mc(args) => args,
-        _ => return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))),
-    };
-
     // ==================== RUN INFERENCE SCRIPT ==================== //
     let stderr_output_file = work_temp_dir.path().join("stderr.txt");
     let inference_start_time = Instant::now();
@@ -130,8 +151,8 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result<JobSuccessR
             maybe_ik3: &mc_args.maybe_ik3,
             maybe_smoothing1: &mc_args.maybe_smoothing1,
             maybe_smoothing2: &mc_args.maybe_smoothing2,
-            maybe_size1: &mc_args.maybe_size1,
-            maybe_size2: &mc_args.maybe_size2,
+            maybe_size1: &maybe_width,
+            maybe_size2: &maybe_height,
             stderr_output_file: &stderr_output_file,
         });
 
@@ -255,3 +276,12 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result<JobSuccessR
         inference_duration,
     })
 }
+
+// Fail open if this doesn't work
+fn try_get_video_info(video_filename: &PathBuf) -> Option<Mp4Info> {
+    let video_file = File::open(video_filename).ok()?;
+    let file_length = video_file.metadata().ok()?.len();
+    let reader = BufReader::new(video_file);
+    let info = get_mp4_info(reader, file_length).ok()?;
+    Some(info)
+}
\ No newline at end of file
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs
index 3a01c23f09..6eac332b77 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs
@@ -18,7 +18,7 @@ use mimetypes::mimetype_to_extension::mimetype_to_extension;
 use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
 use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType};
 use tokens::tokens::media_files::MediaFileToken;
-use videos::get_mp4_info::{get_mp4_info, get_mp4_info_for_bytes};
+use videos::get_mp4_info::{get_mp4_info, get_mp4_info_for_bytes, get_mp4_info_for_bytes_and_len};
 
 use crate::http_server::endpoints::engine::create_scene_handler::{CreateSceneError, CreateSceneSuccessResponse};
 use crate::http_server::endpoints::media_files::upload_video::drain_multipart_request::drain_multipart_request;
@@ -176,7 +176,7 @@ pub async fn upload_video_media_file_handler(
 
   // ==================== FRAME RATE ==================== //
 
-  let mp4_info = get_mp4_info_for_bytes(file_bytes.as_ref(), file_size_bytes as u64)
+  let mp4_info = get_mp4_info_for_bytes(file_bytes.as_ref())
       .map_err(|err| {
         warn!("Error reading mp4 info: {:?}", err);
         VideoMediaFileUploadError::ServerError