process

1 files changed, 79 insertions, 0 deletions

diff --git a/scripts/process.py b/scripts/process.py
new file mode 100644
index 0000000..646c19b
--- /dev/null
+++ b/scripts/process.py
@@ -0,0 +1,79 @@
+import argparse
+import glob
+import multiprocessing
+import re
+from functools import partial
+from pathlib import Path
+
+import librosa
+import numpy
+
+from become_yukarin import AcousticConverter
+from become_yukarin import SuperResolution
+from become_yukarin.config.config import create_from_json as create_config
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-cmd', '--conv_model_directory', type=Path, default=Path('/mnt/dwango/hiroshiba/become-yukarin/'))
+parser.add_argument('-srmd', '--sr_model_directory', type=Path, default=Path('/mnt/dwango/hiroshiba/become-yukarin/'))
+parser.add_argument('-iwd', '--input_directory', type=Path, default=Path('./input'))
+parser.add_argument('-owd', '--output_directory', type=Path, default=Path('./output'))
+parser.add_argument('-it', '--iteration', type=int)
+parser.add_argument('-g', '--gpu', type=int)
+args = parser.parse_args()
+
+conv_model = args.conv_model_directory  # type: Path
+sr_model = args.sr_model_directory  # type: Path
+input_directory = args.input_directory
+output_directory = args.output_directory
+it = args.iteration
+gpu = args.gpu
+
+def extract_number(f):
+    s = re.findall("\d+", str(f))
+    return int(s[-1]) if s else -1
+
+def load_acoustic_converter():
+  conv_config = create_config(conv_model / 'config.json')
+  if it is not None:
+      conv_model_path = conv_model / 'predictor_{}.npz'.format(it)
+  else:
+      conv_model_paths = conv_model.glob('predictor_*.npz')
+      conv_model_path = list(sorted(conv_model_paths, key=extract_number))[-1]
+  print(conv_model_path)
+  return AcousticConverter(conv_config, conv_model_path, gpu=gpu)
+
+def load_super_resolution():
+  sr_config = create_config(sr_model / 'config.json')
+  sr_model_paths = base_model.glob('predictor*.npz')
+  sr_model_path = list(sorted(sr_model_paths, key=extract_number))[-1]
+  print(sr_model_path)
+  return SuperResolution(sr_config, sr_model_path, gpu=gpu)
+
+def process(p: Path, acoustic_converter: AcousticConverter, super_resolution: SuperResolution, output: Path):
+  try:
+    print(str(p))
+    input = acoustic_converter(p)
+    wave = super_resolution(input.spectrogram, acoustic_feature=input, sampling_rate=param.voice_param.sample_rate)
+    librosa.output.write_wav(str(output / p.stem) + '.wav', wave.wave, wave.sampling_rate, norm=True)
+  except:
+    import traceback
+    print('error!', str(p))
+    print(traceback.format_exc())
+
+def run():
+  input_paths = list(input_directory.glob('*.wav'))
+
+  acoustic_converter = load_acoustic_converter()
+  super_resolution = load_super_resolution()
+
+  output = output_directory.absolute()
+  output.mkdir(exist_ok=True)
+
+  process_partial = partial(process, acoustic_converter=acoustic_converter, super_resolution=super_resolution, output=output)
+  if gpu is None:
+      pool = multiprocessing.Pool()
+      pool.map(process_partial, input_paths)
+  else:
+      list(map(process_partial, input_paths))
+
+run()