import os
from shutil import copyfile
from tqdm import tqdm  # Import tqdm for progress bars

def organize_images_by_class(input_folder, output_folder):
    # Create output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Get a list of all spectrogram images
    spectrogram_files = [f for f in os.listdir(input_folder) if f.endswith('.png')]

    # Organize files by class names
    print("Organizing files by class names...")
    for file in tqdm(spectrogram_files, desc="Copying Files"):
        # Extract class name from the filename
        class_name = " ".join(file.split(" ")[:-1])  # Assuming class name is everything except the last part
        class_folder = os.path.join(output_folder, class_name)

        # Create class folder if it doesn't exist
        if not os.path.exists(class_folder):
            os.makedirs(class_folder)

        # Copy the file to the respective class folder
        copyfile(os.path.join(input_folder, file), os.path.join(class_folder, file))

if __name__ == "__main__":
    input_folder = 'spectrograms'  # Input folder containing spectrogram images
    output_folder = 'organized_spectrograms'  # Output folder for organized images

    organize_images_by_class(input_folder, output_folder)