Update app.py

app.py

@@ -1,64 +1,183 @@
-import os
+import gradio as gr
+import numpy as np
 import torch
-import librosa
 import soundfile as sf
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
-import gradio as gr
+import librosa
+from matplotlib import pyplot as plt
+from transformers import AutoFeatureExtractor, AutoModelForAudioFrameClassification
+from recitations_segmenter import segment_recitations, clean_speech_intervals
+import io
+from PIL import Image
+import tempfile
+import os
+import zipfile
 
-MODEL_ID = "xLeonSTES/quran-to-text-base"
-SAMPLE_RATE = 16000
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# 🔹 ASR client to connect to Space B
+from gradio_client import Client, handle_file
 
-@torch.no_grad()
-def load_model():
-    processor = AutoProcessor.from_pretrained(MODEL_ID)
-    model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_ID)
-    model.to(DEVICE)
-    model.eval()
-    return processor, model
+# ======================
+# Setup device and model
+# ======================
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-processor, model = load_model()
+print(f"Loading segmentation model on {device}...")
+processor = AutoFeatureExtractor.from_pretrained("obadx/recitation-segmenter-v2")
+model = AutoModelForAudioFrameClassification.from_pretrained(
+    "obadx/recitation-segmenter-v2",
+    torch_dtype=dtype,
+    device_map=device
+)
+print("Segmentation model loaded successfully!")
 
-def resample_to_16k(path):
+# 🔹 ASR Space (Space B)
+asr_client = Client("aboalaa1472/Quran_ASR")  # لو Space B Private: pass hf_token="HF_xxx"
+
+# ======================
+# Utils
+# ======================
+def read_audio(path, sampling_rate=16000):
     audio, sr = sf.read(path)
-    if audio.ndim > 1:
+    if len(audio.shape) > 1:
         audio = audio.mean(axis=1)
-    if sr != SAMPLE_RATE:
-        audio = librosa.resample(audio.astype('float32'), orig_sr=sr, target_sr=SAMPLE_RATE)
-    return audio, SAMPLE_RATE
-
-def transcribe_audio(path):
-    audio, sr = resample_to_16k(path)
-    audio = audio / (max(abs(audio)) + 1e-9)
-    inputs = processor(audio, sampling_rate=SAMPLE_RATE, return_tensors="pt")
-    input_features = inputs.input_features.to(DEVICE)
-
-    with torch.no_grad():
-        generated_ids = model.generate(input_features)
-
-    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return text
-
-def run(uploaded_audio, mic_audio):
-    path = mic_audio or uploaded_audio
-    if not path:
-        return "No audio provided"
+    if sr != sampling_rate:
+        audio = librosa.resample(audio, orig_sr=sr, target_sr=sampling_rate)
+    return torch.tensor(audio).float()
+
+def get_interval(x, intervals, idx, sr=16000):
+    start = int(intervals[idx][0] * sr)
+    end = int(intervals[idx][1] * sr)
+    return x[start:end]
+
+def plot_signal(x, intervals, sr=16000):
+    fig, ax = plt.subplots(figsize=(20, 4))
+    if isinstance(x, torch.Tensor):
+        x = x.numpy()
+    ax.plot(x, linewidth=0.5)
+    for s, e in intervals:
+        ax.axvline(x=s * sr, color='red', alpha=0.4)
+        ax.axvline(x=e * sr, color='red', alpha=0.4)
+    plt.tight_layout()
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    img = Image.open(buf)
+    plt.close()
+    return img
+
+# ======================
+# Main processing
+# ======================
+def process_audio(audio_file, min_silence_ms, min_speech_ms, pad_ms):
+    if audio_file is None:
+        return None, "⚠️ ارفع ملف صوتي", None, []
+
     try:
-        return transcribe_audio(path)
+        wav = read_audio(audio_file)
+
+        sampled_outputs = segment_recitations(
+            [wav],
+            model,
+            processor,
+            device=device,
+            dtype=dtype,
+            batch_size=4,
+        )
+
+        clean_out = clean_speech_intervals(
+            sampled_outputs[0].speech_intervals,
+            sampled_outputs[0].is_complete,
+            min_silence_duration_ms=min_silence_ms,
+            min_speech_duration_ms=min_speech_ms,
+            pad_duration_ms=pad_ms,
+            return_seconds=True,
+        )
+
+        intervals = clean_out.clean_speech_intervals
+        plot_img = plot_signal(wav, intervals)
+
+        temp_dir = tempfile.mkdtemp()
+        segment_files = []
+        full_asr_text = []
+
+        result_text = f"✅ عدد المقاطع: {len(intervals)}\n\n"
+
+        for i in range(len(intervals)):
+            seg = get_interval(wav, intervals, i)
+            if isinstance(seg, torch.Tensor):
+                seg = seg.cpu().numpy()
+
+            seg_path = os.path.join(temp_dir, f"segment_{i+1:03d}.wav")
+            sf.write(seg_path, seg, 16000)
+            segment_files.append(seg_path)
+
+            # 🔹 ASR call to Space B
+            asr_text = asr_client.predict(
+                uploaded_audio=handle_file(seg_path),
+                mic_audio=handle_file(seg_path),
+                api_name="/run"
+            )
+
+            full_asr_text.append(asr_text)
+
+            result_text += (
+                f"🎵 مقطع {i+1} "
+                f"({intervals[i][0]:.2f}s → {intervals[i][1]:.2f}s)\n"
+                f"📜 {asr_text}\n\n"
+            )
+
+        result_text += "\n🧾 النص الكامل:\n"
+        result_text += " ".join(full_asr_text)
+
+        # ZIP
+        zip_path = os.path.join(temp_dir, "segments.zip")
+        with zipfile.ZipFile(zip_path, 'w') as zipf:
+            for f in segment_files:
+                zipf.write(f, os.path.basename(f))
+
+        return plot_img, result_text, zip_path, segment_files
+
     except Exception as e:
-        return f"Error: {e}"
+        return None, f"❌ خطأ: {str(e)}", None, []
 
-with gr.Blocks(title="Quran ASR") as demo:
-    gr.Markdown("# Quran ASR — Diacritized Transcription\nUpload or record audio, then press Convert.")
+# ======================
+# Gradio UI
+# ======================
+with gr.Blocks(title="Quran Segmentation + ASR") as demo:
+    gr.Markdown("## 🕌 تقطيع التلاوات + ASR (Quran Text)")
 
     with gr.Row():
         with gr.Column():
-            upload = gr.Audio(type="filepath", label="Upload Audio")
-            mic = gr.Audio(type="filepath", label="Microphone Recording")
-            btn = gr.Button("Convert")
+            audio_input = gr.Audio(type="filepath", label="📤 ارفع التلاوة")
+            min_silence = gr.Slider(10, 500, 30, step=10, label="Min Silence (ms)")
+            min_speech = gr.Slider(10, 500, 30, step=10, label="Min Speech (ms)")
+            padding = gr.Slider(0, 200, 30, step=10, label="Padding (ms)")
+            btn = gr.Button("🚀 ابدأ")
+
         with gr.Column():
-            out = gr.Textbox(label="Output Text", lines=10)
+            plot_out = gr.Image(label="📈 الإشارة")
+            text_out = gr.Textbox(lines=20, label="📜 النص")
+
+    zip_out = gr.File(label="📦 تحميل المقاطع")
+
+    segment_outputs = [gr.Audio(visible=False) for _ in range(50)]
+
+    def process_and_show(audio, ms, sp, pad):
+        plot, text, zipf, segments = process_audio(audio, ms, sp, pad)
+        outputs = [plot, text, zipf]
+        for i in range(50):
+            if i < len(segments):
+                outputs.append(gr.Audio(value=segments[i], visible=True))
+            else:
+                outputs.append(gr.Audio(visible=False))
+        return outputs
 
-    btn.click(run, inputs=[upload, mic], outputs=[out])
+    btn.click(
+        process_and_show,
+        inputs=[audio_input, min_silence, min_speech, padding],
+        outputs=[plot_out, text_out, zip_out] + segment_outputs
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()