Corrige salida de vídeo en Gradio usando streaming e OpenCV VideoWriter

app.py

@@ -2,45 +2,71 @@ import gradio as gr
 import tempfile
 import shutil
 import os
+import cv2
 from ultralytics import YOLO
 
-# Evita el error de OpenMP
-os.environ["OMP_NUM_THREADS"] = "1"
-
 # ────────────────────────────
-# Modelo
+# Configuración
 # ────────────────────────────
-model = YOLO("best.pt")  # se carga una vez al arrancar
-
-def infer(video):
-    """Corre la detección sobre el vídeo subido y devuelve el vídeo con boxes"""
-    tmp = tempfile.mkdtemp()
-    inp = os.path.join(tmp, "in.mp4")
-    shutil.copy(video, inp)
-    results = model.predict(source=inp, save=True, conf=0.25, iou=0.45)
-    # results es una lista; cogemos el primer objeto y su carpeta de guardado
-    save_dir = results[0].save_dir
-    return os.path.join(save_dir, "in.mp4")
+os.environ["OMP_NUM_THREADS"] = "1"  # para evitar warnings de OpenMP
+model = YOLO("best.pt")               # carga el modelo solo una vez
+
+def infer(video_path):
+    """
+    Procesa el vídeo en modo streaming, anota cada frame
+    y genera un nuevo MP4 con las detecciones.
+    """
+    # 1) Copiamos el input a un tmp para asegurar ruta persistente
+    tmpdir = tempfile.mkdtemp()
+    in_vid = os.path.join(tmpdir, "in.mp4")
+    shutil.copy(video_path, in_vid)
+
+    # 2) Preparamos el writer de OpenCV en None
+    out_vid = os.path.join(tmpdir, "out.mp4")
+    writer = None
+    fps = 30  # asume 30 FPS; ajústalo si tu vídeo es distinto
+
+    # 3) Streaming: riene un generador de resultados por frame
+    results = model.predict(source=in_vid, conf=0.25, iou=0.45, stream=True)
+
+    for r in results:
+        # r.orig_img es el frame original en BGR
+        frame = r.orig_img
+        # r.plot() devuelve una copia anotada del frame
+        annotated = r.plot()
+
+        # Inicializa el VideoWriter cuando conozcamos tamaño:
+        if writer is None:
+            h, w = annotated.shape[:2]
+            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+            writer = cv2.VideoWriter(out_vid, fourcc, fps, (w, h))
+
+        writer.write(annotated)
+
+    # 4) Cerramos el writer
+    if writer is not None:
+        writer.release()
+
+    return out_vid
 
 def show_classes():
-    """Devuelve las clases aprendidas por el modelo"""
+    """Devuelve las clases que el modelo conoce."""
     names = model.names
     if isinstance(names, dict):
-        # dict: iteramos por clave ordenada y recogemos los valores
-        class_list = [names[i] for i in sorted(names.keys())]
+        # ordenamos por clave
+        class_list = [names[k] for k in sorted(names)]
     else:
-        # lista de strings
         class_list = names
     return ", ".join(class_list)
 
 # ────────────────────────────
-# UI en Blocks
+# Interfaz Gradio
 # ────────────────────────────
 with gr.Blocks(title="Kesherat · Inspección de palas eólicas") as demo:
-    gr.Markdown("## Inspección de palas eólicas con YOLO")
+    gr.Markdown("## Inspección de palas eólicas con YOLOv8")
 
-    video_in = gr.Video(label="Sube tu video de inspección")
-    video_out = gr.Video(label="Video con defectos detectados")
+    video_in = gr.Video(label="Sube tu vídeo de inspección")
+    video_out = gr.Video(label="Vídeo anotado")
 
     btn_detect = gr.Button("Detectar defectos")
     btn_detect.click(fn=infer, inputs=video_in, outputs=video_out)
@@ -50,7 +76,4 @@ with gr.Blocks(title="Kesherat · Inspección de palas eólicas") as demo:
     btn_classes.click(fn=show_classes, outputs=txt_classes)
 
 if __name__ == "__main__":
-    demo.launch(
-        share=True,
-        allowed_paths=["/home/user/.pyenv/runs/detect/predict"]
-    )
+    demo.launch()