CollentData.py

import librosa
import matplotlib.pyplot as plt
import librosa.display
import yt_dlp
import pandas as pd
import numpy as np
import os

class CollectData:
    def __init__(self, filename, evaluation, url):
        self.filename = filename
        self.url = url
        self.evaluation = evaluation
        self.melspec_db = None
        self.y = None
        self.sr = None
    
    def url2mp3(self):
        # ディレクトリの存在を確認し、存在しない場合は作成
        if not os.path.exists('music/'):
            os.makedirs('music/')

        ydl_opts = {
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
            'outtmpl': f'music/{self.filename}.%(ext)s',
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([self.url])
    
    def load_audio(self):
        self.y, self.sr = librosa.load(f'music/{self.filename}.mp3', duration=70)

    def ViewWave(self):
        if self.y is None or self.sr is None:
            self.load_audio()
        plt.plot(self.y)
        plt.show()
        
    def ViewMelspectrogram(self):
        if self.y is None or self.sr is None:
            self.load_audio()
        self.melspec = librosa.feature.melspectrogram(y=self.y, sr=self.sr)
        self.melspec_db = librosa.amplitude_to_db(self.melspec)
        librosa.display.specshow(self.melspec_db, sr=self.sr)
        plt.show()
    
    def SpectralCentroid(self):
        if self.y is None or self.sr is None:
            self.load_audio()
        spectral_centroid = librosa.feature.spectral_centroid(y=self.y, sr=self.sr)[0]
        spectral_centroid = np.convolve(spectral_centroid, np.ones(200)/200, mode='valid')
        times = librosa.frames_to_time(np.arange(len(spectral_centroid)), sr=self.sr)
        plt.figure(figsize=(12, 4))
        plt.plot(times, spectral_centroid)
        plt.ylabel('Hz')
        plt.xlabel('Time')
        plt.title('Spectral Centroid')
        plt.grid(True)
        plt.legend()
        plt.tight_layout()
        plt.show()
    
    def extract_features(self):
        if self.melspec_db is None:
            raise ValueError("データが未生成")

        avg_db = np.mean(self.melspec_db)

        features_df = pd.DataFrame([{
            "filename": self.filename,
            "evaluation": self.evaluation,
            "average_db": avg_db,
        }])
        return features_df

def save_features_to_csv(features_df, csv_filename='data.csv'):
    if not os.path.exists(csv_filename):
        features_df.to_csv(csv_filename, index=False)
    else:
        features_df.to_csv(csv_filename, mode='a', index=False, header=False)


if __name__ == "__main__":
    label = input("ファイル名: ")
    evaluation = input("起きやすさの評価をしてください(1〜5です): ")
    url = input('YouTube URL: ')

    collectData = CollectData(label, evaluation, url)
    collectData.url2mp3()
    collectData.SpectralCentroid()
    collectData.ViewMelspectrogram()

    features = collectData.extract_features()
    save_features_to_csv(features)