From 6afe1c168be5fcf13ee0220623053606b2eaeceb Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Tue, 25 May 2021 01:50:02 -0400
Subject: [PATCH] save duration millis

---
 migrations/2021-05-15-071601_assets/up.sql  |  1 +
 models/Wav2Lip/vocodes_process_upload.py    | 10 +++++++++-
 src/job_queries/w2l_download_job_queries.rs |  9 ++++++---
 src/w2l_download_job.rs                     |  4 +++-
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/migrations/2021-05-15-071601_assets/up.sql b/migrations/2021-05-15-071601_assets/up.sql
index a96a3a2693..1c797e8d60 100644
--- a/migrations/2021-05-15-071601_assets/up.sql
+++ b/migrations/2021-05-15-071601_assets/up.sql
@@ -190,6 +190,7 @@ CREATE TABLE w2l_templates (
   frame_height INT(10) NOT NULL DEFAULT 0,
   frame_count INT(10) NOT NULL DEFAULT 0,
   fps FLOAT NOT NULL DEFAULT 0.0,
+  duration_millis INT(1) NOT NULL DEFAULT 0,
 
   -- Where the file was originally downloaded (if it was downloaded)
   original_download_url VARCHAR(512) DEFAULT NULL,
diff --git a/models/Wav2Lip/vocodes_process_upload.py b/models/Wav2Lip/vocodes_process_upload.py
index 4359c1971e..d766af1ec1 100644
--- a/models/Wav2Lip/vocodes_process_upload.py
+++ b/models/Wav2Lip/vocodes_process_upload.py
@@ -271,7 +271,9 @@ def main(tempdir):
     is_video = False
     frame_w = 0
     frame_h = 0
-    fps = 0
+    frame_count = 0
+    fps = 0.0
+    duration_millis = 0.0
 
     if not os.path.isfile(args.image_or_video_filename):
         raise ValueError('--image_or_video_filename argument must be a valid path to video/image file')
@@ -282,6 +284,7 @@ def main(tempdir):
         fps = args.fps
         frame_h, frame_w = full_frames[0].shape[:-1]
         is_video = False
+        frame_count = len(full_frames)
 
     else:
         video_stream = cv2.VideoCapture(args.image_or_video_filename)
@@ -313,8 +316,12 @@ def main(tempdir):
 
         if len(full_frames) < 2:
             is_video = False
+            frame_count = len(full_frames)
         else:
             is_video = True
+            frame_count = len(full_frames)
+            duration_seconds = frame_count / fps
+            duration_millis = int(duration_seconds * 1000)
 
     print("Number of frames available for inference: "+str(len(full_frames)), flush=True)
     print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True)
@@ -339,6 +346,7 @@ def main(tempdir):
 
     if is_video:
         metadata['fps'] = fps
+        metadata['duration_millis'] = duration_millis
 
     with open(args.output_metadata_filename, 'w') as json_file:
         json.dump(metadata, json_file)
diff --git a/src/job_queries/w2l_download_job_queries.rs b/src/job_queries/w2l_download_job_queries.rs
index 67d66a0343..01de3299de 100644
--- a/src/job_queries/w2l_download_job_queries.rs
+++ b/src/job_queries/w2l_download_job_queries.rs
@@ -120,7 +120,8 @@ pub async fn insert_w2l_template(pool: &MySqlPool,
                                  frame_width: u32,
                                  frame_height: u32,
                                  frame_count: u64,
-                                 fps: f32)
+                                 fps: f32,
+                                 duration_millis: u64)
   -> AnyhowResult<u64>
 {
   let model_token = random_token(32);
@@ -149,7 +150,8 @@ SET
   frame_width = ?,
   frame_height = ?,
   frame_count = ?,
-  fps = ?
+  fps = ?,
+  duration_millis = ?
         "#,
       model_token,
       template_type,
@@ -166,7 +168,8 @@ SET
       frame_width,
       frame_height,
       frame_count,
-      fps
+      fps,
+      duration_millis
     )
     .execute(pool)
     .await;
diff --git a/src/w2l_download_job.rs b/src/w2l_download_job.rs
index 21c8c6bbca..e020182757 100644
--- a/src/w2l_download_job.rs
+++ b/src/w2l_download_job.rs
@@ -227,6 +227,7 @@ struct FileMetadata {
   pub height: u32,
   pub num_frames: u64,
   pub fps: Option<f32>,
+  pub duration_millis: Option<u64>,
 }
 
 fn read_metadata_file(filename: &str) -> AnyhowResult<FileMetadata> {
@@ -381,7 +382,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
     file_metadata.width,
     file_metadata.height,
     file_metadata.num_frames,
-    file_metadata.fps.unwrap_or(0.0f32))
+    file_metadata.fps.unwrap_or(0.0f32),
+    file_metadata.duration_millis.unwrap_or(0))
     .await?;
 
   info!("Job {} complete success! Downloaded, processed, and uploaded. Saved model record: {}",