update_speaker_id_to_json

app.py

@@ -11,6 +11,7 @@ from segmentation import reorganize_transcription_c_unit
 from annotation import annotate_maze_for_mazewhisper
 from morpheme import stanza_v1
 from morpheme import annotate_morpheme
+from speaker import assign_speaker
 
 
 
@@ -43,13 +44,17 @@ def process_audio():
     result, session_id = translate_audio_file(model="mazeWhisper", audio_path = audio_path, device=device, original_filename=filename)
 
     cunit_count, ignored_count = reorganize_transcription_c_unit(session_id, segment_batchalign)
-    print(f"Created {cunit_count} C-units, ignored {ignored_count} boundaries")
+    # print(f"Created {cunit_count} C-units, ignored {ignored_count} boundaries")
 
+    print("Processing speaker identification ... ...")
+    assign_speaker(session_id = session_id)
+
+    print("Processing maze detection ... ...")
     annotate_maze_for_mazewhisper(session_id)
 
+    print("Processing morpheme detection ... ...")
     annotate_morpheme(session_id = session_id, morpheme_function = stanza_v1)
 
-
     # annotate_pauses(session_id, pause_threshold)
     # annotate_repetitions(session_id)
     # # annotate_syllables(session_id)
@@ -58,8 +63,6 @@ def process_audio():
     # annotate_morpheme(session_id)
     # annotate_morpheme_omission(session_id)
 
-    
-
     json_path = f"session_data/{session_id}/transcription_cunit.json"
     if not os.path.isfile(json_path):
         return jsonify({'error': f"Annotation file {json_path} not found"}), 500

requirements.txt

@@ -21,5 +21,3 @@ matplotlib>=3.3.0
 seaborn>=0.11.0
 
 # install ffmpeg
-librosa>=0.8.0
-transformers>=4.0.0

speaker/speaker_identification.py

@@ -1,6 +1,6 @@
 from typing import List, Union, Optional
 import os
-
+import json
 import numpy as np
 import librosa
 from transformers import pipeline
@@ -21,6 +21,7 @@ def get_predictor():
     if _PREDICTOR_INSTANCE is None:
         _PREDICTOR_INSTANCE = Predictor()
     return _PREDICTOR_INSTANCE
+
 class Predictor:
     def __init__(self, model_path: Optional[str] = None):
         """
@@ -100,7 +101,7 @@ def assign_speaker_for_audio_list(audio_list: List[Union[str, np.ndarray]]) -> L
         
     Returns:
         List[str]: List of speaker IDs corresponding to each audio segment.
-                 "Speaker_id_0" for child, "Speaker_id_1" for adult.
+                 "Child" for child, "Examiner" for adult.
     """
     if not audio_list:
         return []
@@ -111,12 +112,85 @@ def assign_speaker_for_audio_list(audio_list: List[Union[str, np.ndarray]]) -> L
     # Get list of 0 (child) or 1 (adult)
     numeric_labels = predictor.predict(audio_list)
     
-    # Map to Speaker_id_0 and Speaker_id_1, preserving order
-    speaker_ids = [f"Speaker_id_{label}" if label in (0,1) else "Unknown" for label in numeric_labels]
+    # Map to Child and Examiner, preserving order
+    speaker_ids = ["Child" if label == 0 else "Examiner" if label == 1 else "Unknown" for label in numeric_labels]
     return speaker_ids
 
 
 # you don't have to implement this function
 def assign_speaker(session_id: str):
 
-    return
+    base_dir = os.path.join("session_data", session_id)
+    json_path = os.path.join(base_dir, "transcription_cunit.json")
+    wav_path = os.path.join(base_dir, "audio.wav")
+
+    with open(json_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    segments = data.get("segments", [])
+
+    if not segments:
+        return
+
+    audio, sr = librosa.load(wav_path, sr=DEFAULT_SAMPLE_RATE, mono=True)
+    n_samples = len(audio)
+    dur_sec = n_samples / float(DEFAULT_SAMPLE_RATE)
+
+    model_inputs: List[np.ndarray] = []
+    model_indices: List[int] = []
+    prefilled_unknown: List[int] = []
+
+    for i, seg in enumerate(segments):
+        start = seg.get("start")
+        end = seg.get("end")
+
+        if (
+            start is None or end is None
+            or not isinstance(start, (int, float))
+            or not isinstance(end, (int, float))
+            or end <= start
+            or start >= dur_sec
+        ):
+            prefilled_unknown.append(i)
+            continue
+
+        s = max(0.0, float(start))
+        e = min(float(end), dur_sec)
+
+        if e <= s:
+            prefilled_unknown.append(i)
+            continue
+
+        s_idx = int(round(s * DEFAULT_SAMPLE_RATE))
+        e_idx = int(round(e * DEFAULT_SAMPLE_RATE))
+
+        s_idx = max(0, min(s_idx, n_samples))
+        e_idx = max(0, min(e_idx, n_samples))
+
+        if e_idx <= s_idx:
+            prefilled_unknown.append(i)
+            continue
+
+        snippet = audio[s_idx:e_idx]
+
+        if snippet.size == 0:
+            prefilled_unknown.append(i)
+            continue
+
+        model_inputs.append(snippet)
+        model_indices.append(i)
+
+    speakers = ["Unknown"] * len(segments)
+    if model_inputs:
+        predicted = assign_speaker_for_audio_list(model_inputs)  # ["Child"/"Examiner"/"Unknown"]
+        for seg_idx, spk in zip(model_indices, predicted):
+            speakers[seg_idx] = spk
+
+    for seg_idx in prefilled_unknown:
+        speakers[seg_idx] = "Unknown"
+
+    for i, seg in enumerate(segments):
+        seg["speaker"] = speakers[i]
+
+
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)

transcription/transcription.py

@@ -298,6 +298,11 @@ def translate_audio_file(model: str = "mazeWhisper", audio_path: str = "", devic
         
         audio = load_audio(audio_path)
         
+        # Save the entire audio as audio.wav in the session directory
+        audio_output_path = session_dir / "audio.wav"
+        sf.write(audio_output_path, audio, SAMPLE_RATE)
+        print(f"Audio saved: {audio_output_path}")
+        
         print("Starting transcription...")
         result = pipeline.transcribe(audio_path, verbose=True)