scanner.py

import pandas as pd
import argparse
from plex_actions import PlexActions
from tqdm import tqdm

def main():
    parser = argparse.ArgumentParser(description="Create a metadata CSV from the ChromaDB database.")
    parser.add_argument("--output_csv", default="music_metadata.csv", help="The name of the output CSV file.")
    parser.add_argument("--exclude_styles", nargs='+', default=["alternative/indie rock", "alternative pop/rock", "indie rock"], help="List of styles to exclude from the dataset.")
    args = parser.parse_args()

    print("Connecting to Plex server...")
    try:
        plex = PlexActions()
        print("Fetching all albums from Plex music library...")
        albums = plex.get_all_albums()
        print(f"Found {len(albums)} albums in the library.")
    except Exception as e:
        print(f"Failed to connect to Plex: {e}")
        return

    print("Scanning tracks for style/style tags...")
    file_data = []
    for album in tqdm(albums):
        try:
            if album.styles:
                styles = [style.tag for style in album.styles]
            elif album.artist() and album.artist().styles:
                styles = [style.tag for style in album.artist().styles]
            else:
                styles = []

            if len(styles) > 0:
                tracks = plex.get_tracks_from_album(album)
                for track in tracks:
                    filepath = track['filepath']
                    
                    for style in styles:
                        file_data.append({'file_path': filepath, 'style': style.strip()})
        except Exception as e:
            print(f"  Warning: Error processing album '{album.title}': {e}")
            continue
    
    print(f"Finished scanning. Found {len(file_data)} track-style pairs.")

    if len(file_data) == 0:
        print("No files with style tags found in the database.")
        return

    df = pd.DataFrame(file_data)
    
    # --- Data Cleaning ---
    # 1. Consolidate styles (e.g., 'Alt-Rock' and 'Alternative Rock')
    df['style'] = df['style'].str.lower().str.strip()
    
    # Example consolidation:
    # df['style'] = df['style'].replace({'alt-rock': 'alternative', 'indie rock': 'indie'})
    
    # 1.5. Exclude specified styles
    if args.exclude_styles:
        excluded_styles = [g.lower().strip() for g in args.exclude_styles]
        initial_count = len(df)
        df = df[~df['style'].isin(excluded_styles)]
        excluded_count = initial_count - len(df)
        print(f"\nExcluded {excluded_count} tracks from styles: {', '.join(excluded_styles)}")
    
    # 2. Remove styles with too few samples (e.g., < 100 songs)
    style_counts = df['style'].value_counts()
    if style_counts.empty:
        print("No styles found after initial processing.")
        return
        
    valid_styles = style_counts[style_counts >= 100].index
    df = df[df['style'].isin(valid_styles)]

    if df.empty:
        print("No tracks remaining after filtering for styles with at least 100 samples.")
        print("\nGenre counts before filtering:")
        print(style_counts)
        return

    df.to_csv(args.output_csv, index=False)
    print(f"\nScan complete. Found {len(df)} tagged files across {len(valid_styles)} styles.")
    print(f"Saved metadata to {args.output_csv}.")
    print("\nTop 10 styles found:")
    print(style_counts.head(10))

if __name__ == "__main__":
    main()