AI-Song-Rater/wavToStars.py at main · PeterM-13/AI-Song-Rater · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import numpy as np
import librosa.display, os
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet import preprocess_input
import shutil

# This script creates spectrograms from WAV files and classifies them into star ratings using a pre-trained model.

def create_spectrogram(audio_file, image_file):
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    fig.subplots_adjust(left=0, right=1, bottom=0, top=1)

    y, sr = librosa.load(audio_file)
    ms = librosa.feature.melspectrogram(y=y, sr=sr)
    log_ms = librosa.power_to_db(ms, ref=np.max)
    librosa.display.specshow(log_ms, sr=sr)

    fig.savefig(image_file)
    plt.close(fig)

def create_pngs_from_wavs(input_path, output_path):
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    dir = os.listdir(input_path)

    for i, file in enumerate(dir):
        input_file = os.path.join(input_path, file)
        output_file = os.path.join(output_path, file.replace('.wav', '.png'))
        create_spectrogram(input_file, output_file)


class SongClassifier:
    def __init__(
        self,
        downloads_dir='./Downloads',
        one_star_dir='./1_Star',
        two_star_dir='./2_Stars',
        three_star_dir='./3_Stars',
        four_star_dir='./4_Stars',
        five_star_dir='./5_Stars',
        spectrograms_tmp_dir='./0_tmp_spectrograms',
        model_path=None
    ):
        self.downloads_dir = downloads_dir
        self.three_star_dir = three_star_dir
        self.four_star_dir = four_star_dir
        self.five_star_dir = five_star_dir
        self.two_star_dir = two_star_dir
        self.one_star_dir = one_star_dir
        self.spectrograms_tmp_dir = spectrograms_tmp_dir
        self.class_labels = ['1 Star', '2 Stars', '3 Stars', '4 Stars', '5 Stars']

        for d in [self.three_star_dir, self.four_star_dir, self.five_star_dir, self.two_star_dir, self.one_star_dir, self.spectrograms_tmp_dir]:
            if not os.path.exists(d):
                os.makedirs(d)

        self.base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
        if model_path is None:
            raise ValueError("model_path must be provided")
        self.trained_model = load_model(model_path)

    def classify_and_move(self):
        for file in os.listdir(self.downloads_dir):
            if not file.lower().endswith('.wav'):
                continue
            audio_path = os.path.join(self.downloads_dir, file)
            spectro_path = os.path.join(self.spectrograms_tmp_dir, file.replace('.wav', '.png'))
            create_spectrogram(audio_path, spectro_path)
            x = image.load_img(spectro_path, target_size=(224, 224))
            x = image.img_to_array(x)
            x = np.expand_dims(x, axis=0)
            x = preprocess_input(x)
            y = self.base_model.predict(x)
            predictions = self.trained_model.predict(y)
            top_label = np.argmax(predictions[0])
            confidence = float(np.max(predictions[0]))
            confidence_str = f"{confidence:.2f}"
            # Create a new file name with confidence
            new_file_name = f"{confidence_str} {file}"
            # Move the file to the appropriate directory based on the classification
            if top_label == 3:  # 4 stars (0-based index)
                shutil.move(audio_path, os.path.join(self.four_star_dir, new_file_name))
            elif top_label == 4:  # 5 stars (0-based index)
                shutil.move(audio_path, os.path.join(self.five_star_dir, new_file_name))
            elif top_label == 2:  # 3 stars (0-based index)
                shutil.move(audio_path, os.path.join(self.three_star_dir, new_file_name))
            elif top_label == 1:  # 2 stars (0-based index)
                shutil.move(audio_path, os.path.join(self.two_star_dir, new_file_name))
            elif top_label == 0:  # 1 star (0-based index)
                shutil.move(audio_path, os.path.join(self.one_star_dir, new_file_name))

# Example usage:
# classifier = SongClassifier()
# classifier.classify_and_move()