Update app/services/processor.py

app/services/processor.py

@@ -72,32 +72,17 @@ class ProcessingResult:
     emotion_timeline: List[EmotionPoint] = None
     emotion_changes: List[EmotionChange] = None
 
-
-def pad_and_refine_tensor(
-    waveform: torch.Tensor,
-    sr: int,
-    start_s: float,
-    end_s: float,
-    pad_ms: int = 250,
-) -> Tuple[float, float]:
-
-    total_len = waveform.shape[1]
-    s = max(int((start_s - pad_ms / 1000) * sr), 0)
-    e = min(int((end_s + pad_ms / 1000) * sr), total_len)
-
-    if e <= s:
-        return start_s, end_s
-
-    return s / sr, e / sr
-
-
 def normalize_asr_result(result: dict):
 
     words = []
 
     for w in result.get("words", []):
 
-        word = w.get("word", "").strip()
+        word = (
+            w.get("word", "")
+            .strip()
+        )
+
         if not word:
             continue
 
@@ -107,12 +92,15 @@ def normalize_asr_result(result: dict):
                 "start": float(w["start"]),
                 "end": float(w["end"]),
                 "speaker": w.get("speaker"),
+                "confidence": float(
+                    w.get("confidence", 1.0)
+                ),
             }
         )
 
     text = result.get("text", "").strip()
-    return text, words
 
+    return text, words
 
 def guess_speaker_by_overlap(start, end, diar_segments):
 
@@ -163,65 +151,44 @@ def format_timestamp(seconds: float) -> str:
     return f"{m:02d}:{s:06.3f}"
 
 
-def extract_mfcc_segment(
-    audio: np.ndarray,
-    sr: int,
-    start: float,
-    end: float,
-    duration=5,
-):
-
-    start_sample = int(start * sr)
-    end_sample = int(end * sr)
-
-    segment = audio[start_sample:end_sample]
-
-    if len(segment) == 0:
-        return None
-
-    target_len = int(sr * duration)
-
-    if len(segment) < target_len:
-        segment = np.pad(segment, (0, target_len - len(segment)), mode="symmetric")
-    else:
-        segment = segment[:target_len]
-
-    mfcc = librosa.feature.mfcc(
-        y=segment, sr=sr, n_mfcc=128, n_fft=2048, hop_length=512
-    )
-
-    return mfcc
-
-
 def merge_consecutive_segments(
     segments: List[SpeakerSegment],
-    max_gap: float = 0.8,
-    min_duration: float = 0.15,
+    max_gap: float = 0.80,
+    max_overlap: float = 0.15,
 ) -> List[SpeakerSegment]:
-    """Merge consecutive segments from same speaker."""
+
     if not segments:
         return []
 
-    merged = []
-    current = SpeakerSegment(
-        start=segments[0].start, end=segments[0].end, speaker=segments[0].speaker
+    segments = sorted(
+        segments,
+        key=lambda x: x.start
     )
 
+    merged = [segments[0]]
+
     for seg in segments[1:]:
-        seg_dur = seg.end - seg.start
+
+        prev = merged[-1]
+
+        gap = seg.start - prev.end
+
+        overlap = prev.end - seg.start
+
         if (
-            seg.speaker == current.speaker
-            and (seg.start - current.end) <= max_gap
-            or seg_dur < min_duration
+            seg.speaker == prev.speaker
+            and gap <= max_gap
+            and overlap <= max_overlap
         ):
-            # Merge: extend current segment
-            current.end = seg.end
+
+            prev.end = max(
+                prev.end,
+                seg.end
+            )
+
         else:
-            # New speaker or gap too large
-            merged.append(current)
-            current = SpeakerSegment(start=seg.start, end=seg.end, speaker=seg.speaker)
+            merged.append(seg)
 
-    merged.append(current)
     return merged
 
 
@@ -242,8 +209,7 @@ class Processor:
         audio_path: Path,
         model_name: str = "PhoWhisper Lora Finetuned",
         language="vi",
-        merge_segments: bool = True,
-        backend: str = "whisper",
+        merge_segments: bool = True
     ) -> ProcessingResult:
 
         import asyncio
@@ -282,8 +248,10 @@ class Processor:
 
         diarization_segments = [
             SpeakerSegment(
-                *pad_and_refine_tensor(waveform, sr, s.start, s.end),
+                start=new_start,
+                end=new_end,
                 speaker=s.speaker,
+                confidence=getattr(s, "confidence", 1.0),
             )
             for s in diarization_segments
         ]
@@ -301,74 +269,29 @@ class Processor:
 
         speakers = list(speaker_map.values())
 
-        # 5. NORMALIZE ROLES
-        speakers = diarization.speakers or []
-        roles = diarization.roles or {}
+        raw_roles = diarization.roles or {}
+
+        roles = {}
 
-        # Default fallback
-        for label in speakers:
-            roles.setdefault(label, "KH")
+        for raw_spk, label in speaker_map.items():
+            roles[label] = raw_roles.get(raw_spk, "KH")
+        
 
         logger.info(f"roles(mapped) = {roles}")
 
         # 7: Transcribe segments after diarization
         logger.info("Step 7: Running ASR with external VAD batch...")
 
-        # asr_result = await TranscriptionService.transcribe_with_words_async(
-        #     audio_array=y,
-        #     model_name=model_name,
-        #     language=language,
-        #     vad_options=True
-        # )
-
-        # text, raw_words = normalize_asr_result(asr_result)
-
-        logger.info(f"ASR backend = {backend}")
-
-        if backend == "whisper":
-            asr_result = await TranscriptionService.transcribe_with_words_async(
-                audio_array=y,
-                model_name=model_name,
-                language=language,
-                vad_options=True,
-            )
-            text, raw_words = normalize_asr_result(asr_result)
-
-        elif backend == "w2v":
-
-            logger.info("Running W2V per diarization segment...")
-
-            raw_words = []
-
-            for seg in diarization_segments:
-
-                start_sample = int(seg.start * sr)
-                end_sample = int(seg.end * sr)
-
-                chunk = y[start_sample:end_sample]
-
-                if len(chunk) == 0:
-                    continue
-
-                text = TranscriptionService.infer_w2v(chunk)
-
-                if not text:
-                    continue
-
-                raw_words.append(
-                    {
-                        "word": text,
-                        "start": seg.start,
-                        "end": seg.end,
-                        "speaker": seg.speaker,
-                    }
-                )
-
-        else:
-            raise ValueError(f"Unsupported backend: {backend}")
+        asr_result = await TranscriptionService.transcribe_with_words_async(
+            audio_array=y,
+            model_name=model_name,
+            language=language,
+            vad_options=False
+        )
 
-        processed_segments: List[TranscriptSegment] = []
+        text, raw_words = normalize_asr_result(asr_result)
 
+    
         if not raw_words:
             processed_segments = [
                 TranscriptSegment(
@@ -400,33 +323,17 @@ class Processor:
                         start=w["start"],
                         end=w["end"],
                         speaker=spk,
+                        confidence=w.get("confidence", 1.0)
                     )
                 )
 
             word_objs.sort(key=lambda x: x.start)
 
             # ===== ALIGNMENT =====
-            # aligned_segments = AlignmentService.align_precision(
-            #     word_objs,
-            #     diarization_segments
-            # )
-            if backend == "w2v":
-                aligned_segments = [
-                    TranscriptSegment(
-                        start=w["start"],
-                        end=w["end"],
-                        speaker=speaker_map.get(w["speaker"], "Speaker 1"),
-                        role=roles.get(
-                            speaker_map.get(w["speaker"], "Speaker 1"), "KH"
-                        ),
-                        text=w["word"],
-                    )
-                    for w in raw_words
-                ]
-            else:
-                aligned_segments = AlignmentService.align_precision(
-                    word_objs, diarization_segments
-                )
+            aligned_segments = AlignmentService.align_precision(
+                word_objs,
+                diarization_segments
+            )
 
             processed_segments = []
 
@@ -463,7 +370,6 @@ class Processor:
                         )
                     )
 
-        processed_segments = cls._merge_adjacent_segments(processed_segments)
         processed_segments.sort(key=lambda x: x.start)
 
         # 8 : Predict emotion segments
@@ -497,43 +403,6 @@ class Processor:
             emotion_changes=emotion_changes,
         )
 
-    @staticmethod
-    def _merge_adjacent_segments(
-        segments: List[TranscriptSegment],
-        max_gap_s: float = 0.8,
-        max_segment_duration: float = 9.0,
-    ) -> List[TranscriptSegment]:
-        """
-        Merge adjacent segments if:
-        - same speaker
-        - gap <= max_gap_s
-        """
-        if not segments:
-            return segments
-
-        segments = sorted(segments, key=lambda s: s.start)
-        merged = [segments[0]]
-
-        for seg in segments[1:]:
-            prev = merged[-1]
-
-            gap = seg.start - prev.end
-            combined_duration = seg.end - prev.start
-
-            if (
-                seg.speaker == prev.speaker
-                and seg.role == prev.role
-                and gap <= max_gap_s
-                and combined_duration <= max_segment_duration
-                and not overlap_prefix(seg.text, prev.text)
-            ):
-                # MERGE
-                prev.text = f"{prev.text} {seg.text}".strip()
-                prev.end = max(prev.end, seg.end)
-            else:
-                merged.append(seg)
-
-        return merged
 
     @staticmethod
     def _predict_emotion_segments(