From 61bd5dd88dd262a1fedc46a2e47668bef9eefa8b Mon Sep 17 00:00:00 2001
From: Tilman Kamp <5991088+tilmankamp@users.noreply.github.com>
Date: Thu, 9 Jul 2020 11:55:26 +0200
Subject: [PATCH] Resolves #3144 - Add augmentation support to build_sdb.py

---
 bin/build_sdb.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/bin/build_sdb.py b/bin/build_sdb.py
index ef4810e7..ac7be8af 100755
--- a/bin/build_sdb.py
+++ b/bin/build_sdb.py
@@ -8,6 +8,7 @@ import argparse
 import progressbar
 
 from deepspeech_training.util.audio import (
+    AUDIO_TYPE_PCM,
     AUDIO_TYPE_OPUS,
     AUDIO_TYPE_WAV,
     change_audio_types,
@@ -17,17 +18,28 @@ from deepspeech_training.util.sample_collections import (
     DirectSDBWriter,
     samples_from_sources,
 )
+from deepspeech_training.util.augmentations import (
+    parse_augmentations,
+    apply_sample_augmentations,
+    SampleAugmentation
+)
 
 AUDIO_TYPE_LOOKUP = {"wav": AUDIO_TYPE_WAV, "opus": AUDIO_TYPE_OPUS}
 
 
 def build_sdb():
     audio_type = AUDIO_TYPE_LOOKUP[CLI_ARGS.audio_type]
+    augmentations = parse_augmentations(CLI_ARGS.augment)
+    if any(not isinstance(a, SampleAugmentation) for a in augmentations):
+        print("Warning: Some of the augmentations cannot be applied by this command.")
     with DirectSDBWriter(
         CLI_ARGS.target, audio_type=audio_type, labeled=not CLI_ARGS.unlabeled
     ) as sdb_writer:
         samples = samples_from_sources(CLI_ARGS.sources, labeled=not CLI_ARGS.unlabeled)
-        bar = progressbar.ProgressBar(max_value=len(samples), widgets=SIMPLE_BAR)
+        num_samples = len(samples)
+        if augmentations:
+            samples = apply_sample_augmentations(samples, audio_type=AUDIO_TYPE_PCM, augmentations=augmentations)
+        bar = progressbar.ProgressBar(max_value=num_samples, widgets=SIMPLE_BAR)
         for sample in bar(
             change_audio_types(samples, audio_type=audio_type, bitrate=CLI_ARGS.bitrate, processes=CLI_ARGS.workers)
         ):
@@ -67,6 +79,11 @@ def handle_args():
         help="If to build an SDB with unlabeled (audio only) samples - "
         "typically used for building noise augmentation corpora",
     )
+    parser.add_argument(
+        "--augment",
+        action='append',
+        help="Add an augmentation operation",
+    )
     return parser.parse_args()