Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import json | |
| import numpy as np | |
| import os | |
| from datetime import datetime | |
| from model import Image2Phoneme | |
| from utils import ctc_post_process, audio_to_mel, mel_to_image, text_to_phonemes | |
| import soundfile as sf | |
| import shutil | |
| import time | |
| # Configuration | |
| DEVICE = torch.device("cpu") | |
| PHMAP = "phoneme_to_id.json" | |
| AUDIO_DIR = "audio_inputs" | |
| # Ensure audio directory exists | |
| os.makedirs(AUDIO_DIR, exist_ok=True) | |
| # Load phoneme vocabulary | |
| try: | |
| vocab = json.load(open(PHMAP, "r")) | |
| id_to_ph = {v: k for k, v in vocab.items()} | |
| except FileNotFoundError: | |
| raise FileNotFoundError(f"Phoneme mapping file not found at {PHMAP}") | |
| # Build model | |
| vocab_size = max(vocab.values()) + 1 | |
| model = Image2Phoneme(vocab_size=vocab_size).to(DEVICE) | |
| try: | |
| ckpt = torch.load("last_checkpoint.pt", map_location=DEVICE, weights_only=True) | |
| model.load_state_dict(ckpt["model_state_dict"]) | |
| model.eval() | |
| except FileNotFoundError: | |
| raise FileNotFoundError(f"Checkpoint file not found at last_checkpoint.pt") | |
| def process_audio(audio_input): | |
| """Process audio to predict phonemes and display mel spectrogram.""" | |
| try: | |
| print(f"Received audio_input before processing: {audio_input}") | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| audio_path = os.path.join(AUDIO_DIR, f"input_{timestamp}.wav") | |
| if audio_input is None: | |
| print("Audio input is None") | |
| return {"error": "No audio input provided"}, None, None, None | |
| if isinstance(audio_input, str): | |
| print(f"Processing uploaded file: {audio_input}") | |
| if not os.path.exists(audio_input): | |
| return {"error": f"Uploaded file not found: {audio_input}"}, None, None, None | |
| if audio_input.endswith(".mp3"): | |
| print("Converting .mp3 to .wav") | |
| from pydub import AudioSegment | |
| audio = AudioSegment.from_mp3(audio_input) | |
| audio_path = audio_path.replace(".wav", "_converted.wav") | |
| audio.export(audio_path, format="wav") | |
| print(f"Converted file saved to: {audio_path}") | |
| else: | |
| shutil.copy(audio_input, audio_path) | |
| print(f"Copied file to: {audio_path}") | |
| else: | |
| raise ValueError("Microphone input not supported in this configuration") | |
| mel_path = audio_to_mel(audio_path) | |
| print(f"Generated mel spectrogram: {mel_path}") | |
| if not os.path.exists(mel_path): | |
| return {"error": f"Mel spectrogram file not found: {mel_path}"}, None, None, None | |
| mel_image_path = mel_to_image(mel_path) | |
| print(f"Generated mel spectrogram image: {mel_image_path}") | |
| if not os.path.exists(mel_image_path): | |
| return {"error": f"Mel spectrogram image not found: {mel_image_path}"}, None, None, None | |
| mel = np.load(mel_path) | |
| print(f"Loaded mel spectrogram shape: {mel.shape}") | |
| mel_tensor = torch.tensor(mel).unsqueeze(0).to(DEVICE) | |
| mel_lens = torch.tensor([mel.shape[1]]).to(DEVICE) | |
| with torch.no_grad(): | |
| ph_pred = model(mel_tensor) | |
| ph_ids = ph_pred.argmax(-1)[0].cpu().numpy() | |
| print(f"Predicted phoneme IDs: {ph_ids}") | |
| ph_seq = [id_to_ph[i] for i in ph_ids if i > 0] | |
| print(f"Raw phonemes: {ph_seq}") | |
| post_processed = ctc_post_process(ph_seq) | |
| print(f"Post-processed phonemes: {post_processed}") | |
| return { | |
| "audio_path": audio_path, | |
| "phonemes": " ".join(ph_seq), | |
| "post_processed_phonemes": " ".join(post_processed) | |
| }, mel_image_path, " ".join(ph_seq), " ".join(post_processed) | |
| except Exception as e: | |
| print(f"Error in process_audio: {str(e)}") | |
| return {"error": f"Processing failed: {str(e)}"}, None, None, None | |
| # Gradio interface | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Speech to Phonemes Converter") | |
| gr.Markdown("Upload audio to predict phonemes and display mel spectrogram. Enter text to convert to phonemes.") | |
| audio_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio (.wav or .mp3)", interactive=True) | |
| text_input = gr.Textbox(label="Enter Text", placeholder="Type a sentence to convert to phonemes") | |
| process_button = gr.Button("Process") | |
| audio_path_output = gr.Textbox(label="Audio Path") | |
| mel_image = gr.Image(label="Mel Spectrogram", type="filepath") | |
| raw_phonemes_audio = gr.Textbox(label="Raw Phonemes (Audio)") | |
| post_processed_phonemes_audio = gr.Textbox(label="Post-Processed Phonemes (Audio)") | |
| text_phonemes_output = gr.Textbox(label="Phonemes (Text)") | |
| def process(audio_input, text_input): | |
| print(f"Processing inputs - Audio: {audio_input}, Text: {text_input}") | |
| if audio_input: | |
| audio_result, mel_image_path, raw_ph_audio, post_ph_audio = process_audio(audio_input) | |
| audio_path = audio_result.get("audio_path", "") | |
| raw_ph_audio = audio_result.get("phonemes", "") | |
| post_ph_audio = audio_result.get("post_processed_phonemes", "") | |
| error_msg = audio_result.get("error", "") | |
| if error_msg: | |
| audio_path = error_msg | |
| raw_ph_audio = "" | |
| post_ph_audio = "" | |
| mel_image_path = None | |
| else: | |
| audio_path = "" | |
| raw_ph_audio = "" | |
| post_ph_audio = "" | |
| mel_image_path = None | |
| text_result = text_to_phonemes(text_input) if text_input and text_input.strip() else {} | |
| text_phonemes = "".join(text_result) | |
| return audio_path, mel_image_path, raw_ph_audio, post_ph_audio, text_phonemes | |
| process_button.click( | |
| fn=process, | |
| inputs=[audio_input, text_input], | |
| outputs=[audio_path_output, mel_image, raw_phonemes_audio, post_processed_phonemes_audio, text_phonemes_output] | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(debug=True) |