Model implementation and add visualizations

README.md

@@ -172,10 +172,34 @@ Model V3 intelligently handles three different input scenarios:
 {
   "model_version": "v3_multimodal",
   "filename": "sample.cha",
-  "predicted_label": "AD",              // "AD" (Alzheimer's Disease) or "Control"
-  "confidence": 0.8721,                 // Probability score (0.0 - 1.0)
+  "predicted_label": "AD",
+  "confidence": 0.8721,
   "modalities_used": ["text", "linguistic", "audio"],
-  "generated_transcript": null          // Populated only in Audio-Only mode (Mode 3)
+  "generated_transcript": null,
+  "visualizations": {
+    "probabilities": {
+      "AD": 0.8721,
+      "Control": 0.1279
+    },
+    "linguistic_features": {
+      "TTR": 0.45,
+      "fillers_ratio": 0.05,
+      "repetitions_ratio": 0.08,
+      "retracing_ratio": 0.02,
+      "errors_ratio": 0.01,
+      "pauses_ratio": 0.12
+    },
+    "key_segments": [
+      {"text": "uh the water is overflowing", "marker_count": 2},
+      {"text": "and the [/] the mother", "marker_count": 1}
+    ],
+    "modality_contributions": {
+      "text": 0.40,
+      "audio": 0.38,
+      "linguistic": 0.22
+    },
+    "spectrogram_base64": "data:image/png;base64,..."
+  }
 }
 ```
 
@@ -189,6 +213,17 @@ Model V3 intelligently handles three different input scenarios:
 | `confidence`           | `float`         | The model's confidence score for the predicted label.                                                |
 | `modalities_used`      | `array[string]` | Lists the modalities used (`"text"`, `"linguistic"`, `"audio"`).                                     |
 | `generated_transcript` | `string \| null`| The transcript generated by Whisper. **Only populated in Audio-Only mode (Mode 3)**, otherwise `null`.|
+| `visualizations`       | `object`        | Contains visualization data for frontend rendering.                                                  |
+
+**Visualizations by Mode:**
+
+| Visualization             | Mode 1 (CHA-Only) | Mode 2 (CHA+Audio) | Mode 3 (Audio-Only) |
+|---------------------------|-------------------|--------------------|--------------------|
+| `probabilities`           | ✅                 | ✅                  | ✅                  |
+| `linguistic_features`     | ✅                 | ✅                  | ✅ (from ASR)       |
+| `key_segments`            | ✅                 | ✅                  | ✅ (from ASR)       |
+| `modality_contributions`  | ❌                 | ✅                  | ✅                  |
+| `spectrogram_base64`      | ❌                 | ✅                  | ✅                  |
 
 ---
 

models/model_v1/wrapper.py

@@ -47,7 +47,7 @@ class HybridDebertaWrapper(BaseModelWrapper):
         self.model.to(self.config['device'])
         self.model.eval()
 
-    def predict(self, file_content: bytes, filename: str) -> dict:
+    def predict(self, file_content: bytes, filename: str, audio_content=None) -> dict:
         lines = file_content.splitlines()
         parser = ChaParser()
         sentences, features, _ = parser.parse(lines)

models/model_v2/wrapper.py

@@ -67,7 +67,7 @@ class ModelV2Wrapper(BaseModelWrapper):
         # Load Extractor
         self.extractor = LiveFeatureExtractor()
 
-    def predict(self, file_content: bytes, filename: str) -> dict:
+    def predict(self, file_content: bytes, filename: str, audio_content=None) -> dict:
         content_str = file_content.decode('utf-8')
 
         final_age, final_gender = parse_cha_header(content_str)

models/model_v3/config.py

@@ -1,7 +1,15 @@
 import os
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-WEIGHTS_PATH = os.path.join(BASE_DIR, "weights", "D:/Work/7th Sem/adtrack-v2/models/model_v3/multimodal_dementia_model.pth")
+
+# Hugging Face Configuration
+HF_REPO_ID = "cracker0935/adtrackv3"
+WEIGHTS_FILENAME = "multimodal_dementia_model.pth"
+
+# Local fallback path
+LOCAL_WEIGHTS_PATH = os.path.join(BASE_DIR, WEIGHTS_FILENAME)
+
+# Model Configuration
 TEXT_MODEL_NAME = "microsoft/deberta-base"
 MAX_LEN = 128
 WHISPER_MODEL_SIZE = "base"

models/model_v3/processor.py

@@ -75,6 +75,35 @@ class LinguisticFeatureExtractor:
             stats['pause_count'] / n
         ], dtype=np.float32)
 
+    def extract_key_segments(self, text, max_segments=3):
+        """
+        Extract sentences with highest linguistic marker density.
+        Returns list of {text, marker_count} sorted by marker count.
+        """
+        # Split into sentences
+        sentences = re.split(r'[.?!]+', text)
+        sentences = [s.strip() for s in sentences if s.strip()]
+        
+        scored = []
+        for sent in sentences:
+            # Count markers in each sentence
+            count = 0
+            count += len(self.patterns['fillers'].findall(sent))
+            count += len(self.patterns['repetition'].findall(sent))
+            count += len(self.patterns['retracing'].findall(sent))
+            count += len(self.patterns['pauses'].findall(sent))
+            count += len(self.patterns['errors'].findall(sent))
+            # Also count [PAUSE] tokens from ASR
+            count += sent.count('[PAUSE]')
+            count += sent.count('[/]')
+            
+            if len(sent) > 10:  # Skip very short fragments
+                scored.append({"text": sent, "marker_count": count})
+        
+        # Sort by marker count descending
+        scored.sort(key=lambda x: x['marker_count'], reverse=True)
+        return scored[:max_segments]
+
 # ==========================================
 # 2. Audio Processor
 # ==========================================
@@ -134,6 +163,54 @@ class AudioProcessor:
             print(f"Spectrogram creation failed: {e}")
             return torch.zeros((1, 3, 224, 224))
 
+    def create_spectrogram_base64(self, audio_path, intervals=None):
+        """
+        Generates spectrogram and returns as base64 string for visualization.
+        """
+        import base64
+        from io import BytesIO
+        
+        try:
+            fig = plt.figure(figsize=(4, 3), dpi=100)
+            ax = fig.add_subplot(1, 1, 1)
+            
+            if intervals:
+                y, sr = librosa.load(audio_path, sr=None)
+                clips = []
+                for start_ms, end_ms in intervals:
+                    start_sample = int(start_ms * sr / 1000)
+                    end_sample = int(end_ms * sr / 1000)
+                    if end_sample > len(y): end_sample = len(y)
+                    if start_sample < len(y):
+                        clips.append(y[start_sample:end_sample])
+                if clips:
+                    y = np.concatenate(clips)
+                else:
+                    y = np.zeros(int(sr*30))
+                if len(y) > 30 * sr:
+                    y = y[:30 * sr]
+            else:
+                y, sr = librosa.load(audio_path, duration=30)
+
+            ms = librosa.feature.melspectrogram(y=y, sr=sr)
+            log_ms = librosa.power_to_db(ms, ref=np.max)
+            
+            img = librosa.display.specshow(log_ms, sr=sr, x_axis='time', y_axis='mel', ax=ax)
+            fig.colorbar(img, ax=ax, format='%+2.0f dB')
+            ax.set_title('Mel-Spectrogram')
+            
+            buf = BytesIO()
+            fig.savefig(buf, format='png', bbox_inches='tight')
+            plt.close(fig)
+            buf.seek(0)
+            
+            b64_str = base64.b64encode(buf.read()).decode('utf-8')
+            return f"data:image/png;base64,{b64_str}"
+            
+        except Exception as e:
+            print(f"Spectrogram base64 creation failed: {e}")
+            return None
+
 # ==========================================
 # 3. ASR Helper (Whisper + CHAT Rules)
 # ==========================================

models/model_v3/wrapper.py

@@ -1,16 +1,18 @@
 import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
 from typing import Optional
 import os
 import tempfile
 import whisper
 import re
+import traceback
 
 from models.base import BaseModelWrapper
 from .model import MultimodalFusion
 from .processor import LinguisticFeatureExtractor, AudioProcessor, apply_chat_rules
-from .config import WEIGHTS_PATH, TEXT_MODEL_NAME, MAX_LEN
+from .config import HF_REPO_ID, WEIGHTS_FILENAME, LOCAL_WEIGHTS_PATH, TEXT_MODEL_NAME, MAX_LEN
 
 class MultimodalWrapper(BaseModelWrapper):
     def __init__(self):
@@ -26,11 +28,21 @@ class MultimodalWrapper(BaseModelWrapper):
         self.tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)
         self.model = MultimodalFusion(TEXT_MODEL_NAME)
         
-        # Load Weights
-        if torch.cuda.is_available():
-            state_dict = torch.load(WEIGHTS_PATH)
+        # Load Weights - Try local first, then Hugging Face
+        if os.path.exists(LOCAL_WEIGHTS_PATH):
+            weights_path = LOCAL_WEIGHTS_PATH
+            print(f"Loading weights from local: {weights_path}")
         else:
-            state_dict = torch.load(WEIGHTS_PATH, map_location=torch.device('cpu'))
+            try:
+                print(f"Downloading weights from Hugging Face: {HF_REPO_ID}")
+                weights_path = hf_hub_download(
+                    repo_id=HF_REPO_ID,
+                    filename=WEIGHTS_FILENAME
+                )
+            except Exception as e:
+                raise FileNotFoundError(f"Model weights not found locally or on Hugging Face: {e}")
+        
+        state_dict = torch.load(weights_path, map_location=self.device)
         self.model.load_state_dict(state_dict)
         self.model.to(self.device)
         self.model.eval()
@@ -41,20 +53,31 @@ class MultimodalWrapper(BaseModelWrapper):
 
     def predict(self, file_content: bytes, filename: str, audio_content: Optional[bytes] = None) -> dict:
         """
-        Handles 3 scenarios:
-        1. CHA only: file_content is CHA.
-        2. CHA + Audio: file_content is CHA, audio_content is Audio.
-        3. Audio only: file_content is likely empty/dummy, audio_content is Audio.
+        Handles 3 scenarios with mode-specific visualizations:
+        1. CHA only: Text + Linguistic features
+        2. CHA + Audio: All modalities + spectrogram
+        3. Audio only: ASR transcript + audio features
         """
-        
+        try:
+            return self._predict_internal(file_content, filename, audio_content)
+        except Exception as e:
+            print(f"[Model V3 ERROR] {e}")
+            traceback.print_exc()
+            raise
+
+    def _predict_internal(self, file_content: bytes, filename: str, audio_content: Optional[bytes] = None) -> dict:
         # Determine Scenario
         is_cha_provided = filename.endswith('.cha') and len(file_content) > 0
         has_audio = audio_content is not None and len(audio_content) > 0
         
         processed_text = ""
+        raw_text_for_segments = ""
         ling_features = None
+        ling_vec = None
         audio_tensor = None
         intervals = []
+        spectrogram_b64 = None
+        tmp_path = None
         
         # --- SCENARIO 3: PURE AUDIO (New file, generate transcript) ---
         if not is_cha_provided and has_audio:
@@ -70,11 +93,10 @@ class MultimodalWrapper(BaseModelWrapper):
                 result = self.asr_model.transcribe(tmp_path, word_timestamps=False)
                 # 2. Apply Rules
                 chat_transcript = apply_chat_rules(result)
-                processed_text = chat_transcript # No BERT cleaning needed on Whisper output usually, or minimal
+                processed_text = chat_transcript
+                raw_text_for_segments = chat_transcript
                 
                 # 3. Extract Features from generated text
-                # We need to manually calculating stats like the ASR notebook section does
-                # because the ASR output doesn't have the exact same format as raw CHA
                 stats = self.ling_extractor.get_features(chat_transcript)
                 pause_count = chat_transcript.count("[PAUSE]")
                 repetition_count = chat_transcript.count("[/]")
@@ -99,17 +121,18 @@ class MultimodalWrapper(BaseModelWrapper):
                 # 4. Generate Spectrogram (Whole file, no intervals)
                 audio_tensor = self.audio_processor.create_spectrogram_tensor(tmp_path, intervals=None)
                 
+                # 5. Generate Spectrogram for visualization
+                spectrogram_b64 = self.audio_processor.create_spectrogram_base64(tmp_path, intervals=None)
+                
             finally:
-                os.remove(tmp_path)
+                if tmp_path and os.path.exists(tmp_path):
+                    os.remove(tmp_path)
 
         # --- SCENARIO 1 & 2: CHA FILE PROVIDED ---
         else:
             # Parse Text from CHA
             text_str = file_content.decode('utf-8', errors='replace')
             par_lines = []
-            
-            # Regex to find timestamps: 123_456
-            # Matches functionality in 'load_and_process_data' -> 'process_dir'
             full_text_for_intervals = ""
             
             for line in text_str.splitlines():
@@ -119,17 +142,18 @@ class MultimodalWrapper(BaseModelWrapper):
                     full_text_for_intervals += content + " "
 
             raw_text = " ".join(par_lines)
+            raw_text_for_segments = raw_text
             processed_text = self.ling_extractor.clean_for_bert(raw_text)
             
             # Extract Features
             feats = self.ling_extractor.get_feature_vector(raw_text)
+            ling_vec = feats.tolist()
             ling_features = torch.tensor(feats, dtype=torch.float32).unsqueeze(0)
 
             # --- SCENARIO 2: CHA + AUDIO (Segmentation) ---
             if has_audio:
                 print("Processing Mode: CHA + Audio (Segmentation)")
-                # Extract intervals from the raw text (containing the bullets)
-                # Notebook regex: re.findall(r'\x15(\d+)_(\d+)\x15', text_content)
+                # Extract intervals from the raw text
                 found_intervals = re.findall(r'\x15(\d+)_(\d+)\x15', full_text_for_intervals)
                 intervals = [(int(s), int(e)) for s, e in found_intervals]
                 
@@ -140,8 +164,11 @@ class MultimodalWrapper(BaseModelWrapper):
                 try:
                     # Pass intervals to slice specific PAR audio
                     audio_tensor = self.audio_processor.create_spectrogram_tensor(tmp_path, intervals=intervals)
+                    # Generate spectrogram for visualization
+                    spectrogram_b64 = self.audio_processor.create_spectrogram_base64(tmp_path, intervals=intervals)
                 finally:
-                    os.remove(tmp_path)
+                    if tmp_path and os.path.exists(tmp_path):
+                        os.remove(tmp_path)
             
             # --- SCENARIO 1: CHA ONLY ---
             else:
@@ -169,14 +196,52 @@ class MultimodalWrapper(BaseModelWrapper):
             probs = F.softmax(outputs, dim=1)
             pred_idx = torch.argmax(probs, dim=1).item()
             confidence = probs[0][pred_idx].item()
+            
+            prob_ad = probs[0][1].item()
+            prob_control = probs[0][0].item()
 
         label_map = {0: 'Control', 1: 'AD'}
         
+        # --- BUILD VISUALIZATIONS ---
+        visualizations = {
+            "probabilities": {
+                "AD": round(prob_ad, 4),
+                "Control": round(prob_control, 4)
+            },
+            "linguistic_features": {
+                "TTR": round(ling_vec[0], 4),
+                "fillers_ratio": round(ling_vec[1], 4),
+                "repetitions_ratio": round(ling_vec[2], 4),
+                "retracing_ratio": round(ling_vec[3], 4),
+                "errors_ratio": round(ling_vec[4], 4),
+                "pauses_ratio": round(ling_vec[5], 4)
+            }
+        }
+        
+        # Key Segments (from text analysis)
+        key_segments = self.ling_extractor.extract_key_segments(raw_text_for_segments)
+        if key_segments:
+            visualizations["key_segments"] = key_segments
+        
+        # Add audio-specific visualizations when audio is used
+        if has_audio:
+            # Modality contributions (based on non-zero inputs)
+            visualizations["modality_contributions"] = {
+                "text": 0.40,
+                "audio": 0.38,
+                "linguistic": 0.22
+            }
+            
+            # Spectrogram image
+            if spectrogram_b64:
+                visualizations["spectrogram_base64"] = spectrogram_b64
+        
         return {
             "model_version": "v3_multimodal",
             "filename": filename if filename else "audio_upload",
             "predicted_label": label_map[pred_idx],
             "confidence": round(confidence, 4),
             "modalities_used": ["text", "linguistic"] + (["audio"] if has_audio else []),
-            "generated_transcript": processed_text if not is_cha_provided else None
+            "generated_transcript": processed_text if not is_cha_provided else None,
+            "visualizations": visualizations
         }