cleanups, add demucs option

transcribe.py

@@ -26,12 +26,12 @@ def prepare_annotations(result, join_sentences=False):
     if join_sentences:
         return prepare_joint_annotations(result)
     annotations = []
-    for segments in result["segments"]:
+    for segment in result["segments"]:
         annotations.append(
             {
                 "in": segment["start"],
-                "end": segment["end"] + 0.3,
-                "value": segment["text"],
+                "out": segment["end"] + 0.3,
+                "value": segment["text"].strip(),
             }
         )
     return annotations
@@ -92,20 +92,30 @@ def prepare_joint_annotations(result, target_length=200):
     return annotations
 
 
-def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=False, model="small"):
-    language = None
-    if "language" not in item.data:
-        language = None
-    else:
-        language = ox.iso.langTo2Code(item.data["language"][0])
-    if not item.streams():
-        logger.error("skip item without media %s: %s", item.public_id)
-        return False
-    src = item.streams()[0].media.path
+def run_demucs(src, output):
+    cmd = [
+        "/opt/whisper-timestamped/bin/demucs"
+        "--two-stems", "vocals",
+        "-o", output,
+        src
+    ]
+    subprocess.check_call(cmd)
+    wav = glob("%s/htdemucs/*/vocals.wav" % output)[0]
+    return wav
 
+
+def run_whisper(src, language=None, translate=False, gpu=False, model="small", demucs=False):
     tmp = tempfile.mkdtemp()
-    output = os.path.join(tmp, "output.json")
 
+    if demucs:
+        try:
+            src = run_demucs(src, tmp)
+        except:
+            logger.error("failed to run demucs for %s", src)
+            shutil.rmtree(tmp)
+            return None
+
+    output = os.path.join(tmp, "output.json")
     run_py = os.path.join(os.path.dirname(os.path.abspath(__file__)), "run_whisper.py")
     cmd = ["/opt/whisper-timestamped/bin/python", run_py]
     cmd += ["--model", model]
@@ -120,13 +130,30 @@ def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=Fa
     try:
         subprocess.check_call(cmd)
     except:
-        logger.error(
-            "failed to extract subtitles from item %s\n%s", item.public_id, cmd
-        )
-        return False
+        logger.error("failed to run: %s", cmd)
+        shutil.rmtree(tmp)
+        return None
+
     with open(output) as fd:
         response = json.load(fd)
-    shutil.rmtree(tmp)
+
+    #shutil.rmtree(tmp)
+    return response
+
+def extract_subtitles(item, user, layer, translate, gpu=False, join_sentences=False, model="small"):
+    language = None
+    if "language" not in item.data:
+        language = None
+    else:
+        language = ox.iso.langTo2Code(item.data["language"][0])
+    if not item.streams():
+        logger.error("skip item without media %s: %s", item.public_id)
+        return False
+    src = item.streams()[0].media.path
+    response = run_whisper(src, language, translate, gpu, model)
+    if not response:
+        logger.error("extract failed for  %s", item.public_id)
+        return False
     annotations = prepare_annotations(response, join_sentences=join_sentences)
     if not annotations:
         return False