Upload inference.py with huggingface_hub

inference.py

@@ -26,6 +26,21 @@ _cnn_available = None
 
 TARGET_SR = 16000
 
+# 4-class murmur timing labels
+CLASS_NAMES = ["Normal", "Systolic Murmur", "Diastolic Murmur", "Continuous Murmur"]
+NUM_CLASSES  = 4
+
+# Brief clinical notes per murmur type (shown in UI + PDF)
+MURMUR_TYPE_NOTES = {
+    "Normal":           "No murmur detected. Heart sounds are within normal limits.",
+    "Systolic Murmur":  "Systolic murmur (S1→S2). Common causes: mitral insufficiency, "
+                        "pulmonic or aortic stenosis, VSD. Recommend echocardiography.",
+    "Diastolic Murmur": "Diastolic murmur (S2→S1). Uncommon in dogs — often indicates "
+                        "aortic insufficiency. Specialist evaluation strongly advised.",
+    "Continuous Murmur":"Continuous (machinery) murmur throughout the cardiac cycle. "
+                        "Classic finding in patent ductus arteriosus (PDA). Urgent referral advised.",
+}
+
 print("CardioScreen AI engine loaded (lightweight mode)", flush=True)
 
 
@@ -329,7 +344,7 @@ def _load_cnn_model():
         import torch.nn as nn
 
         class HeartSoundCNN(nn.Module):
-            def __init__(self):
+            def __init__(self, num_classes=NUM_CLASSES):
                 super().__init__()
                 self.features = nn.Sequential(
                     nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
@@ -337,7 +352,7 @@ def _load_cnn_model():
                     nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
                     nn.AdaptiveAvgPool2d((1, 1)),
                 )
-                self.classifier = nn.Sequential(nn.Dropout(0.3), nn.Linear(128, 2))
+                self.classifier = nn.Sequential(nn.Dropout(0.3), nn.Linear(128, num_classes))
 
             def forward(self, x):
                 x = self.features(x)
@@ -383,7 +398,10 @@ def _load_cnn_model():
 
 
 def predict_cnn(y, sr):
-    """Classify audio using the trained Mel-spectrogram CNN."""
+    """
+    Classify audio using the trained Mel-spectrogram CNN (4-class).
+    Returns Normal / Systolic Murmur / Diastolic Murmur / Continuous Murmur.
+    """
     if not _load_cnn_model():
         return None
 
@@ -391,7 +409,7 @@ def predict_cnn(y, sr):
 
     # Config must match training
     N_MELS, N_FFT, HOP = 64, 1024, 512
-    CLIP_SEC = 5
+    CLIP_SEC   = 5
     target_len = sr * CLIP_SEC
 
     # Split into 5-sec clips
@@ -406,11 +424,10 @@ def predict_cnn(y, sr):
     target_t = int(np.ceil(CLIP_SEC * sr / HOP))
     probs = []
     for clip in clips:
-        S = librosa.feature.melspectrogram(y=clip, sr=sr, n_mels=N_MELS, n_fft=N_FFT, hop_length=HOP)
+        S    = librosa.feature.melspectrogram(y=clip, sr=sr, n_mels=N_MELS, n_fft=N_FFT, hop_length=HOP)
         S_db = librosa.power_to_db(S, ref=np.max)
         S_db = (S_db - S_db.min()) / (S_db.max() - S_db.min() + 1e-8)
 
-        # Pad/truncate time axis
         if S_db.shape[1] < target_t:
             S_db = np.pad(S_db, ((0, 0), (0, target_t - S_db.shape[1])))
         else:
@@ -419,28 +436,44 @@ def predict_cnn(y, sr):
         tensor = torch.FloatTensor(S_db).unsqueeze(0).unsqueeze(0)  # (1,1,64,T)
         with torch.no_grad():
             logits = _cnn_model(tensor)
-            p = torch.softmax(logits, dim=1)[0]  # [P(normal), P(murmur)]
+            p = torch.softmax(logits, dim=1)[0]  # shape: (NUM_CLASSES,)
             probs.append(p.numpy())
 
     # Average probabilities across clips
-    avg_prob = np.mean(probs, axis=0)
-    normal_p = float(avg_prob[0])
-    murmur_p = float(avg_prob[1])
-    # Optimized threshold (0.30) — validated via sweep on patient-level split
-    # Gives 96.3% sensitivity + 96.0% specificity (both > 90%)
+    avg_prob = np.mean(probs, axis=0)  # (NUM_CLASSES,)
+
+    # --- Murmur detection threshold (binary: Normal vs. any murmur type) ---
+    # P(any murmur) = 1 - P(Normal).  Threshold 0.30 keeps high sensitivity.
+    normal_p   = float(avg_prob[0])
+    murmur_p   = float(1.0 - normal_p)   # P(any murmur type)
     MURMUR_THRESHOLD = 0.30
-    is_murmur = murmur_p > MURMUR_THRESHOLD
+    is_murmur  = murmur_p > MURMUR_THRESHOLD
+
+    # --- Murmur type: argmax over 4 classes ---
+    predicted_class = int(np.argmax(avg_prob))
+    # If we detect a murmur but the model's top class is Normal (border case),
+    # fall back to the highest-probability murmur subclass.
+    if is_murmur and predicted_class == 0:
+        predicted_class = int(np.argmax(avg_prob[1:])) + 1
+
+    murmur_type       = CLASS_NAMES[predicted_class]
+    type_confidence   = float(avg_prob[predicted_class])
+    overall_label     = murmur_type if is_murmur else "Normal"
+    overall_conf      = round(murmur_p if is_murmur else normal_p, 4)
 
     return {
-        "label": "Murmur" if is_murmur else "Normal",
-        "confidence": round(murmur_p if is_murmur else normal_p, 4),
-        "is_disease": bool(is_murmur),
-        "method": "CNN (Mel-Spectrogram)",
+        "label":       overall_label,
+        "confidence":  overall_conf,
+        "is_disease":  bool(is_murmur),
+        "murmur_type": murmur_type,
+        "murmur_type_confidence": round(type_confidence, 4),
+        "murmur_type_note": MURMUR_TYPE_NOTES.get(murmur_type, ""),
+        "method":      "CNN (Mel-Spectrogram, 4-class)",
         "clips_analyzed": len(clips),
         "all_classes": [
-            {"label": "Normal", "probability": round(normal_p, 4)},
-            {"label": "Murmur", "probability": round(murmur_p, 4)},
-        ]
+            {"label": CLASS_NAMES[i], "probability": round(float(avg_prob[i]), 4)}
+            for i in range(NUM_CLASSES)
+        ],
     }
 
 
@@ -683,22 +716,33 @@ def predict_audio(audio_bytes: bytes):
         # Combined summary — CNN is the sole decision-maker
         dsp_disease = dsp_result["is_disease"]
         cnn_disease = cnn_result["is_disease"] if cnn_result else dsp_disease
-        is_disease = cnn_disease  # top-level flag driven by CNN only
+        is_disease  = cnn_disease  # top-level flag driven by CNN only
+
+        # Murmur type from CNN (None if no CNN or no murmur)
+        murmur_type      = None
+        murmur_type_conf = None
+        murmur_type_note = None
+        if cnn_result and cnn_disease:
+            murmur_type      = cnn_result.get("murmur_type", "Murmur")
+            murmur_type_conf = cnn_result.get("murmur_type_confidence")
+            murmur_type_note = cnn_result.get("murmur_type_note", "")
 
         if quality["grade"] == "Poor":
-            summary = "⚠️ Poor recording quality — results may be unreliable, please re-record"
+            summary   = "⚠️ Poor recording quality — results may be unreliable, please re-record"
             agreement = "poor_quality"
         elif cnn_disease and dsp_disease:
-            summary = "⚠️ Murmur detected — confirmed by both CNN and DSP analysis"
+            type_str  = f" ({murmur_type})" if murmur_type else ""
+            summary   = f"⚠️ Murmur detected{type_str} — confirmed by both CNN and DSP analysis"
             agreement = "both_murmur"
         elif cnn_disease and not dsp_disease:
-            summary = "⚠️ Murmur detected by CNN — DSP analysis was inconclusive"
+            type_str  = f" ({murmur_type})" if murmur_type else ""
+            summary   = f"⚠️ Murmur detected{type_str} by CNN — DSP analysis was inconclusive"
             agreement = "cnn_only"
         elif not cnn_disease and dsp_disease:
-            summary = "Normal heart sound (CNN) — DSP flagged minor irregularity, likely artifact"
+            summary   = "Normal heart sound (CNN) — DSP flagged minor irregularity, likely artifact"
             agreement = "dsp_only"
         else:
-            summary = "Normal heart sound — no murmur detected"
+            summary   = "Normal heart sound — no murmur detected"
             agreement = "both_normal"
 
         # Downsample waveform for frontend (~800 points)
@@ -711,18 +755,21 @@ def predict_audio(audio_bytes: bytes):
         peak_vis_indices = [int(p // step) for p in peaks if int(p // step) < vis_duration]
 
         return {
-            "bpm": bpm,
-            "heartbeat_count": heartbeat_count,
+            "bpm":              bpm,
+            "heartbeat_count":  heartbeat_count,
             "duration_seconds": round(duration, 1),
-            "is_disease": is_disease,                 # CNN-driven decision
-            "agreement": agreement,                   # how DSP & CNN align
+            "is_disease":       is_disease,           # CNN-driven binary decision
+            "murmur_type":      murmur_type,          # NEW: "Systolic Murmur" / "Diastolic Murmur" / "Continuous Murmur" / None
+            "murmur_type_confidence": murmur_type_conf,
+            "murmur_type_note": murmur_type_note,     # clinical description
+            "agreement":        agreement,
             "clinical_summary": summary,
-            "heart_score": heart_score,
-            "ai_classification": dsp_result,          # backward compatible
-            "dsp_classification": dsp_result,         # explicit DSP (supplementary)
-            "cnn_classification": cnn_result,         # CNN (primary, or None)
-            "signal_quality": quality,
-            "waveform": vis_waveform,
+            "heart_score":      heart_score,
+            "ai_classification":  dsp_result,         # backward compatible
+            "dsp_classification":  dsp_result,        # explicit DSP (supplementary)
+            "cnn_classification":  cnn_result,        # CNN (primary, or None)
+            "signal_quality":   quality,
+            "waveform":         vis_waveform,
             "peak_times_seconds": peak_times_sec,
             "peak_vis_indices": peak_vis_indices,
         }