From 0b9b7f7c4dfb42121089a5f5066cd7e727a42deb Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sun, 11 Dec 2022 21:35:09 -0500
Subject: [PATCH] Calculate wav metadata

---
 voice_conversion/softvc/.gitignore       |  5 +++
 voice_conversion/softvc/cli_inference.py | 48 ++++++++++++++++++++++--
 voice_conversion/softvc/requirements.txt |  7 ++++
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/voice_conversion/softvc/.gitignore b/voice_conversion/softvc/.gitignore
index 7a2662b..a06b66e 100644
--- a/voice_conversion/softvc/.gitignore
+++ b/voice_conversion/softvc/.gitignore
@@ -8,3 +8,8 @@ __pycache__/
 *.swo
 *.swp
 *~
+
+# Default output files
+output.wav
+metadata.json
+
diff --git a/voice_conversion/softvc/cli_inference.py b/voice_conversion/softvc/cli_inference.py
index 35705e2..7490479 100755
--- a/voice_conversion/softvc/cli_inference.py
+++ b/voice_conversion/softvc/cli_inference.py
@@ -3,7 +3,11 @@
 import IPython.display as display
 import argparse
 import importlib.util
-import soundfile as sf
+import json
+import magic
+import os
+import soundfile
+import subprocess
 import sys
 import sys  
 import torch
@@ -29,7 +33,8 @@ parser = argparse.ArgumentParser()
 parser.add_argument('--acoustic_model_filename', type=str, required=True)
 parser.add_argument('--hifigan_model_filename', type=str, required=True)
 parser.add_argument('--source_audio_filename', type=str, required=True)
-parser.add_argument('--output_audio_filename', type=str, required=True)
+parser.add_argument('--output_audio_filename', type=str, default='output.wav')
+parser.add_argument('--output_metadata_filename', type=str, default='metadata.json')
 args = parser.parse_args()
 
 print('\n'.join(f'[arg] {k} = {v}' for k, v in vars(args).items()))
@@ -115,8 +120,43 @@ def do_voice_conversion(hubert, acoustic, hifigan, audio_source):
         end = time()
         print(f"time for wavform {end-start}")
         return target
+
 # ===========================================
 
+def save_output_audio(audio_data, output_audio_filename):
+  soundfile.write(args.output_audio_filename, audio_data.squeeze().cpu(), samplerate=16000)
+
+def save_metadata(audio_filename, output_metadata_filename):
+    # Run ffmpeg's ffprobe
+    command = [
+        "ffprobe",
+        "-loglevel",  "quiet",
+        "-print_format", "json",
+        "-show_format",
+        "-show_streams",
+        audio_filename
+    ]
+
+    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    out, err = pipe.communicate()
+    ffmpeg_metadata = json.loads(out)
+
+    # Collect metadata
+    print(ffmpeg_metadata)
+    mime_type = magic.from_file(audio_filename, mime=True)
+    file_size_bytes = os.path.getsize(audio_filename)
+
+    duration_millis = int(float(ffmpeg_metadata['format']['duration']) * 1000)
+
+    metadata = {
+        'duration_millis': duration_millis,
+        'mimetype': mime_type,
+        'file_size_bytes': file_size_bytes,
+    }
+
+    with open(output_metadata_filename, 'w') as json_file:
+        json.dump(metadata, json_file)
+
 hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
 
 acoustic = load_acoustic_model(args.acoustic_model_filename)
@@ -125,4 +165,6 @@ hifigan = load_hifigan_model(args.hifigan_model_filename)
 source_audio = load_wav_data(args.source_audio_filename)
 target_audio = do_voice_conversion(hubert, acoustic, hifigan, source_audio)
 
-sf.write(args.output_audio_filename, target_audio.squeeze().cpu(), samplerate=16000)
+save_output_audio(target_audio, args.output_audio_filename)
+save_metadata(args.output_audio_filename, args.output_metadata_filename)
+
diff --git a/voice_conversion/softvc/requirements.txt b/voice_conversion/softvc/requirements.txt
index a74e976..7f47736 100644
--- a/voice_conversion/softvc/requirements.txt
+++ b/voice_conversion/softvc/requirements.txt
@@ -1,5 +1,6 @@
 asttokens==2.2.1
 backcall==0.2.0
+beautifulsoup4==4.11.1
 certifi==2022.12.7
 cffi==1.15.1
 charset-normalizer==2.1.1
@@ -7,7 +8,9 @@ contourpy==1.0.6
 cycler==0.11.0
 decorator==5.1.1
 executing==1.2.0
+filelock==3.8.2
 fonttools==4.38.0
+gdown==4.6.0
 idna==3.4
 ipython==8.7.0
 jedi==0.18.2
@@ -27,18 +30,22 @@ pure-eval==0.2.2
 pycparser==2.21
 Pygments==2.13.0
 pyparsing==3.0.9
+PySocks==1.7.1
 python-dateutil==2.8.2
+python-magic==0.4.27
 requests==2.28.1
 scikit-learn==1.1.3
 scipy==1.9.3
 six==1.16.0
 sklearn==0.0.post1
 soundfile==0.11.0
+soupsieve==2.3.2.post1
 stack-data==0.6.2
 threadpoolctl==3.1.0
 torch==1.13.0+cu116
 torchaudio==0.13.0+cu116
 torchvision==0.14.0+cu116
+tqdm==4.64.1
 traitlets==5.6.0
 typing_extensions==4.4.0
 urllib3==1.26.13