add extract-audio and loop through dir for all commands

ffmpeg_handling.py

@@ -14,6 +14,7 @@ import json
 import ffmpeg
 import re
 from pathlib import Path
+VIDEO_EXTS = {".mp4", ".mkv", ".avi", ".mov", ".webm"}
 
 # ----------------------------
 # OUTPUT ENFORCEMENT (clean + centralized)
@@ -24,6 +25,7 @@ REQUIRES_OUTPUT = {
     "convert-video",
     "remove-subs",
     "remove-audio",
+    "extract-audio",
     "downsize",
     "filter-audio",
     "cut",
@@ -220,7 +222,7 @@ def cmd_convert_av(args):
         # Prevents FFmpeg from trying to guess/add extra default streams
         map_metadata=0
     )
-    return [out]
+    return out
 
 def cmd_convert_audio(args):
     """
@@ -242,7 +244,7 @@ def cmd_convert_audio(args):
         map='0',
         map_metadata=0
     )
-    return [out]
+    return out
 
 def cmd_convert_video(args):
     """
@@ -261,7 +263,7 @@ def cmd_convert_video(args):
         map='0',
         map_metadata=0
     )
-    return [out]
+    return out
 
 def cmd_remove_subs(args):
     """
@@ -282,12 +284,11 @@ def cmd_remove_subs(args):
         map='0',
         map_metadata=0
     )
-    return [out]
+    return out
 
 def cmd_remove_audio_track(args):
     """
     1. Removes a specific audio track from the video (Main Output).
-    2. Saves that specific audio track as a standalone MP3 (Extra Output).
     """
     inp = ffmpeg.input(args.input)
     probe = ffmpeg.probe(args.input)
@@ -295,35 +296,47 @@ def cmd_remove_audio_track(args):
     remove_idx = select_audio_stream(args.input, args.audio_index)
 
     streams_to_keep = []
-    extracted_audio = None
 
     for s in probe["streams"]:
         idx = s["index"]
-        if idx == remove_idx:
+        if idx != remove_idx:
-            extracted_audio = inp[str(idx)]
-        else:
             streams_to_keep.append(inp[str(idx)])
 
-    video_out = ffmpeg.output(
+    out = ffmpeg.output(
         *streams_to_keep,
         args.output,
         vcodec="copy",
         acodec="copy",
         scodec= get_scodec(args.output),
     )
+    return out
 
-    outputs = [video_out]
+def cmd_extract_audio_track(args):
+    """
+    1. Saves that specific audio track as a standalone MP3 (Extra Output).
+    """
+    inp = ffmpeg.input(args.input)
+    probe = ffmpeg.probe(args.input)
 
-    if args.output_audio:
+    output_audio = Path(args.output)
-        audio_out = ffmpeg.output(
+    output_audio = str(output_audio.with_name(f"{output_audio.stem}.mp3"))
-            extracted_audio,
-            args.output_audio,
-            acodec='libmp3lame',
-            audio_bitrate='192k'
-        )
-        outputs.append(audio_out)
 
-    return outputs
+    remove_idx = select_audio_stream(args.input, args.audio_index)
+
+    extracted_audio = None
+
+    for s in probe["streams"]:
+        idx = s["index"]
+        if idx == remove_idx:
+            extracted_audio = inp[str(idx)]
+
+    out = ffmpeg.output(
+        extracted_audio,
+        output_audio,
+        acodec='libmp3lame',
+        audio_bitrate='192k'
+    )
+    return out
 
 def cmd_downsize(args):
     """
@@ -356,15 +369,15 @@ def cmd_downsize(args):
         acodec="copy",
         scodec=get_scodec(args.output)
     )
-    return [out]
+    return out
 
 def cmd_filter_audio(args):
     """
     Multi-stage audio mastering chain.
 
     FFmpeg Equivalent:
-    ffmpeg -i [input] -filter_complex "[0:a:X]highpass...,acompressor...,equalizer...[out]" \
+    ffmpeg -i [input] -filter_complex "[0:a:X]highpass...,acompressor...,equalizer...out" \
-           -map 0:v -map [out] -c:v copy -c:a libvorbis [output]
+           -map 0:v -map out -c:v copy -c:a libvorbis [output]
 
     Filter Descriptions:
     1. highpass: Cuts frequencies below 100Hz (removes low-end rumble).
@@ -384,6 +397,7 @@ def cmd_filter_audio(args):
         for s in probe["streams"]
         if s.get("codec_type") == "audio"
     ]
+
     filtered_audios = []
     for idx in audio_streams:
         audio = inp[str(idx)]
@@ -400,15 +414,14 @@ def cmd_filter_audio(args):
             )
             filtered_audios.append(audio_filtered)
         else:
-
             filtered_audios.append(audio)
-    info = ffmpeg.probe(args.input)
+
-    streams = info.get("streams", [])
+    streams = probe.get("streams", [])
     subtitle_streams = [s for s in streams if s.get("codec_type") == "subtitle"]
 
-    subs = {}
+    subs = []
-    if len(subtitle_streams) > 0:
+    if subtitle_streams:
-        subs = {inp["s"]}
+        subs.append(inp["s"])
 
     out = ffmpeg.output(
         inp.video,
@@ -420,7 +433,7 @@ def cmd_filter_audio(args):
         scodec=get_scodec(args.output),
     )
 
-    return [out]
+    return out
 
 def cmd_cut_segment(args):
     """
@@ -452,7 +465,7 @@ def cmd_cut_segment(args):
         map="0"
     )
 
-    return [out]
+    return out
 
 def cmd_set_primary_audio(args):
     """
@@ -514,7 +527,7 @@ def cmd_set_primary_audio(args):
         **out_kwargs
     )
 
-    return [out]
+    return out
 
 def cmd_test(args):
     """
@@ -568,10 +581,6 @@ def cmd_test(args):
         if hasattr(args, "output") and args.output:
             out_files.append(args.output)
 
-        if hasattr(args, "output_audio") and args.output_audio:
-            out_files.append(args.output_audio)
-
-
         for i, f in enumerate(out_files):
             if not os.path.exists(f):
                 print(f"{red_tag("FAIL")} Missing output: {f}")
@@ -708,12 +717,20 @@ def cmd_test(args):
     for test_file in test_files:
         video_streams,audio_streams,subtitle_streams = get_streams(test_file)
         args.input = test_file
-        test_dict = [{"video_count": len(video_streams), "audio_count": 1, "subtitle_count": len(subtitle_streams)},
+        test_dict = [{"video_count": len(video_streams), "audio_count": 1, "subtitle_count": len(subtitle_streams)}]
-                     {"video_count": 0, "audio_count": 1, "subtitle_count": 0}]
         args.output = os.path.join(test_dir, "out_remove_audio" + os.path.splitext(test_file)[1])
-        args.output_audio = os.path.join(test_dir, "extracted_audio_" + os.path.splitext(test_file)[1][1:] + ".mp3")
         ret += run_and_probe("remove-audio-track", cmd_remove_audio_track, args, test_dict)
 
+    # ---------- EXTRACT AUDIO TRACK ----------
+    args = Args()
+    args.audio_index = 1
+
+    for test_file in test_files:
+        args.input = test_file
+        test_dict = [{"video_count": 0, "audio_count": 1, "subtitle_count": 0}]
+        args.output = os.path.join(test_dir, "extracted_audio_" + os.path.splitext(test_file)[1][1:] + ".mp3")
+        ret += run_and_probe("remove-audio-track", cmd_extract_audio_track, args, test_dict)
+
     # ---------- CONVERT AV ----------
     args = Args()
     for test_file in test_files:
@@ -864,10 +881,16 @@ def main():
     # REMOVE AUDIO TRACK
     # ----------------------------
     p_ra = sub.add_parser("remove-audio", aliases=["ra"], parents=[base_p], help="Remove one audio track")
-    p_ra.add_argument("-oa", "--output-audio")
     p_ra.add_argument("-a", "--audio-index", type=int)
     p_ra.set_defaults(func=cmd_remove_audio_track, command="remove-audio")
 
+    # ----------------------------
+    # REMOVE AUDIO TRACK
+    # ----------------------------
+    p_ea = sub.add_parser("extract-audio", aliases=["ea"], parents=[base_p], help="Extract one audio track")
+    p_ea.add_argument("-a", "--audio-index", type=int)
+    p_ea.set_defaults(func=cmd_extract_audio_track, command="extract-audio")
+
     # ----------------------------
     # DOWNSIZE
     # ----------------------------
@@ -931,11 +954,13 @@ def main():
     # ----------------------------
     if args.command != "test":
         args.input = ensure(getattr(args, "input", None), f"Input file ({Path.cwd()})")
-        if not Path(args.input).is_file():
+        if not Path(args.input).expanduser().exists():
             parser.error("Input file does not exist!")
-
     if args.command in REQUIRES_OUTPUT and not getattr(args, "output", None):
-        args.output = ensure(getattr(args, "output", None), "Output file")
+        args.output = ensure(getattr(args, "output", None), "Output file (enter for default file)")
+        if args.output == "" or args.output is None:
+            input_path = Path(args.input)
+            args.output = str(input_path.with_name(f"{input_path.stem}_{args.command}{input_path.suffix}"))
     # ----------------------------
     # COMMAND-SPECIFIC REQUIRED FIELDS
     # ----------------------------
@@ -960,17 +985,40 @@ def main():
         args.height = ensure(getattr(args, "height", None), "height: eg 720")
         args.crf = ensure(getattr(args, "crf", None), "crf: eg 23")
 
-    if args.command in ["filter-audio", "remove-audio", "set-primary-audio"]:
+    if args.command in ["filter-audio", "remove-audio", "extract-audio", "set-primary-audio"]:
         args.audio_index = ensure(getattr(args, "audio_index", None), "Audio index")
+        if args.audio_index is None:
+            args.audio_index = 1
+
 
     # ----------------------------
     # EXECUTE COMMAND
     # ----------------------------
     try:
-        outputs = args.func(args)
+        if args.command == "test":
+            args.func(args)
+            sys.exit(0)
+        outputs = []
+        input_path = Path(args.input)
+        # ---- DIRECTORY ----
+        if input_path.is_dir():
+            for file in input_path.iterdir():
+                if file.suffix.lower() not in VIDEO_EXTS:
+                    continue
+                print(f"Processing File: {file}")
+                out_path = file.with_name(f"{file.stem}_{args.command}{file.suffix}")
+
+                args.input = str(file)
+                args.output = str(out_path)
+                outputs.append(args.func(args))
+        # ---- SINGLE FILE ----
+        elif input_path.is_file():
+            outputs.append(args.func(args))
+        else:
+            raise ValueError(f"Invalid input path: {args.input}")
 
         if outputs is None:
-            return
+            raise ValueError(f"No Output: {outputs}")
 
         if not isinstance(outputs, list):
             outputs = [outputs]