Merge pull request #127 from nulib/testaprocnewspectrogram

Updated spectrogram tool

Author: SoFrans

nulrdcscripts/tools/spectrogramgeneration/main.py

@@ -1,100 +1,92 @@
-import progressbar
 import os
 import subprocess
+import progressbar
 from nulrdcscripts.tools.spectrogramgeneration.params import args
 
 
 def generate_spectrogram(input_path, channels, output_path, spectroname, ffmpegpath):
     """Creates a spectrogram for each audio track in the input"""
-    for index, item in enumerate(channels):
-        spectrogram_resolution = "1928x1080"
+    spectrogram_resolution = "1920x1080"  # Updated resolution
+    bar = progressbar.ProgressBar(max_value=int(channels))
+    bar.start()
+    for index in range(int(channels)):
         output = os.path.join(
-            output_path, spectroname + "_spectrogram0" + str(index + 1) + "_s.png"
+            output_path, f"{spectroname}_spectrogram0{index + 1}_s.png"
         )
         spectrogram_args = [ffmpegpath]
         spectrogram_args += ["-loglevel", "error", "-y"]
         spectrogram_args += ["-i", input_path, "-lavfi"]
-        item = int(item)
-        if item > 1:
+        if int(channels) > 1:
             spectrogram_args += [
-                "[0:a:%(a)s]showspectrumpic=mode=separate:s=%(b)s"
-                % {"a": index, "b": spectrogram_resolution}
+                f"[0:a:{index}]showspectrumpic=mode=separate:s={spectrogram_resolution}"
             ]
         else:
             spectrogram_args += [
-                "[0:a:%(a)s]showspectrumpic=s=%(b)s"
-                % {"a": index, "b": spectrogram_resolution}
+                f"[0:a:{index}]showspectrumpic=s={spectrogram_resolution}"
             ]
         spectrogram_args += [output]
         subprocess.run(spectrogram_args)
-        print("*** Spectrogram Generated ***")
+        bar.update(index + 1)
+        print(f"*** Spectrogram for channel {index + 1} generated ***")
+    bar.finish()
 
 
 def softwarecheck(software, softwarename):
     """Checks that software exists"""
-    if "sox" in software:
-        version = "--version"
-    else:
-        version = "-version"
-
+    version = "--version" if "sox" in software else "-version"
     try:
         subprocess.check_output([software, version]).decode(
             "ascii"
         ).rstrip().splitlines()[0].split()[2]
     except:
-        raise Exception("Cannot locate " + softwarename + ". Check path.")
+        raise Exception(f"Cannot locate {softwarename}. Check path.")
 
 
 def getnumberchannels(ffprobe_path, input_path):
-    command = (
-        ffprobe_path
-        + " "
-        + "-i"
-        + " "
-        + input_path
-        + " "
-        + "-show_entries stream=channels -select_streams a:0 -of compact=p=0:nk=1 -v 0"
-    )
+    command = [
+        ffprobe_path,
+        "-i",
+        input_path,
+        "-show_entries",
+        "stream=channels",
+        "-select_streams",
+        "a:0",
+        "-of",
+        "compact=p=0:nk=1",
+        "-v",
+        "0",
+    ]
     output = subprocess.run(command, capture_output=True, text=True)
-    channels = output.stdout
-    channels = channels.replace("\n", "")
+    channels = output.stdout.strip()
     return channels
 
 
 def checkOrCreateOutput(output_path, input_path):
     if output_path == "NA":
         output_path = input_path
-    else:
-        pass
-
-    outputdir = os.path.isdir(output_path)
-    if outputdir:
-        pass
-    else:
+    if not os.path.isdir(output_path):
         output_path = os.path.dirname(output_path)
     return output_path
 
 
-def callableSpectrogram(ffprobe_path, input_path, output_path):
+def callableSpectrogram(ffprobe_path, input_path, output_path, ffmpeg_path):
     channels = getnumberchannels(ffprobe_path, input_path)
-    spectroname, ext = os.path.splitext(os.path.basename(input_path))
-    generate_spectrogram(input_path, channels, output_path, spectroname)
+    spectroname, _ = os.path.splitext(os.path.basename(input_path))
+    generate_spectrogram(input_path, channels, output_path, spectroname, ffmpeg_path)
 
 
 def main():
-    """Only runs if you are running this command on it's own"""
-    input_path = os.path.abspath(args.input_path)
+    """Only runs if you are running this command on its own"""
+    input_path = args.input_path
     ffprobe_path = args.ffprobe_path
     ffmpeg_path = args.ffmpeg_path
     sox_path = args.sox_path
     softwarecheck(sox_path, "sox")
     softwarecheck(ffprobe_path, "ffprobe")
     softwarecheck(ffmpeg_path, "ffmpeg")
-    channels = getnumberchannels(ffprobe_path, input_path)
     output_path = args.output_path
     output_path = checkOrCreateOutput(output_path, input_path)
-    spectroname, ext = os.path.splitext(os.path.basename(input_path))
-    generate_spectrogram(input_path, channels, output_path, spectroname, ffmpeg_path)
+    callableSpectrogram(ffprobe_path, input_path, output_path, ffmpeg_path)
 
 
 if __name__ == "__main__":

nulrdcscripts/tools/spectrogramgeneration/test.py

@@ -0,0 +1,85 @@
+import progressbar
+import os
+import subprocess
+from nulrdcscripts.tools.spectrogramgeneration.params import args
+
+
+def generate_spectrogram(input_path, channels, output_path, spectroname, ffmpegpath):
+    """Creates a spectrogram for each audio track in the input"""
+    bar = progressbar.ProgressBar(max_value=int(channels))
+    for index in range(int(channels)):
+        spectrogram_resolution = "1928x1080"
+        output = os.path.join(
+            output_path, f"{spectroname}_spectrogram0{index + 1}_s.png"
+        )
+        spectrogram_args = [ffmpegpath]
+        spectrogram_args += ["-loglevel", "error", "-y"]
+        spectrogram_args += ["-i", input_path, "-lavfi"]
+        spectrogram_args += [f"[0:a:{index}]showspectrumpic=s={spectrogram_resolution}"]
+        spectrogram_args += [output]
+        subprocess.run(spectrogram_args)
+        bar.update(index + 1)
+        print(f"*** Spectrogram for channel {index + 1} generated ***")
+    bar.finish()
+
+
+def softwarecheck(software, softwarename):
+    """Checks that software exists"""
+    version = "--version" if "sox" in software else "-version"
+    try:
+        subprocess.check_output([software, version]).decode(
+            "ascii"
+        ).rstrip().splitlines()[0].split()[2]
+    except:
+        raise Exception(f"Cannot locate {softwarename}. Check path.")
+
+
+def getnumberchannels(ffprobe_path, input_path):
+    command = [
+        ffprobe_path,
+        "-i",
+        input_path,
+        "-show_entries",
+        "stream=channels",
+        "-select_streams",
+        "a:0",
+        "-of",
+        "compact=p=0:nk=1",
+        "-v",
+        "0",
+    ]
+    output = subprocess.run(command, capture_output=True, text=True)
+    channels = output.stdout.strip()
+    return channels
+
+
+def checkOrCreateOutput(output_path, input_path):
+    if output_path == "NA":
+        output_path = input_path
+    if not os.path.isdir(output_path):
+        output_path = os.path.dirname(output_path)
+    return output_path
+
+
+def callableSpectrogram(ffprobe_path, input_path, output_path, ffmpeg_path):
+    channels = getnumberchannels(ffprobe_path, input_path)
+    spectroname, _ = os.path.splitext(os.path.basename(input_path))
+    generate_spectrogram(input_path, channels, output_path, spectroname, ffmpeg_path)
+
+
+def main():
+    """Only runs if you are running this command on its own"""
+    input_path = os.path.abspath(args.input_path)
+    ffprobe_path = args.ffprobe_path
+    ffmpeg_path = args.ffmpeg_path
+    sox_path = args.sox_path
+    softwarecheck(sox_path, "sox")
+    softwarecheck(ffprobe_path, "ffprobe")
+    softwarecheck(ffmpeg_path, "ffmpeg")
+    output_path = args.output_path
+    output_path = checkOrCreateOutput(output_path, input_path)
+    callableSpectrogram(ffprobe_path, input_path, output_path, ffmpeg_path)
+
+
+if __name__ == "__main__":
+    main()

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nul-rdc-scripts"
-version = "0.4.1"
+version = "0.5.0"
 description = "Scripts for NUL RDC Digitization Team"
 authors = [
     "Northwestern University Libraries <[email protected]>",
@@ -14,6 +14,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.10"
+progressbar2 = "^4.5.0"
 
 [tool.poetry.dev-dependencies]