Update app.py

app.py

@@ -5,43 +5,46 @@ import torch
 import soundfile as sf
 from moviepy.editor import VideoFileClip, CompositeVideoClip, ImageClip
 from PIL import Image, ImageDraw, ImageFont
-
 from nemo.collections import asr as nemo_asr
-from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
 from ctc_segmentation import ctc_segmentation, CtcSegmentationParameters, prepare_text
 
 
 # =============================
-# LISTE DES MODELES ROBOTSMALI
+# LISTE OFFICIELLE DES MODELES ROBOTSMALI
 # =============================
 
 MODELS = {
-    "Soloni 114M TDT CTC v1": "RobotsMali/soloni-114m-tdt-ctc-v1",
-    "Soloni 350M TDT CTC v1": "RobotsMali/soloni-350m-tdt-ctc-v1",
+    "Soloni V0": "RobotsMali/soloni-114m-tdt-ctc-V0",
+    "Soloni V1": "RobotsMali/soloni-114m-tdt-ctc-V1",
 
-    "Soloba CTC 0.6B v0": "RobotsMali/soloba-ctc-0.6b-v0",
-    "Soloba CTC 0.6B v1": "RobotsMali/soloba-ctc-0.6b-v1",
+    "Soloba V0": "RobotsMali/soloba-ctc-0.6b-V0",
+    "Soloba V1": "RobotsMali/soloba-ctc-0.6b-V1",
 
-    "QuartzNet Bambara v1": "RobotsMali/stt-bm-quartznet15x5-v1",
-    "QuartzNet Bambara v2": "RobotsMali/stt-bm-quartznet15x5-v2"
+    "QuartzNet V0": "RobotsMali/stt-bm-quartznet15x5-V0",
+    "QuartzNet V1": "RobotsMali/stt-bm-quartznet15x5-V1"
 }
 
 
 # =============================
-# FONCTION : EXTRAIRE AUDIO
+# EXTRACTION AUDIO (FIABLE + COMPATIBLE HF & COLAB)
 # =============================
 
 def extract_audio(video_path, wav_path):
-    clip = VideoFileClip(video_path)
-    audio = clip.audio.to_soundarray(fps=16000)
-    if audio.ndim == 2:
-        audio = np.mean(audio, axis=1)
-    sf.write(wav_path, audio, 16000)
-    clip.close()
+    (
+        VideoFileClip(video_path)
+        .audio
+        .write_audiofile(
+            wav_path,
+            fps=16000,
+            codec="pcm_s16le",
+            verbose=False,
+            logger=None
+        )
+    )
 
 
 # =============================
-# FONCTION : TRANSCRIPTION + TIMESTAMP
+# TRANSCRIPTION + ALIGNEMENT
 # =============================
 
 def transcribe(model, device, wav, model_name):
@@ -51,18 +54,19 @@ def transcribe(model, device, wav, model_name):
 
     x = torch.tensor(audio, dtype=torch.float32).unsqueeze(0).to(device)
     ln = torch.tensor([x.shape[1]]).to(device)
+    total_s = len(audio) / sr
 
-    # === Cas 1 : Soloni → timestamps natifs ===
-    if "Soloni" in model_name and hasattr(model, "decode_and_align"):
+    # === Soloni → timestamps natifs ===
+    if "Soloni" in model_name:
         with torch.no_grad():
             proc, plen = model.preprocessor(input_signal=x, input_signal_length=ln)
             hyps = model.decode_and_align(encoder_output=proc, encoded_lengths=plen)
+
         hyp = hyps[0][0] if isinstance(hyps[0], list) else hyps[0]
         return [(w.start_offset_ms/1000, w.end_offset_ms/1000, w.word) for w in hyp.words]
 
-    # === Cas 2 : Soloba & QuartzNet → Forced Alignment CTC ===
-    text = model.transcribe([wav])[0]
-    text = text.strip()
+    # === Soloba / QuartzNet → Forced Alignment CTC ===
+    text = model.transcribe([wav])[0].strip()
     if not text:
         return []
 
@@ -72,37 +76,29 @@ def transcribe(model, device, wav, model_name):
     words = text.split()
     config = CtcSegmentationParameters()
     config.char_list = list(model.tokenizer.vocab.keys())
-    gt, utt = prepare_text(config, words)
+    gt, _ = prepare_text(config, words)
 
     timings, _, _ = ctc_segmentation(config, logits.cpu().numpy()[0], gt)
-    total_s = len(audio) / sr
     tps = total_s / logit_len.cpu().numpy()[0]
 
-    word_times = []
-    for i, w in enumerate(words):
-        s = timings[i] * tps
-        e = timings[i+1] * tps if i+1 < len(timings) else total_s
-        word_times.append((s, e, w))
-
-    # Groupage lisible : 3-5 mots par ligne
-    grouped, block = [], []
-    for w in word_times:
-        block.append(w)
-        if len(block) >= 4:
-            grouped.append(block)
-            block = []
-    if block:
-        grouped.append(block)
+    aligned = [(timings[i] * tps,
+                timings[i+1] * tps if i+1 < len(timings) else total_s,
+                words[i]) for i in range(len(words))]
 
-    subs = []
-    for g in grouped:
-        subs.append((g[0][0], g[-1][1], " ".join([w[2] for w in g])))
+    grouped, temp = [], []
+    for w in aligned:
+        temp.append(w)
+        if len(temp) >= 4:
+            grouped.append(temp)
+            temp = []
+    if temp:
+        grouped.append(temp)
 
-    return subs
+    return [(g[0][0], g[-1][1], " ".join([w[2] for w in g])) for g in grouped]
 
 
 # =============================
-# FONCTION : INCRUSTATION SOUS-TITRES
+# INCRUSTATION SOUS-TITRES (SANS IMAGEMAGICK)
 # =============================
 
 def burn(video, subs):
@@ -121,10 +117,7 @@ def burn(video, subs):
         bbox = draw.textbbox((0,0), text, font=font)
         tw, th = bbox[2]-bbox[0], bbox[3]-bbox[1]
         draw.text(((W-tw)//2, (int(H*0.12)-th)//2), text, font=font, fill="white")
-
-        layers.append(ImageClip(np.array(img))
-                      .set_start(s).set_duration(e-s)
-                      .set_position(("center", int(H*0.85))))
+        layers.append(ImageClip(np.array(img)).set_start(s).set_duration(e-s).set_position(("center", int(H*0.85))))
 
     final = CompositeVideoClip([clip] + layers)
     out = "RobotsMali_Subtitled.mp4"
@@ -145,22 +138,23 @@ def pipeline(video_file, model_name):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = nemo_asr.models.ASRModel.from_pretrained(MODELS[model_name]).to(device)
 
-    wav = "temp.wav"
+    wav = "audio.wav"
     extract_audio(video_file, wav)
     subs = transcribe(model, device, wav, model_name)
-    out = burn(video_file, subs)
-    return "✅ Sous-titres générés avec succès.", out
+    final = burn(video_file, subs)
+
+    return "✅ Sous-titres générés.", final
 
 
 # =============================
-# INTERFACE GRADIO
+# INTERFACE (inchangée)
 # =============================
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎙️ RobotsMali Subtitle Generator")
+    gr.Markdown("# 🎙️ **RobotsMali - Sous-titrage Bambara Automatique**")
 
-    video = gr.Video(label="Importer une vidéo")
-    model = gr.Dropdown(list(MODELS.keys()), value="Soloni 114M TDT CTC v1", label="Sélection du modèle")
+    video = gr.Video(label="Vidéo")
+    model = gr.Dropdown(list(MODELS.keys()), value="Soloni V1", label="Modèle")
     btn = gr.Button("⚡ Générer les sous-titres")
     status = gr.Markdown()
     out = gr.Video(label="Résultat")