| import os |
| import json |
| from pathlib import Path |
| import whisperx |
| import soundfile as sf |
| import numpy as np |
| import re |
| import torch |
| from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline |
| import sys |
|
|
| from dotenv import load_dotenv |
| load_dotenv() |
| token = os.getenv("HF_TOKEN") |
|
|
| print("Start Preprocessing ... ...") |
|
|
| sys.path.append('./CrisperWhisper/') |
| from utils import adjust_pauses_for_hf_pipeline_output |
|
|
| def generate_session_id(): |
| session_root = "session_data" |
| if not os.path.exists(session_root): |
| os.makedirs(session_root) |
| return "000001" |
| |
| existing_ids = [d for d in os.listdir(session_root) |
| if os.path.isdir(os.path.join(session_root, d)) and d.isdigit()] |
| if existing_ids: |
| new_id = max(int(x) for x in existing_ids) + 1 |
| else: |
| new_id = 1 |
| return f"{new_id:06d}" |
|
|
| def assign_speakers(segments, diarization_segments): |
| speaker_map = {} |
| for segment in segments: |
| segment_start = segment["start"] |
| segment_end = segment["end"] |
| max_overlap = 0 |
| assigned_speaker = "Unknown" |
| |
| for _, diar in diarization_segments.iterrows(): |
| speaker = diar["speaker"] |
| diar_start = diar["start"] |
| diar_end = diar["end"] |
| overlap_start = max(segment_start, diar_start) |
| overlap_end = min(segment_end, diar_end) |
| overlap_duration = max(0, overlap_end - overlap_start) |
| if overlap_duration > max_overlap: |
| max_overlap = overlap_duration |
| assigned_speaker = speaker |
| |
| speaker_map[segment_start] = assigned_speaker |
| return speaker_map |
|
|
| def load_audio_for_split(input_audio_file): |
|
|
| if input_audio_file.lower().endswith('.mp3'): |
| from pydub import AudioSegment |
| audio_seg = AudioSegment.from_file(input_audio_file) |
| sr = audio_seg.frame_rate |
| samples = np.array(audio_seg.get_array_of_samples()).astype(np.float32) |
| samples = samples / 32768.0 |
| if audio_seg.channels > 1: |
| samples = samples.reshape((-1, audio_seg.channels)) |
| return samples, sr |
| else: |
| return sf.read(input_audio_file) |
|
|
| def split_segment_by_sentences(segment): |
|
|
| text = segment["text"] |
| words = segment["words"] |
| start_time = segment["start"] |
| end_time = segment["end"] |
| speaker = segment["speaker"] |
| |
| sentences = [s.strip() for s in text.split('.') if s.strip()] |
| |
| if len(sentences) <= 1: |
| return [segment] |
| |
| new_segments = [] |
| word_index = 0 |
| |
| for i, sentence in enumerate(sentences): |
| if not sentence: |
| continue |
| |
| sentence_words = [] |
| sentence_text_clean = re.sub(r'[^\w\s]', '', sentence.lower()) |
| sentence_word_tokens = sentence_text_clean.split() |
| |
| matched_words = 0 |
| sentence_start = None |
| sentence_end = None |
| |
| temp_word_index = word_index |
| while temp_word_index < len(words) and matched_words < len(sentence_word_tokens): |
| word_obj = words[temp_word_index] |
| word_text_clean = re.sub(r'[^\w\s]', '', word_obj["word"].lower()) |
| |
| if word_text_clean == sentence_word_tokens[matched_words]: |
| if sentence_start is None: |
| sentence_start = word_obj["start"] |
| sentence_end = word_obj["end"] |
| sentence_words.append(word_obj) |
| matched_words += 1 |
| elif word_text_clean in sentence_word_tokens[matched_words:]: |
| sentence_words.append(word_obj) |
| if sentence_start is None: |
| sentence_start = word_obj["start"] |
| sentence_end = word_obj["end"] |
| |
| temp_word_index += 1 |
| |
| if sentence_start is None or sentence_end is None: |
| total_duration = end_time - start_time |
| sentence_duration = total_duration / len(sentences) |
| sentence_start = start_time + i * sentence_duration |
| sentence_end = start_time + (i + 1) * sentence_duration |
| |
| if i == len(sentences) - 1: |
| sentence_end = end_time |
| |
| word_index = temp_word_index |
| |
| new_segment = { |
| "start": round(sentence_start, 3), |
| "end": round(sentence_end, 3), |
| "speaker": speaker, |
| "text": sentence + ".", |
| "words": sentence_words |
| } |
| new_segments.append(new_segment) |
| |
| return new_segments |
|
|
| def process_audio_file(input_audio_file, num_speakers, device="cuda"): |
| |
| print("Loading WhisperX model (English)...") |
| model = whisperx.load_model("medium", device, language="en") |
| |
| audio = whisperx.load_audio(input_audio_file) |
| |
| print("Transcribing audio with WhisperX...") |
| result = model.transcribe(audio) |
| |
| print("Performing forced alignment with WhisperX...") |
| alignment_model, metadata = whisperx.load_align_model(language_code="en", device=device) |
| result_aligned = whisperx.align(result["segments"], alignment_model, metadata, audio, device, return_char_alignments=True) |
| |
| print("Detecting speakers with WhisperX...") |
| diarization_model = whisperx.DiarizationPipeline(use_auth_token=token, |
| device=device) |
| diarization_segments = diarization_model(audio) |
| |
| speaker_map = assign_speakers(result_aligned["segments"], diarization_segments) |
| for segment in result_aligned["segments"]: |
| segment["speaker"] = speaker_map.get(segment["start"], "Unknown") |
| segment.pop("chars", None) |
|
|
| session_id = generate_session_id() |
| session_dir = os.path.join("session_data", session_id) |
| os.makedirs(session_dir, exist_ok=True) |
| |
| data, sr = load_audio_for_split(input_audio_file) |
| |
| for segment in result_aligned["segments"]: |
| start_time = segment["start"] |
| end_time = segment["end"] |
| speaker = segment["speaker"] |
| start_sample = int(start_time * sr) |
| end_sample = int(end_time * sr) |
| segment_audio = data[start_sample:end_sample] |
| |
| segment_filename = f"{session_id}-{start_time:.2f}-{end_time:.2f}-{speaker}.wav" |
| segment_filepath = os.path.join(session_dir, segment_filename) |
| sf.write(segment_filepath, segment_audio, sr) |
| print(f"Saved segment: {segment_filepath}") |
| |
|
|
| transcript_path = os.path.join(session_dir, f"{session_id}_transcription.txt") |
| with open(transcript_path, "w", encoding="utf-8") as f: |
| for segment in result_aligned["segments"]: |
| f.write(f"[{segment['start']} - {segment['end']}] (Speaker {segment['speaker']}): {segment['text']}\n") |
| |
|
|
| del model |
| torch.cuda.empty_cache() |
| |
|
|
| print("Loading CrisperWhisper model...") |
| device_str = "cuda:0" if torch.cuda.is_available() else "cpu" |
| torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 |
|
|
| """ Use local Crisper Whisper Model |
| local_model_dir = "./CrisperWhisper_local" |
| |
| cw_model = AutoModelForSpeechSeq2Seq.from_pretrained( |
| local_model_dir, |
| torch_dtype=torch_dtype, |
| low_cpu_mem_usage=True, |
| use_safetensors=True |
| ) |
| cw_model.to(device_str) |
| |
| processor = AutoProcessor.from_pretrained(local_model_dir) |
| |
| """ |
| hf_model_id = "nyrahealth/CrisperWhisper" |
|
|
| cw_model = AutoModelForSpeechSeq2Seq.from_pretrained( |
| hf_model_id, |
| torch_dtype=torch_dtype, |
| low_cpu_mem_usage=True, |
| use_safetensors=True, |
| token=token |
| ) |
| cw_model.to(device_str) |
|
|
| processor = AutoProcessor.from_pretrained(hf_model_id, token=token) |
|
|
| |
| |
| asr_pipeline = pipeline( |
| "automatic-speech-recognition", |
| model=cw_model, |
| tokenizer=processor.tokenizer, |
| feature_extractor=processor.feature_extractor, |
| chunk_length_s=30, |
| batch_size=4, |
| return_timestamps='word', |
| torch_dtype=torch_dtype, |
| device=0 if torch.cuda.is_available() else -1, |
| generate_kwargs={"language": "en"} |
| ) |
| |
| segments_cw = [] |
| skipped_segments = [] |
| segment_files = [f for f in os.listdir(session_dir) if f.endswith('.wav')] |
| for seg_file in sorted(segment_files): |
|
|
| match = re.match(r'^(\d+)-(\d+\.\d+)-(\d+\.\d+)-(.+)\.wav$', seg_file) |
| if not match: |
| continue |
|
|
| seg_session_id = match.group(1) |
| start_time = float(match.group(2)) |
| end_time = float(match.group(3)) |
| speaker = match.group(4) |
| seg_path = os.path.join(session_dir, seg_file) |
| |
| print(f"Processing segment with CrisperWhisper: {seg_path}") |
| try: |
| cw_output = asr_pipeline(seg_path) |
| cw_result = adjust_pauses_for_hf_pipeline_output(cw_output) |
| except Exception as e: |
| print(f"[Warning] CrisperWhisper error, skiped this segment: {seg_path}\nError Message: {e}") |
| skipped_segments.append(seg_path) |
| continue |
| |
| text = cw_result.get('text', '').strip() |
| if not text: |
| print(f"********** No text returned, skiped this segment: {seg_path} **********") |
| skipped_segments.append(seg_path) |
| continue |
| |
| chunks = cw_result.get('chunks', []) |
| words_info = [] |
| for i, chunk in enumerate(chunks): |
| word_text = chunk['text'].strip() |
| if not word_text: |
| continue |
|
|
| chunk_start, chunk_end = chunk['timestamp'] |
| |
| if chunk_start is None: |
| if i == 0: |
| chunk_start = 0.0 |
| else: |
| chunk_start = words_info[-1]['end'] - start_time |
|
|
| if chunk_end is None: |
| if i < len(chunks) - 1: |
| next_chunk_start, _ = chunks[i+1]['timestamp'] |
| if next_chunk_start is None: |
| next_chunk_start = chunk_start |
| chunk_end = next_chunk_start |
| else: |
| chunk_end = end_time - start_time |
|
|
| word_start = round(start_time + chunk_start, 3) |
| word_end = round(start_time + chunk_end, 3) |
| words_info.append({ |
| "word": word_text, |
| "start": word_start, |
| "end": word_end |
| }) |
| |
| segment_entry = { |
| "start": round(start_time, 3), |
| "end": round(end_time, 3), |
| "speaker": speaker, |
| "text": text, |
| "words": words_info |
| } |
| |
| print(f"Post-processing: splitting segment by sentences...") |
| split_segments = split_segment_by_sentences(segment_entry) |
| segments_cw.extend(split_segments) |
| |
| segments_cw = sorted(segments_cw, key=lambda x: x["start"]) |
| cw_json_path = os.path.join(session_dir, f"{session_id}_transcriptionCW.json") |
| with open(cw_json_path, "w", encoding="utf-8") as f: |
| json.dump({"segments": segments_cw}, f, ensure_ascii=False, indent=4) |
| print(f"CrisperWhisper transcription saved to: {cw_json_path}") |
| |
| if skipped_segments: |
| skipped_file = os.path.join(session_dir, "skipped_segments.txt") |
| with open(skipped_file, "w", encoding="utf-8") as f: |
| for s in sorted(skipped_segments): |
| f.write(s + "\n") |
| print(f"Skipped segments recorded in: {skipped_file}") |
| |
| return session_id |
|
|
| if __name__ == "__main__": |
| session = process_audio_file("/home/easgrad/shuweiho/workspace/volen/SATE_docker_test/input/454.mp3", num_speakers=2, device="cuda") |
| print("Processing complete. Session ID:", session) |
