Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 import torchaudio
 import numpy as np
 from datasets import load_dataset
+import torch.nn.functional as F
 
 # ---------------------------
 # Load Dataset for Label Reference
@@ -23,13 +24,19 @@ def fake_quality_score(mel_spec):
 
 # ---------------------------
 # Audio Preprocessing
-# NOTE: Gradio with type="numpy" gives (sample_rate, data)
 # ---------------------------
+
+TARGET_SR = 44100
+N_FFT = 1024
+HOP_LENGTH = 512
+N_MELS = 64
+
 mel_transform = torchaudio.transforms.MelSpectrogram(
-    sample_rate=44100,
-    n_fft=1024,
-    hop_length=512,
-    n_mels=64
+    sample_rate=TARGET_SR,
+    n_fft=N_FFT,
+    hop_length=HOP_LENGTH,
+    n_mels=N_MELS,
+    center=False  # we will handle padding manually
 )
 
 def preprocess_audio(audio):
@@ -48,20 +55,27 @@ def preprocess_audio(audio):
 
     # If shape is (samples, channels), transpose to (channels, samples)
     if waveform.ndim == 2 and waveform.shape[0] < waveform.shape[1]:
-        # shape (samples, channels) -> (channels, samples)
         waveform = waveform.transpose(0, 1)
 
     # Convert to mono if stereo or more channels
     if waveform.shape[0] > 1:
         waveform = waveform.mean(dim=0, keepdim=True)
 
-    # Resample to 44.1 kHz if needed
-    if sr != 44100:
-        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=44100)
+    # Resample to TARGET_SR if needed
+    if sr != TARGET_SR:
+        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=TARGET_SR)
         waveform = resampler(waveform)
-        sr = 44100
+        sr = TARGET_SR
+
+    # --- NEW: Ensure minimum length for STFT / MelSpectrogram ---
+    min_len = N_FFT  # at least one window
+    current_len = waveform.shape[-1]
+    if current_len < min_len:
+        pad_amount = min_len - current_len
+        # Pad at the end with zeros
+        waveform = F.pad(waveform, (0, pad_amount))
 
-    # Mel-spectrogram
+    # Mel-spectrogram (no internal centering/padding)
     mel = mel_transform(waveform)
     mel_db = torchaudio.transforms.AmplitudeToDB()(mel)
     return mel_db
@@ -71,7 +85,7 @@ def preprocess_audio(audio):
 # ---------------------------
 def analyze_piano(audio):
     if audio is None:
-        return "Please upload or record a piano audio clip."
+        return "Please upload or record a piano audio clip (at least 1–2 seconds)."
 
     try:
         mel = preprocess_audio(audio)
@@ -102,7 +116,7 @@ demo = gr.Interface(
     ),
     outputs=gr.Textbox(label="AI Analysis Output"),
     title="AI Piano Sound Analyzer 🎹",
-    description="Upload a short piano recording to get a predicted piano type and estimated sound-quality score."
+    description="Upload a short piano recording (around 1–3 seconds) to get a predicted piano type and estimated sound-quality score."
 )
 
 if __name__ == "__main__":