Create app.py

Initial release of the RobotsMali Caption Studio.
Adds full Bambara speech-to-text video subtitling pipeline using the 6 official ASR models (Soloni, Soloba, QuartzNet). Includes timestamp alignment, hard subtitle rendering and Gradio web interface.

app.py

@@ -0,0 +1,130 @@
+import os, warnings, tempfile
+warnings.filterwarnings("ignore")
+
+# Autoriser ImageMagick
+os.system('sed -i "s/rights=\\"none\\"/rights=\\"read|write\\"/g" /etc/ImageMagick-6/policy.xml || true')
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
+from nemo.collections import asr as nemo_asr
+from nemo.collections.asr.parts.ds_cluster import convert_rnnt_to_word_timestamps
+
+SR = 16000
+
+ASR_MODELS = {
+    "Soloba CTC 0.6B V0": "RobotsMali/soloba-ctc-0.6b-v0",
+    "Soloba CTC 0.6B V1": "RobotsMali/soloba-ctc-0.6b-v1",
+    "Soloni 114M TDT CTC V0": "RobotsMali/soloni-114m-tdt-ctc-V0",
+    "Soloni 114M TDT CTC V1": "RobotsMali/soloni-114m-tdt-ctc-v1",
+    "QuartzNet BM V0": "RobotsMali/stt-bm-quartznet15x5-V0",
+    "QuartzNet BM V1": "RobotsMali/stt-bm-quartznet15x5-V1",
+}
+
+_MODEL_CACHE = {}
+
+def load_asr(model_key):
+    if model_key in _MODEL_CACHE:
+        return _MODEL_CACHE[model_key]
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = nemo_asr.models.ASRModel.from_pretrained(model_name=ASR_MODELS[model_key])
+    model = model.to(device).eval()
+    _MODEL_CACHE[model_key] = (model, device)
+    return model, device
+
+def extract_audio(video_path, wav_path):
+    with VideoFileClip(video_path) as clip:
+        audio = clip.audio.to_soundarray(fps=SR)
+    if audio.ndim > 1:
+        audio = audio.mean(axis=1)
+    sf.write(wav_path, audio.astype(np.float32), SR)
+    return len(audio) / SR
+
+def transcribe_with_timestamps(model, device, wav_path, model_key):
+    audio, sr = sf.read(wav_path)
+    if audio.ndim > 1:
+        audio = audio.mean(axis=1)
+    total_duration = len(audio) / sr
+    audio_t = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
+    audio_len = torch.tensor([audio_t.shape[1]]).to(device)
+
+    # Soloni → alignement natif
+    if "Soloni" in model_key:
+        with torch.no_grad():
+            proc, proc_len = model.preprocessor(audio_t, audio_len)
+            hyps = model.decode_and_align(encoder_output=proc, encoded_lengths=proc_len)
+        hyp = hyps[0][0] if isinstance(hyps[0], list) else hyps[0]
+        words = convert_rnnt_to_word_timestamps(hyp, model.tokenizer)
+        return [(w.start_time, w.end_time, w.word) for w in words]
+
+    # Soloba / QuartzNet → alignement à débit contrôlé
+    text = model.transcribe([wav_path], batch_size=1)[0]
+    words = text.split()
+    if len(words) == 0:
+        return []
+    wps = max(1.9, len(words) / total_duration)
+    t = 0
+    timeline = []
+    for w in words:
+        dur = 1.0 / wps
+        start, end = t, min(total_duration, t + dur)
+        timeline.append((start, end, w))
+        t = end
+        if t >= total_duration: break
+    return timeline
+
+def burn_subtitles(video_path, subs):
+    clip = VideoFileClip(video_path)
+    W, H = clip.size
+    textclips = []
+    for start, end, word in subs:
+        tc = (TextClip(word.upper(), fontsize=int(H/20), color="white",
+                       stroke_color="black", stroke_width=2, method="caption",
+                       size=(int(W*0.9), None))
+                .set_start(start)
+                .set_duration(end-start)
+                .set_position(("center", int(H*0.88))))
+        textclips.append(tc)
+    final = CompositeVideoClip([clip] + textclips)
+    out = "output_captioned.mp4"
+    final.write_videofile(out, codec="libx264", audio_codec="aac", verbose=False, logger=None)
+    return out
+
+def pipeline(video, model_key, progress=gr.Progress()):
+    progress(0.2, "🧠 Chargement du modèle…")
+    model, device = load_asr(model_key)
+
+    with tempfile.TemporaryDirectory() as td:
+        wav_path = f"{td}/audio.wav"
+        progress(0.4, "🎧 Extraction audio…")
+        _ = extract_audio(video, wav_path)
+
+        progress(0.65, "✍🏾 Transcription Bambara…")
+        subs = transcribe_with_timestamps(model, device, wav_path, model_key)
+
+    progress(0.85, "🎞️ Incrustation des sous-titres…")
+    out = burn_subtitles(video, subs)
+
+    progress(1.0, "✅ Terminé.")
+    return f"✅ Sous-titres générés avec **{model_key}**", out
+
+CSS = """
+body { background:#F5F8FC; font-family:Inter, sans-serif; color:#222; }
+.gradio-container { max-width:880px; margin:auto; background:white; border-radius:16px; padding:28px; box-shadow:0 8px 24px rgba(0,0,0,0.08)}
+h1 { text-align:center; color:#007BFF; font-weight:800; }
+p { text-align:center; color:#5A6B85; }
+.gr-button { background:#007BFF !important; color:white !important; font-weight:700; border-radius:8px; }
+"""
+
+with gr.Blocks(title="RobotsMali Caption Studio", css=CSS) as demo:
+    gr.Markdown("<h1>RobotsMali Caption Studio</h1><p>Sous-titrage Bambara Automatique</p><hr>")
+    video_in = gr.File(label="🎥 Importer une vidéo", type="filepath")
+    model_sel = gr.Dropdown(list(ASR_MODELS.keys()), value="Soloni 114M TDT CTC V1", label="🧠 Modèle ASR")
+    btn = gr.Button("🚀 Générer les sous-titres")
+    status = gr.Markdown()
+    video_out = gr.Video()
+    btn.click(pipeline, inputs=[video_in, model_sel], outputs=[status, video_out])
+
+demo.launch()