Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr 
+import gradio as gr
 import numpy as np
 import librosa
 import librosa.display
@@ -6,14 +6,17 @@ import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import pywt
-import os
-import json
-import random
-import tempfile
+import io
 from PIL import Image
 from tensorflow.keras.models import load_model
 from sklearn.preprocessing import StandardScaler
+import json
+import random
+import plotly.express as px
 
+# ================================
+# CẤU HÌNH
+# ================================
 SAMPLE_RATE = 22050
 MAX_DURATION = 5
 TIME_STEPS = 20
@@ -36,6 +39,9 @@ with open("label_map.json", "r") as f:
     label_map = json.load(f)
 index_to_label = {v: k for k, v in label_map.items()}
 
+# ================================
+# HÀM TIỀN XỬ LÝ
+# ================================
 def denoise_wavelet(signal, wavelet='db8', level=4):
     coeffs = pywt.wavedec(signal, wavelet, level=level)
     sigma = np.median(np.abs(coeffs[-1])) / 0.6745
@@ -52,35 +58,38 @@ def cat_2s_ngau_nhien(y, sr, duration=2):
     start = random.randint(0, len(y) - duration * sr)
     return y[start:start + duration * sr]
 
+# ================================
+# VẼ ẢNH (numpy array)
+# ================================
+def fig_to_numpy(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", dpi=90, bbox_inches="tight")
+    buf.seek(0)
+    img = Image.open(buf)
+    plt.close(fig)
+    return np.array(img)
+
 def tao_anh_mel(file_path):
     y, sr = librosa.load(file_path, sr=None, mono=True)
     y = cat_2s_ngau_nhien(y, sr)
     S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128)
     S_dB = librosa.power_to_db(S, ref=np.max)
     fig, ax = plt.subplots(figsize=(6, 3))
-    img = librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel', ax=ax, cmap='magma')
+    librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel', ax=ax, cmap='magma')
     ax.set_title("Phổ tần Mel", fontsize=10)
-    fig.colorbar(img, ax=ax)
-    plt.tight_layout()
-    path = os.path.join(tempfile.gettempdir(), "mel.png")
-    fig.savefig(path, dpi=80)
-    plt.close()
-    return Image.open(path)
+    return fig_to_numpy(fig)
 
 def tao_wavelet_transform(file_path):
     y, sr = librosa.load(file_path, sr=None, mono=True)
     y = cat_2s_ngau_nhien(y, sr)
     coef, _ = pywt.cwt(y, scales=np.arange(1, 128), wavelet='morl', sampling_period=1/sr)
     fig, ax = plt.subplots(figsize=(6, 3))
-    ax.imshow(np.abs(coef), extent=[0, len(y)/sr, 1, 128], cmap='plasma', aspect='auto', origin='lower')
+    ax.imshow(np.abs(coef), extent=[0, len(y)/sr, 1, 128],
+              cmap='plasma', aspect='auto', origin='lower')
     ax.set_title("Phổ sóng con (Wavelet)")
     ax.set_xlabel("Thời gian (s)")
     ax.set_ylabel("Tần số (scale)")
-    plt.tight_layout()
-    path = os.path.join(tempfile.gettempdir(), "wavelet.png")
-    fig.savefig(path, dpi=80)
-    plt.close()
-    return Image.open(path)
+    return fig_to_numpy(fig)
 
 def tao_waveform_image(file_path):
     y, sr = librosa.load(file_path, sr=None, mono=True)
@@ -90,28 +99,80 @@ def tao_waveform_image(file_path):
     ax.set_title("Biểu đồ Sóng Âm (Waveform)")
     ax.set_xlabel("Thời gian (s)")
     ax.set_ylabel("Biên độ")
-    plt.tight_layout()
-    path = os.path.join(tempfile.gettempdir(), "waveform.png")
-    fig.savefig(path, dpi=80)
-    plt.close()
-    return Image.open(path)
+    return fig_to_numpy(fig)
 
-def bao_san_sang(file_path):
-    if not file_path:
-        return ""
-    return "<b style='color:green;'>✅ Âm thanh đã sẵn sàng. Nhấn kiểm tra ngay!</b>"
+def tao_waveform_denoise(file_path):
+    y, sr = librosa.load(file_path, sr=None, mono=True)
+    y = cat_2s_ngau_nhien(y, sr)
+    y_denoised = denoise_wavelet(y)
+
+    fig, ax = plt.subplots(3, 2, figsize=(10, 8))
+
+    # 1. Waveform
+    librosa.display.waveshow(y, sr=sr, ax=ax[0,0], color='red')
+    ax[0,0].set_title("Waveform - Trước lọc")
+    librosa.display.waveshow(y_denoised, sr=sr, ax=ax[0,1], color='green')
+    ax[0,1].set_title("Waveform - Sau lọc")
+
+    # 2. FFT
+    freqs = np.fft.rfftfreq(len(y), 1/sr)
+    fft_y = np.abs(np.fft.rfft(y))
+    fft_y_denoised = np.abs(np.fft.rfft(y_denoised))
 
-def sinh_anh(file_path):
+    ax[1,0].plot(freqs, fft_y, color='red')
+    ax[1,0].set_xlim(0, 8000)
+    ax[1,0].set_title("FFT - Trước lọc")
+
+    ax[1,1].plot(freqs, fft_y_denoised, color='green')
+    ax[1,1].set_xlim(0, 8000)
+    ax[1,1].set_title("FFT - Sau lọc")
+
+    # 3. Spectrogram
+    D1 = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
+    D2 = librosa.amplitude_to_db(np.abs(librosa.stft(y_denoised)), ref=np.max)
+
+    img1 = librosa.display.specshow(D1, sr=sr, x_axis='time', y_axis='log', ax=ax[2,0], cmap="magma")
+    ax[2,0].set_title("Spectrogram - Trước lọc")
+    fig.colorbar(img1, ax=ax[2,0], format="%+2.0f dB")
+
+    img2 = librosa.display.specshow(D2, sr=sr, x_axis='time', y_axis='log', ax=ax[2,1], cmap="magma")
+    ax[2,1].set_title("Spectrogram - Sau lọc")
+    fig.colorbar(img2, ax=ax[2,1], format="%+2.0f dB")
+
+    plt.tight_layout()
+    return fig_to_numpy(fig)
+
+# ================================
+# VẼ BIỂU ĐỒ Top-3 (Plotly)
+# ================================
+def ve_top3_chart(probs):
+    labels = [index_to_label[i] for i in range(len(probs))]
+    values = probs * 100
+    top_idx = np.argsort(values)[::-1][:3]
+    fig = px.pie(
+        values=[values[i] for i in top_idx],
+        names=[labels[i] for i in top_idx],
+        title="Top-3 dự đoán"
+    )
+    return fig
+
+# ================================
+# DỰ ĐOÁN
+# ================================
+def bao_san_sang(file_path):
     if not file_path:
-        return None, None, None
-    mel_img = tao_anh_mel(file_path)
-    wavelet_img = tao_wavelet_transform(file_path)
-    waveform_img = tao_waveform_image(file_path)
-    return mel_img, wavelet_img, waveform_img
+        return "", None, None, None, None
+    return (
+        "<b style='color:green;'>✅ Âm thanh đã sẵn sàng. Nhấn kiểm tra ngay!</b>",
+        tao_anh_mel(file_path),
+        tao_wavelet_transform(file_path),
+        tao_waveform_image(file_path),
+        tao_waveform_denoise(file_path)
+    )
 
 def du_doan(file_path):
     if not file_path:
-        return "<b style='color:red;'>❌ Chưa có âm thanh.</b>"
+        return "<b style='color:red;'>❌ Chưa có âm thanh.</b>", None
 
     signal, sr = librosa.load(file_path, sr=SAMPLE_RATE, mono=True)
     signal, _ = librosa.effects.trim(signal)
@@ -125,10 +186,11 @@ def du_doan(file_path):
     X_input = create_sequences(mfcc, time_steps=TIME_STEPS)
 
     if len(X_input) == 0:
-        return "<b style='color:red;'>⚠️ Âm thanh quá ngắn để phân tích.</b>"
+        return "<b style='color:red;'>⚠️ Âm thanh quá ngắn để phân tích.</b>", None
 
     y_preds = model.predict(X_input, verbose=0)
     avg_probs = np.mean(y_preds, axis=0)
+
     pred_index = np.argmax(avg_probs)
     confidence = avg_probs[pred_index] * 100
     pred_label = "HƯ HỎNG KHÁC" if confidence < 50 else index_to_label[pred_index]
@@ -142,14 +204,18 @@ def du_doan(file_path):
     for i, prob in enumerate(avg_probs):
         html += f"<span style='color:#000'>- {index_to_label[i]}: {prob*100:.1f}%</span><br>"
     html += "</div>"
-    return html
 
-def reset_output():
-    return "", None, None, None, ""
+    return html, ve_top3_chart(avg_probs)
 
-def chon_file(f1, f2):
-    return f1 if f1 else f2
+# ================================
+# RESET
+# ================================
+def reset_output():
+    return "", None, None, None, None, "", None
 
+# ================================
+# GIAO DIỆN
+# ================================
 with gr.Blocks(css="""
 #check-btn {
     background: #007acc;
@@ -200,33 +266,40 @@ with gr.Blocks(css="""
     output_html = gr.HTML()
 
     with gr.Accordion("📊 Phân tích Âm Thanh", open=False):
-        mel_output = gr.Image(label="")
-        wavelet_output = gr.Image(label="")
-        waveform_output = gr.Image(label="")
-
-    audio_file.change(fn=bao_san_sang, inputs=audio_file, outputs=thong_bao_ready)
-    audio_mic.change(fn=bao_san_sang, inputs=audio_mic, outputs=thong_bao_ready)
-
-    audio_file.change(fn=sinh_anh, inputs=audio_file, outputs=[mel_output, wavelet_output, waveform_output])
-    audio_mic.change(fn=sinh_anh, inputs=audio_mic, outputs=[mel_output, wavelet_output, waveform_output])
-
-    btn_check.click(fn=lambda f1, f2: du_doan(chon_file(f1, f2)),
-                    inputs=[audio_file, audio_mic],
-                    outputs=output_html)
-
+        mel_output = gr.Image(label="Mel Spectrogram", type="numpy")
+        wavelet_output = gr.Image(label="Wavelet Transform", type="numpy")
+        waveform_output = gr.Image(label="Waveform", type="numpy")
+        waveform_denoise_output = gr.Image(label="So sánh chi tiết trước & sau lọc nhiễu (Waveform + FFT + Spectrogram)", type="numpy")
+        top3_chart = gr.Plot(label="Top-3 dự đoán")
+
+    # --- Upload/ghi âm → chỉ báo sẵn sàng + vẽ ảnh
+    audio_file.change(
+        fn=bao_san_sang,
+        inputs=audio_file,
+        outputs=[thong_bao_ready, mel_output, wavelet_output, waveform_output, waveform_denoise_output]
+    )
+
+    audio_mic.change(
+        fn=bao_san_sang,
+        inputs=audio_mic,
+        outputs=[thong_bao_ready, mel_output, wavelet_output, waveform_output, waveform_denoise_output]
+    )
+
+    # --- Nút kiểm tra → chỉ dự đoán
+    btn_check.click(
+        fn=lambda f1, f2: du_doan(f1 if f1 else f2),
+        inputs=[audio_file, audio_mic],
+        outputs=[output_html, top3_chart]
+    )
+
+    # --- Reset khi clear
     audio_file.clear(fn=reset_output, outputs=[
-        thong_bao_ready,
-        mel_output,
-        wavelet_output,
-        waveform_output,
-        output_html
+        thong_bao_ready, mel_output, wavelet_output, waveform_output, waveform_denoise_output,
+        output_html, top3_chart
     ])
     audio_mic.clear(fn=reset_output, outputs=[
-        thong_bao_ready,
-        mel_output,
-        wavelet_output,
-        waveform_output,
-        output_html
+        thong_bao_ready, mel_output, wavelet_output, waveform_output, waveform_denoise_output,
+        output_html, top3_chart
     ])
 
-demo.launch()
+demo.launch()