Upload inference.py with huggingface_hub

inference.py

@@ -629,7 +629,7 @@ def _load_gru_model():
 
         class HeartSoundGRU(nn.Module):
             def __init__(self, input_dim=129, hidden_dim=64, num_layers=2,
-                         num_classes=NUM_CLASSES, dropout=0.4):
+                         num_classes=2, dropout=0.4):
                 super().__init__()
                 self.input_norm = nn.LayerNorm(input_dim)
                 self.gru = nn.GRU(input_dim, hidden_dim, num_layers,
@@ -677,6 +677,7 @@ def predict_gru(y, sr):
     N_FFT_G, HOP_G, CLIP_SEC = 256, 64, 5
     target_len = GRU_SR * CLIP_SEC
     clips = [y_4k[s:s+target_len] for s in range(0, len(y_4k)-target_len+1, target_len)] if len(y_4k) >= target_len else [np.pad(y_4k, (0, target_len-len(y_4k)))]
+    GRU_BINARY_NAMES = ["Normal", "Murmur"]
     probs = []
     for clip in clips:
         S = np.abs(librosa.stft(clip, n_fft=N_FFT_G, hop_length=HOP_G)) ** 2
@@ -688,10 +689,12 @@ def predict_gru(y, sr):
             probs.append(torch.softmax(_gru_model(t), 1)[0].numpy())
     avg = np.mean(probs, 0)
     pred = int(np.argmax(avg))
+    is_murmur = pred == 1
+    label = GRU_BINARY_NAMES[pred]
     return {
-        "label": CLASS_NAMES[pred], "confidence": round(float(avg[pred]), 4),
-        "is_disease": pred != 0, "method": "Bi-GRU (McDonald et al., Cambridge 2024)",
-        "all_classes": [{"label": CLASS_NAMES[i], "probability": round(float(avg[i]), 4)} for i in range(NUM_CLASSES)],
+        "label": label, "confidence": round(float(avg[pred]), 4),
+        "is_disease": is_murmur, "method": "Bi-GRU Binary (McDonald et al., Cambridge 2024)",
+        "all_classes": [{"label": GRU_BINARY_NAMES[i], "probability": round(float(avg[i]), 4)} for i in range(2)],
     }