songSplitter.py

import whisper
import json
import os
from pydub import AudioSegment
import argparse
import sys
import shutil


def transcribe_audio(model, audio_path):
    """
    Transcribes the audio file using the provided Whisper model
    and returns a list of words with timestamps.
    """
    print(f"Starting transcription for: {audio_path}")
    try:
        result = model.transcribe(audio_path, word_timestamps=True)
        word_list = []
        if 'segments' not in result:
            print("Warning: No segments found in transcription result.", file=sys.stderr)
            return []

        for segment in result['segments']:
            if 'words' not in segment:
                continue
            for word in segment['words']:
                # Ensure essential keys exist
                if 'word' in word and 'start' in word:
                    word_list.append({
                        "word": word['word'].strip(),
                        "timestamp": round(word['start'], 3)
                    })
                else:
                     print(f"Warning: Skipping word entry due to missing data: {word}", file=sys.stderr)

        print(f"Transcription complete. Found {len(word_list)} words.")
        return word_list
    except Exception as e:
        print(f"Error during transcription: {e}", file=sys.stderr)
        return []


def split_audio_by_words(word_list, audio_path, output_dir):
    """
    Splits the audio file into segments based on word timestamps
    and saves them to the output directory.
    """
    if not word_list:
        print("Word list is empty, skipping audio splitting.", file=sys.stderr)
        return

    print(f"Loading audio file: {audio_path}")
    try:
        audio = AudioSegment.from_mp3(audio_path)
    except Exception as e:
        print(f"Error loading audio file {audio_path}: {e}", file=sys.stderr)
        return

    print(f"Creating output directory: {output_dir}")
    if os.path.exists(output_dir):
        print(f"Cleaning existing output directory: {output_dir}")
        try:
            shutil.rmtree(output_dir)
        except Exception as e:
            print(f"Error removing directory {output_dir}: {e}", file=sys.stderr)
    
    os.makedirs(output_dir, exist_ok=True)

    print("Processing and exporting audio segments...")
    num_words = len(word_list)
    for i in range(num_words):
        current_word = word_list[i]
        start_ms = current_word['timestamp'] * 1000

        if i < num_words - 1:
            next_word = word_list[i+1]
            end_ms_default = next_word['timestamp'] * 1000
            # Set the end time to the minimum of the next word's start or start_ms + 2 seconds
            end_ms = min(end_ms_default, start_ms + 2000)
            # Ensure end_ms is not before start_ms (can happen with overlapping timestamps)
            if end_ms <= start_ms:
                 print(f"Warning: Adjusted end time for word '{current_word['word']}' ({i}) due to overlap. Using start_ms + 100ms.", file=sys.stderr)
                 end_ms = start_ms + 100 # Use a small default duration
        else:
            # Handle the last word: extend by 1 second or up to audio length
            end_ms = min(start_ms + 1000, len(audio))

        try:
            segment = audio[start_ms:end_ms]
            output_filename = os.path.join(output_dir, f"{i:03d}.mp3")
            segment.export(output_filename, format="mp3")
        except Exception as e:
             print(f"Error processing or exporting segment {i} ({current_word['word']}): {e}", file=sys.stderr)

    print(f"Audio splitting complete. {num_words} segments generated.")


def generate_splitsong_json(word_list, output_dir, output_json_path):
    """
    Generates a JSON file mapping lowercase words to their
    corresponding audio segment filenames.
    """
    if not word_list:
        print("Word list is empty, skipping JSON generation.", file=sys.stderr)
        return

    print(f"Generating JSON mapping file: {output_json_path}")
    output_data = []
    for i, word_data in enumerate(word_list):
        entry = {
            "word": word_data["word"].lower(),
            "sound": os.path.join(output_dir, f"{i:03d}.mp3")
        }
        output_data.append(entry)

    try:
        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)
        print("JSON mapping file generated successfully.")
    except Exception as e:
        print(f"Error writing JSON file {output_json_path}: {e}", file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description="Transcribe an audio file, split it by words, and generate a JSON map.")
    parser.add_argument("audio_file", help="Path to the input audio file (e.g., song.mp3)")
    parser.add_argument("-o", "--output_dir", default="output", help="Directory to save segmented audio files (default: output)")
    parser.add_argument("-j", "--json_path", default="splitsong.json", help="Path to save the output JSON mapping (default: splitsong.json)")
    parser.add_argument("-m", "--model", default="medium", help="Whisper model name (e.g., tiny, base, small, medium, large) (default: medium)")
    args = parser.parse_args()

    if not os.path.exists(args.audio_file):
        print(f"Error: Input audio file not found: {args.audio_file}", file=sys.stderr)
        sys.exit(1)

    # --- Load Model ---
    print(f"Loading Whisper model: {args.model}...")
    try:
        model = whisper.load_model(args.model)
    except Exception as e:
        print(f"Error loading Whisper model '{args.model}': {e}", file=sys.stderr)
        print("Please ensure the model name is correct and you have enough resources.", file=sys.stderr)
        sys.exit(1)

    # --- Step 1: Transcription ---
    word_list = transcribe_audio(model, args.audio_file)
    if not word_list:
        print("Transcription failed or resulted in an empty word list. Exiting.", file=sys.stderr)
        sys.exit(1) # Exit if transcription fails

    # --- Step 2: Audio Splitting ---
    split_audio_by_words(word_list, args.audio_file, args.output_dir)

    # --- Step 3: JSON Generation ---
    generate_splitsong_json(word_list, args.output_dir, args.json_path)

    print("\nProcessing complete.")
    print(f"Segmented audio saved to: {args.output_dir}")
    print(f"JSON mapping saved to: {args.json_path}")

if __name__ == "__main__":
    main()