Update app.py

app.py

@@ -1,16 +1,16 @@
 import os, warnings, tempfile
 warnings.filterwarnings("ignore")
 
-# Autoriser ImageMagick sur n'importe quelle distribution
-for path in ["/etc/ImageMagick/policy.xml", "/etc/ImageMagick-6/policy.xml"]:
-    if os.path.exists(path):
-        os.system(f'sed -i "s/rights=\\"none\\"/rights=\\"read|write\\"/g" {path}')
+# Débloquer ImageMagick (compatible tous environnements)
+for p in ["/etc/ImageMagick/policy.xml", "/etc/ImageMagick-6/policy.xml"]:
+    if os.path.exists(p):
+        os.system(f'sed -i "s/rights=\\"none\\"/rights=\\"read|write\\"/g" "{p}"')
 
 import gradio as gr
 import numpy as np
 import soundfile as sf
 import torch
-from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
+from moviepy.editor import VideoFileClip, CompositeVideoClip, TextClip
 from nemo.collections import asr as nemo_asr
 
 SR = 16000
@@ -26,98 +26,106 @@ ASR_MODELS = {
 
 _CACHE = {}
 
-def load_model(model_key):
-    if model_key in _CACHE:
-        return _CACHE[model_key]
+def load_model(name):
+    if name in _CACHE:
+        return _CACHE[name]
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = nemo_asr.models.ASRModel.from_pretrained(model_name=ASR_MODELS[model_key]).to(device).eval()
-    _CACHE[model_key] = (model, device)
+    model = nemo_asr.models.ASRModel.from_pretrained(model_name=ASR_MODELS[name]).to(device).eval()
+    _CACHE[name] = (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):
+    # Extraction audio stable (pas to_soundarray)
+    try:
+        clip = VideoFileClip(video_path)
+        clip.audio.write_audiofile(
+            wav_path, fps=SR, codec="pcm_s16le", verbose=False, logger=None
+        )
+        clip.close()
+    except:
+        os.system(f"ffmpeg -i '{video_path}' -ac 1 -ar {SR} -vn -y '{wav_path}' >/dev/null 2>&1")
     audio, sr = sf.read(wav_path)
-    if audio.ndim > 1:
-        audio = audio.mean(axis=1)
-    total_duration = len(audio) / sr
+    return len(audio)/sr
 
-    audio_t = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
-    length_t = torch.tensor([audio_t.shape[1]]).to(device)
+def transcribe(model, device, wav_path, model_key):
+    audio, sr = sf.read(wav_path)
+    if audio.ndim > 1: audio = audio.mean(axis=1)
+    total_s = len(audio)/sr
+
+    x = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
+    ln = torch.tensor([x.shape[1]]).to(device)
 
-    # === Soloni → timestamps exacts ===
+    # Timestamps Soloni (decode_and_align → hyp.words)
     if "Soloni" in model_key:
         with torch.no_grad():
-            proc, proc_len = model.preprocessor(audio_t, length_t)
-            hyps = model.decode_and_align(encoder_output=proc, encoded_lengths=proc_len)
+            proc, plen = model.preprocessor(x, ln)
+            hyps = model.decode_and_align(proc, plen)
         hyp = hyps[0][0] if isinstance(hyps[0], list) else hyps[0]
-        if hasattr(hyp, "words") and hyp.words:
+        if hasattr(hyp, "words"):
             return [(w.start_offset_ms/1000, w.end_offset_ms/1000, w.word) for w in hyp.words]
 
-    # === Soloba + QuartzNet → alignement fluide ===
+    # CTC models → alignement dynamique
     text = model.transcribe([wav_path])[0]
     words = text.split()
-    if not words:
-        return []
+    if not words: return []
 
-    wps = max(2.0, len(words)/total_duration)
+    wps = max(2.2, len(words)/total_s)
     subs, t = [], 0
     for w in words:
-        dur = 1.0 / wps
-        subs.append((t, min(total_duration, t+dur), w))
-        t += dur
-        if t >= total_duration: break
+        d = 1/wps
+        subs.append((t, min(total_s, t+d), w))
+        t += d
+        if t >= total_s: break
     return subs
 
-def burn_subtitles(video_path, subs):
+def burn(video_path, subs):
     clip = VideoFileClip(video_path)
     W, H = clip.size
-    overlays = [
-        TextClip(w.upper(), fontsize=int(H/20), color="white",
-                 stroke_color="black", stroke_width=2,
-                 method="caption", size=(int(W*0.9), None)
+
+    layers = []
+    for s,e,w in subs:
+        txt = TextClip(w.upper(), fontsize=int(H/20), color="white",
+                       stroke_color="black", stroke_width=2,
+                       method="caption", size=(int(W*0.9), None)
         ).set_start(s).set_duration(e-s).set_position(("center", int(H*0.88)))
-        for s, e, w in subs
-    ]
-    final = CompositeVideoClip([clip] + overlays)
-    out = "output_captioned.mp4"
+        layers.append(txt)
+
+    final = CompositeVideoClip([clip] + layers)
+    out = "RobotsMali_Subtitled.mp4"
     final.write_videofile(out, codec="libx264", audio_codec="aac", verbose=False, logger=None)
     return out
 
-def pipeline(video, model_key, progress=gr.Progress()):
+def pipeline(video, model_name, progress=gr.Progress()):
     progress(0.2, "Chargement du modèle…")
-    model, device = load_model(model_key)
+    model, device = load_model(model_name)
+
     with tempfile.TemporaryDirectory() as td:
         wav = f"{td}/audio.wav"
-        progress(0.4, "Extraction audio…")
+        progress(0.45, "Extraction audio…")
         extract_audio(video, wav)
+
         progress(0.7, "Transcription en Bambara…")
-        subs = transcribe_with_timestamps(model, device, wav, model_key)
+        subs = transcribe(model, device, wav, model_name)
+
     progress(0.9, "Incrustation des sous-titres…")
-    out = burn_subtitles(video, subs)
+    out = burn(video, subs)
+
     progress(1.0, "✅ Terminé")
-    return f"✅ Sous-titres générés avec **{model_key}**", out
+    return f"✅ Sous-titré avec **{model_name}**", out
 
 CSS = """
-body { background:#F5F8FC; font-family:Inter, sans-serif; }
-h1 { text-align:center; font-weight:800; color:#007BFF; }
-p { text-align:center; color:#5A6B85; margin-bottom:24px; }
-.gr-button { background:#007BFF !important; color:white !important; border-radius:8px; }
+body { background:#F6F9FF; font-family:Inter, sans-serif; }
+h1 { text-align:center; font-weight:800; color:#006CFF; }
+.gr-button { background:#007BFF !important; color:white !important; border-radius:8px; font-weight:700; }
 """
 
-with gr.Blocks(css=CSS) as demo:
+with gr.Blocks(css=CSS, title="RobotsMali Caption Studio") as demo:
     gr.Markdown("<h1>RobotsMali Caption Studio</h1><p>Sous-titrage Automatique en Bambara</p>")
     video = gr.File(label="🎥 Importer une vidéo", type="filepath")
     model = gr.Dropdown(list(ASR_MODELS.keys()), value="Soloni 114M TDT CTC V1")
-    btn = gr.Button("🚀 Générer")
+    run = gr.Button("🚀 Générer les sous-titres")
     status = gr.Markdown()
-    out_video = gr.Video()
-    btn.click(pipeline, inputs=[video, model], outputs=[status, out_video])
+    output = gr.Video()
+    run.click(pipeline, inputs=[video, model], outputs=[status, output])
 
 demo.launch()