syncing latest revisions - dk123

video_transcription/ai_transcriber_v2/main.py

@@ -27,6 +27,8 @@ from diarizer import diarize_audio, merge_diarization_with_transcript
 import tracker
 from tracker import JobStatus
 
+from tqdm import tqdm
+
 def save_srt_with_speakers(segments, output_path):
     """Helper to save SRT with speaker labels prepended to text."""
     def format_timestamp(seconds: float):
@@ -50,7 +52,7 @@ def save_srt_with_speakers(segments, output_path):
             f.write(f"{i}\n")
             f.write(f"{start} --> {end}\n")
             f.write(f"{text}\n\n")
-    print(f"SRT saved to: {output_path}")
+    tracker.logger.info(f"SRT saved to: {output_path}")
 
 def process_file(file_path, args, source_lang=None, loaded_model=None, service_status=None):
     tracker.logger.info(f"=== Processing: {file_path} ===")
@@ -324,34 +326,35 @@ def main():
     loaded_model = load_whisper_model(args.model)
     # -----------------------
 
-    # Initialize Graceful Exit Handler
-    killer = GracefulKiller()
-
+    # --- Collect All Files ---
+    all_files = []
     video_extensions = ('.mp4', '.mkv', '.mov', '.avi', '.webm', '.flv', '.wmv', '.m4v')
     
     for input_path in valid_inputs:
+        if os.path.isfile(input_path):
+            all_files.append(input_path)
+        elif os.path.isdir(input_path):
+            for root, dirs, files in os.walk(input_path):
+                for file in files:
+                    if file.lower().endswith(video_extensions):
+                        all_files.append(os.path.join(root, file))
+    
+    if not all_files:
+        print("No video files found to process.")
+        return
+
+    # --- Batch Process with Progress Bar ---
+    pbar = tqdm(all_files, desc="Batch Progress", unit="file", dynamic_ncols=True, leave=True)
+    
+    for file_path in pbar:
         if killer.kill_now:
             break
-
-        if os.path.isfile(input_path):
-            process_file(input_path, args, source_lang, loaded_model=loaded_model, service_status=service_status)
-        elif os.path.isdir(input_path):
-            found = False
-            for root, dirs, files in os.walk(input_path):
-                if killer.kill_now:
-                    break
-
-                for file in files:
-                    if killer.kill_now:
-                        break
-
-                    if file.lower().endswith(video_extensions):
-                        found = True
-                        process_file(os.path.join(root, file), args, source_lang, loaded_model=loaded_model, service_status=service_status)
-            if not found:
-                print(f"No video files found in {input_path}")
-        else:
-            print(f"Error: Invalid input path '{input_path}'")
+        
+        # Update progress bar description with current file
+        filename = os.path.basename(file_path)
+        pbar.set_description(f"File: {filename[:30]}")
+        
+        process_file(file_path, args, source_lang, loaded_model=loaded_model, service_status=service_status)
     
     if killer.kill_now:
         print("\n🛑 Process stopped by user. Progress saved in database.")