Upload 11 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+audios_samples/classical.00000.wav filter=lfs diff=lfs merge=lfs -text
+audios_samples/country.00031.wav filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,137 @@
+import streamlit as st
+import numpy as np
+import torch
+import os
+
+from utils import *
+from model.model import CVAE
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+AUDIO_SAMPLES_DIR = "audios_samples"
+
+st.set_page_config(
+    page_title="Audio Reconstruction",
+    page_icon="./static/aivn_favicon.png",
+)
+
+st.image("./static/aivn_logo.png", width=300)
+
+st.title('New Genres Audio Reconstruction')
+
+@st.cache_data 
+def load_models():
+    st.spinner('Đang tải mô hình...')
+    # lưu mô hình để tránh tải lại
+    model = CVAE(64, 128, 256, 130, len(uni_genres_list)).to(device)
+    model.load_state_dict(torch.load('model/model_256.pth', map_location=torch.device('cpu')))
+    model.eval()
+    return model
+
+
+def gen_audio(model, audio_source, genres_list, fixed_length_seconds=3):
+    with st.spinner('Đang xử lý âm thanh...'):
+        audio_data, sr = load_and_resample_audio(audio_source)
+        n_frames = len(audio_data)
+        segment_length_frame = int(fixed_length_seconds * sr)
+        n_segments = n_frames // segment_length_frame
+        
+        split_audio_text_placeholder = st.empty()
+        split_audio_text_placeholder.text("Đang chia nhỏ audio... ✂")
+        progress_bar_placeholder = st.empty()
+        progress_bar = progress_bar_placeholder.progress(0)
+        
+        audios = []
+        for i in range(n_segments):
+            start = i * segment_length_frame
+            end = (i + 1) * segment_length_frame
+            segment = audio_data[start:end]
+            mel_spec = audio_to_melspec(segment, sr, to_db=True)
+            mel_spec_norm = normalize_melspec(mel_spec)
+            mel_spec = torch.tensor(mel_spec, dtype=torch.float32)
+            mel_spec_norm = torch.tensor(mel_spec_norm, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+            audios.append((mel_spec_norm, mel_spec))
+            progress_bar.progress(int((i + 1) / n_segments * 100))
+        
+        progress_bar_placeholder.empty()
+        split_audio_text_placeholder.empty()
+        
+        audios_input = torch.cat([audio[0] for audio in audios], dim=0)
+        
+        genres_input = onehot_encode(tokenize(genres_list), len(uni_genres_list))
+        genres_input = torch.tensor(genres_input, dtype=torch.long).unsqueeze(0).unsqueeze(0)
+        genres_input = genres_input.repeat(audios_input.shape[0], 1, 1)
+        
+        with st.spinner('Mô hình đang nấu ăn... 🍳🍴'):
+            recons, _, _ = model(audios_input, genres_input)
+        
+        recon_audio_text_placeholder = st.empty()
+        recon_audio_text_placeholder.text("Đang dựng lại audio video... 🎵")
+        progress_bar_placeholder = st.empty()
+        progress_bar = progress_bar_placeholder.progress(0)
+        recon_audios = []
+        for i in range(len(recons)):
+            spec_denorm = denormalize_melspec(recons[i].detach().numpy().squeeze(), audios[i][1])
+            audio_reconstructed = melspec_to_audio(spec_denorm)
+            recon_audios.append(audio_reconstructed)
+            progress_bar.progress(int((i + 1) / len(recons) * 100))
+        recon_audios = np.concatenate(recon_audios)
+        progress_bar_placeholder.empty()
+        recon_audio_text_placeholder.empty()
+        
+        return recon_audios
+   
+     
+def run():    
+    model = load_models()
+    uploaded_audio = st.file_uploader("Tải lên 1 audio (chỉ xử lý 15s đầu tiên)", type=['wav', 'mp3'])
+    
+    select_audio = st.selectbox(
+        "Hoặc chọn 1 audio mẫu dưới dây:",
+        options=[""] + [f"{file} - được lấy từ GTZAN Dataset" for file in os.listdir(AUDIO_SAMPLES_DIR) if file.endswith(('.wav', '.mp3'))],
+        index=0,
+        format_func=lambda x: "Không chọn audio mẫu" if x == "" else x
+    )
+
+    if uploaded_audio is not None or select_audio != "":
+        if uploaded_audio is not None:
+            st.audio(uploaded_audio, format='audio/wav')
+        else:
+            uploaded_audio = os.path.join(AUDIO_SAMPLES_DIR, select_audio.replace(" - được lấy từ GTZAN Dataset", ""))
+            st.audio(uploaded_audio, format='audio/wav')
+        
+        genres_list = st.multiselect('Chọn thể loại', uni_genres_list)
+        
+        if st.button('Xử lý Âm Thanh'):
+            result = gen_audio(model, uploaded_audio, genres_list)
+            st.write('Kết quả:')
+            st.audio(result, format='audio/wav', sample_rate=22050)
+
+run()
+
+st.markdown(
+    """
+    <style>
+    .footer {
+        position: fixed;
+        bottom: 0;
+        left: 0;
+        width: 100%;
+        background-color: #f1f1f1;
+        text-align: center;
+        padding: 10px 0;
+        font-size: 14px;
+        color: #555;
+    }
+    </style>
+    
+    <div class="footer">
+        <div>
+            <a href="https://ieeexplore.ieee.org/document/1021072">*GTZAN Dataset</a>
+        </div>
+        <div>
+            2024 AI VIETNAM | Made by <a href="https://github.com/Koii2k3/Music-Reconstruction" target="_blank">Koii2k3</a>
+        </div>
+    </div>
+    """,
+    unsafe_allow_html=True
+)

audios_samples/classical.00000.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8add3a7d1add1e157ce5de91be72372773b1ad2779742532c4f0ad1c7316f2a4
+size 1323632

audios_samples/country.00031.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb99205cc54237a0462c593bb472af9ad977806495671d5aa1e231aca9885ccc
+size 1323632

model/model.py

@@ -0,0 +1,91 @@
+import torch
+import numpy as np
+import torch.nn as nn
+
+class CVAE(nn.Module):
+    def __init__(self, d_model, latent_dim, n_frames, n_mels, n_genres):
+        super(CVAE, self).__init__()
+        self.d_model = d_model
+        self.latent_dim = latent_dim
+        self.n_frames = int(np.ceil(n_frames / 2**3))
+        self.n_mels = int(np.ceil(n_mels / 2**3))
+        self.n_genres = n_genres
+        print(self.n_frames, self.n_mels)
+        
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.Conv2d(1 + self.n_genres, d_model, kernel_size=3, stride=2, padding=1),  # [B, d, ceil(n_mels/2), ceil(n_frame/2)]
+            nn.BatchNorm2d(d_model),
+            nn.SiLU(),
+            nn.Dropout2d(0.05),
+
+            nn.Conv2d(d_model, d_model * 2, kernel_size=3, stride=2, padding=1),  # [B, 2*d, ceil(n_mels/2**2), ceil(n_frames/2**2)]
+            nn.BatchNorm2d(d_model * 2),
+            nn.SiLU(),
+            nn.Dropout2d(0.1),
+
+            nn.Conv2d(d_model * 2, d_model * 4, kernel_size=3, stride=2, padding=1),  # [B, 4*d, ceil(n_mels/2**3), ceil(n_frames/2**3)]
+            nn.BatchNorm2d(d_model * 4),
+            nn.SiLU(),
+            nn.Dropout2d(0.15),
+            
+            nn.AdaptiveAvgPool2d((1, 1)),  # [B, 4*d, 1, 1]
+            nn.Flatten()
+        )
+
+        # Latent space
+        self.fc_mu = nn.Linear(d_model * 4, latent_dim)
+        self.fc_logvar = nn.Linear(d_model * 4, latent_dim)
+
+        # Decoder
+        self.decoder_input = nn.Linear(latent_dim + self.n_genres, d_model * 4 * self.n_frames * self.n_mels)  # [B, 4*d, ceil(n_mels/2**3), ceil(n_frames/2**3)]
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(d_model * 4, d_model * 2, kernel_size=3, stride=2, padding=1, output_padding=(1, 0)),  # [B, 2*d, ceil(n_mels/2**2), ceil(n_frames/2**2)]
+            nn.BatchNorm2d(d_model * 2),
+            nn.SiLU(),
+            nn.Dropout2d(0.1),
+
+            nn.ConvTranspose2d(d_model * 2, d_model, kernel_size=3, stride=2, padding=1, output_padding=(1, 0)),  # [B, d, ceil(n_mels/2), ceil(n_frame/2)]
+            nn.BatchNorm2d(d_model),
+            nn.SiLU(),
+            nn.Dropout2d(0.05),
+
+            nn.ConvTranspose2d(d_model, 1, kernel_size=3, stride=2, padding=1, output_padding=1),  # [B, 1, n_mels, n_frame]
+            nn.Sigmoid()
+        )
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward(self, x, genres_input):
+        ori_genres_embed = genres_input.view(genres_input.size(0), -1)
+        genres_embed = ori_genres_embed.unsqueeze(-1).unsqueeze(-1)
+        genres_embed = genres_embed.expand(-1, -1, x.size(2), x.size(3))
+        x_genres = torch.cat((x, genres_embed), dim=1)
+ 
+        h = x_genres
+        shortcuts = []
+        for block in self.encoder:
+            h = block(h)
+            if isinstance(block, nn.SiLU):
+                shortcuts.append(h)  # skip-connection
+        
+        mu = self.fc_mu(h)
+        logvar = self.fc_logvar(h)
+        
+        z = self.reparameterize(mu, logvar)
+        z_genres = torch.cat((z, ori_genres_embed), dim=1)
+        
+        h_dec = self.decoder_input(z_genres)
+        h_dec = h_dec.view(-1, self.d_model * 4, self.n_frames, self.n_mels)
+        
+        for block in self.decoder:
+            if isinstance(block, nn.ConvTranspose2d) and shortcuts:
+                shortcut = shortcuts.pop()  # skip-connection
+                h_dec = h_dec + shortcut
+            h_dec = block(h_dec)
+
+        recon = h_dec[:, :, :x.size(2), :x.size(3)]
+        return recon, mu, logvar

model/model_256.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aade85a5d0f5a11612c2e4fd17a72842936b3a18e8c9c28462f4ae64cd1f9755
+size 89088777

utils.py

@@ -0,0 +1,78 @@
+import librosa
+import numpy as np
+from sklearn.preprocessing import MinMaxScaler
+
+uni_genres_list = ['House', 'Soundtrack', 'Composed Music', 'Drone', 'Instrumental', 'Ambient Electronic', 'Blues', 'Easy Listening', 'Classical', 'Jazz', 'Christmas', 'Electronic', 'Ambient', 'Lo-fi Instrumental', 'Lounge', 'Contemporary Classical', 'Indie-Rock', 'Dance', 'New Age', 'Halloween', 'Lo-fi Electronic', '20th Century Classical', 'Piano', 'Chill-out', 'Pop']
+genres2idx = {genre: idx for idx, genre in enumerate(uni_genres_list)}
+idx2genres = {idx: genre for genre, idx in genres2idx.items()}
+
+def tokenize(genres):
+    return [genres2idx[genre] for genre in genres if genre in genres2idx]
+
+def detokenize_tolist(tokens):
+    return [idx2genres[token] for token in tokens if token in idx2genres]
+
+def onehot_encode(tokens, max_genres):
+    onehot = np.zeros(max_genres)
+    onehot[tokens] = 1
+    return onehot
+
+def onehot_decode(onehot):
+    return [idx for idx, val in enumerate(onehot) if val == 1]
+
+def load_and_resample_audio(file_path, target_sr=22050, max_duration=15):
+    audio, sr = librosa.load(file_path, sr=None)
+    if sr != target_sr:
+        audio = librosa.resample(audio, orig_sr=sr, target_sr=target_sr)
+    if len(audio) > target_sr * max_duration:
+        audio = audio[:target_sr * max_duration]
+    return audio, target_sr
+
+def audio_to_melspec(audio, sr, n_mels=256, n_fft=2048, hop_length=512, to_db=False):
+    spec = librosa.feature.melspectrogram(y=audio,
+                                          sr=sr,
+                                          n_fft=n_fft,
+                                          hop_length=hop_length,
+                                          win_length=None,
+                                          window='hann',
+                                          center=True,
+                                          pad_mode='reflect',
+                                          power=2.0,
+                                          n_mels=n_mels)
+    
+    if to_db:
+        spec = librosa.power_to_db(spec, ref=np.max)
+    
+    return spec
+
+# Normalize the Mel spectrogram
+def normalize_melspec(melspec, norm_range=(0, 1)):
+    scaler = MinMaxScaler(feature_range=norm_range)
+    melspec = melspec.T
+    melspec_normalized = scaler.fit_transform(melspec)
+    return melspec_normalized.T
+ 
+# Denormalize the Mel spectrogram
+def denormalize_melspec(melspec_normalized, original_melspec, norm_range=(0, 1)):
+    scaler = MinMaxScaler(feature_range=norm_range)
+    melspec = original_melspec.T
+    scaler.fit(melspec)
+    melspec_denormalized = scaler.inverse_transform(melspec_normalized.T)
+    return melspec_denormalized.T
+
+# Function to convert Mel spectrogram back to audio
+def melspec_to_audio(melspec, sr=22050, n_fft=2048, hop_length=512, n_iter=64):
+    if np.any(melspec < 0):
+        melspec = librosa.db_to_power(melspec)
+    
+    audio_reconstructed = librosa.feature.inverse.mel_to_audio(melspec,
+                                                              sr=sr,
+                                                              n_fft=n_fft,
+                                                              hop_length=hop_length,
+                                                              win_length=None,
+                                                              window='hann',
+                                                              center=True,
+                                                              pad_mode='reflect',
+                                                              power=2.0,  # Ensure the correct inverse transformation
+                                                              n_iter=n_iter)
+    return audio_reconstructed