autochart.py

#!/usr/bin/env python3
import argparse
import os
import sys

# Ensure local modules can be found
sys.path.append(os.getcwd())

try:
    from infer.autochart_lib import AutoChart
except ImportError:
    sys.path.append(os.path.join(os.getcwd(), "infer"))
    try:
        from autochart_lib import AutoChart
    except ImportError:
        print("Error: Could not import AutoChart class.")
        sys.exit(1)


def get_version():
    try:
        with open(os.path.join(os.path.dirname(__file__), "VERSION"), encoding="utf-8") as f:
            return f.read().strip()
    except FileNotFoundError:
        return "0.0.0-dev"


def main():
    print(f"AutoChart v{get_version()} - Dance Dance Convolution")
    parser = argparse.ArgumentParser()
    parser.add_argument("--version", action="version", version=f"%(prog)s {get_version()}")
    parser.add_argument("input_paths", type=str, nargs="+", help="Input MP3/OGG/WAV files or directories")
    parser.add_argument("--out_dir", type=str, default="output", help="Output directory")
    parser.add_argument("--models_dir", type=str, required=True, help="Directory containing trained models")
    parser.add_argument("--ffr_dir", type=str, help="Directory containing FFR models")
    parser.add_argument("--google_key", type=str, help="Google API Key")
    parser.add_argument("--cx", type=str, help="Google Custom Search Engine ID")

    args = parser.parse_args()
    ac = AutoChart(args.models_dir, args.ffr_dir, args.google_key, args.cx)

    for input_path in args.input_paths:
        if os.path.isdir(input_path):
            print(f"Batch processing directory: {input_path}")
            files = []
            for root, _, filenames in os.walk(input_path):
                for filename in filenames:
                    if filename.lower().endswith((".mp3", ".ogg", ".wav")):
                        files.append(os.path.join(root, filename))

            print(f"Found {len(files)} audio files in {input_path}.")
            for file_path in files:
                try:
                    ac.process_song(file_path, args.out_dir)
                except Exception as e:
                    print(f"Failed to process {file_path}: {e}")
        else:
            try:
                ac.process_song(input_path, args.out_dir)
            except Exception as e:
                print(f"Failed to process {input_path}: {e}")


if __name__ == "__main__":
    main()