Upload main_router.py

main_process/main_router.py

@@ -0,0 +1,211 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+from pathlib import Path
+from typing import Counter,List, Dict
+import ast
+import json
+import torch
+from svision_client import extract_scenes, add_ocr_and_faces, keyframes_every_second_extraction, extract_descripcion_escena
+from asr_client import extract_audio_from_video, diarize_audio, transcribe_long_audio, transcribe_short_audio, identificar_veu
+
+ROOT = Path("/tmp/engine")
+ROOT.mkdir(parents=True, exist_ok=True)
+
+# Aqui te explico lo que necesitaría
+    # Ruta del video: llamandola como video_path
+    # Ruta del json que contiene la información de los castings: llamándola como castings_path
+
+    # El castings_path es un json con la siguiente estructura:
+    # {"face_col": [{"nombre": "Ana", "embedding": [...]}, ...],
+    #  "voice_col": [{"nombre": "Ana", "embedding": [...]}, ...]}
+
+    # Importante: los embeddings de face_col son de dimensión 512 y los de voice_col de dimensión 192 (para obtenerlos simplemente es llamar al endpoint creado en svision y asr respectivamente).
+
+# Entonces aquí cargarías el json y extraerías las colecciones de caras y voces:
+face_col = [{"nombre": "Ana", "embedding": [0.1]*512}]
+voice_col = [{"nombre": "Ana", "embedding": [0.2]*192}]
+
+# Y ya lo demás es cómo lo tengo.
+
+# Al final hay dos path que es lo que luego necesitaría para hacer luego la versión o de Salamandra o la mixture_of_experts:
+    # - El path del srt inicial generado
+    # - El path del json con toda la información generada
+
+transcripcion_inicial = "/home/acasado/bsc/hugging_face_bsc/engine/results/transcription_initial.srt"
+informacion_json = "/home/acasado/bsc/hugging_face_bsc/engine/results/informacion.json"
+
+def map_identities_per_second(frames_per_second, intervals):
+    for seg in intervals:
+        seg_start = seg["start"]
+        seg_end = seg["end"]
+
+        identities = []
+        for f in frames_per_second:
+            if seg_start <= f["start"] <= seg_end:
+                for face in f.get("faces", []):
+                    identities.append(face)
+
+        seg["counts"] = dict(Counter(identities))
+
+    return intervals
+
+def _fmt_srt_time(seconds: float) -> str:
+    """Formatea segundos en el formato SRT HH:MM:SS,mmm"""
+    h = int(seconds // 3600)
+    m = int((seconds % 3600) // 60)
+    s = int(seconds % 60)
+    ms = int((seconds - int(seconds)) * 1000)
+    return f"{h:02}:{m:02}:{s:02},{ms:03}"
+
+def generate_srt_from_segments(segments: List[Dict], output_srt_path: str):
+    """
+    Genera un archivo SRT a partir de segmentos de diarización.
+
+    Args:
+        segments (List[Dict]): Lista de diccionarios con 'start', 'end', 'speaker' y 'transcription'.
+        output_srt_path (str): Ruta donde se guardará el archivo SRT.
+    """
+    out = Path(output_srt_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    with out.open("w", encoding="utf-8-sig") as f:
+        for i, seg in enumerate(segments, 1):
+            start = seg.get("start", 0.0)
+            end = seg.get("end", 0.0)
+            transcription = seg.get("transcription", "").strip()
+            speaker_info = seg.get("speaker", {})
+            speaker = speaker_info.get("identity", "Desconegut")
+
+            text = f"[{speaker}]: {transcription}" if speaker else transcription
+
+            f.write(f"{i}\n")
+            f.write(f"{_fmt_srt_time(start)} --> {_fmt_srt_time(end)}\n")
+            f.write(f"{text}\n\n")
+
+def pipeline_preprocessing_vision(video_path: str):
+    """
+    Pipeline que toma un video y realiza todo el preprocesamiento del video de la parte de vision.
+    """
+
+    print(f"Procesando video para visión: {video_path}")
+
+    print("Extrayendo escenas...")
+    threshold: float = 30.0
+    offset_frames: int = 3
+    crop_ratio: float = 0.1
+    result_extract_scenes = extract_scenes(video_path, threshold, offset_frames, crop_ratio)
+    print(result_extract_scenes)
+    # Obtener las rutas de las imágenes y la información de las escenas
+    escenas = result_extract_scenes[0] if len(result_extract_scenes) > 0 else []
+    escenas_paths = [f["image"] for f in escenas]
+    print(escenas_paths)
+    info_escenas = result_extract_scenes[1] if len(result_extract_scenes) > 1 else []
+    print(info_escenas)
+
+    print("Extrayendo imagenes por segundo...")
+    result_extract_per_second = keyframes_every_second_extraction(video_path)
+    # Obtener las rutas de las imágenes y la información de las escenas
+    images_per_second = result_extract_per_second[0] if len(result_extract_per_second) > 0 else []
+    images_per_second_paths = [f["image"] for f in images_per_second]
+    info_images_per_second = result_extract_per_second[1] if len(result_extract_per_second) > 1 else []
+
+    print("Aumentamos la información de las escenas viendo quién aparece en cada escena y detectando OCR...")
+    info_escenas_completa = []
+    for imagen_escena, info_escena in zip(escenas_paths, info_escenas):
+        result_add_ocr_and_faces = add_ocr_and_faces(imagen_escena, info_escena, face_col)
+        info_escenas_completa.append(result_add_ocr_and_faces)
+    
+    print("Aumentamos la información de las imagenes por segundo viendo quién aparece en cada escena y detectando OCR...")
+    info_images_per_second_completa = []
+    for imagen_segundo, info_segundo in zip(images_per_second_paths, info_images_per_second):
+        result_add_ocr_and_faces =add_ocr_and_faces(imagen_segundo, info_segundo, face_col)
+        info_images_per_second_completa.append(result_add_ocr_and_faces)
+    print(info_escenas_completa)
+    
+    print("Ahora se va a tratar los OCR (se sustituirán ciertas escenas por alguna de las imágenes por segundo si tienen mejor OCR)...")
+    # Se hará lo último
+
+    print("Combinando información de escenas e imágenes por segundo...")
+    info_escenas_completa = map_identities_per_second(info_images_per_second_completa, info_escenas_completa)
+    print(info_escenas_completa)
+
+    print("Ahora se incluyen en los diccionarios de las escenas la descripciones de estas.")
+    for escena_path, info_escena in zip(escenas_paths, info_escenas_completa):
+        descripcion_escena = extract_descripcion_escena(escena_path)
+        lista = ast.literal_eval(descripcion_escena)
+        frase = lista[0]
+        info_escena["descripcion"] = frase
+        del descripcion_escena
+        torch.cuda.empty_cache()
+            
+    return info_escenas_completa, info_images_per_second_completa
+
+def pipeline_preprocessing_audio(video_path: str):
+    """
+    Pipeline que toma un video y realiza todo el preprocesamiento del video de la parte de audio.
+    """
+    print(f"Procesando video para audio: {video_path}")
+
+    print("Extrayendo audio del video...")
+    audio_video = extract_audio_from_video(video_path)
+    print(audio_video)
+
+    print("Diartizando el audio...")
+    diarization_audio = diarize_audio(audio_video)
+    print(diarization_audio)
+    clips_path = diarization_audio[0]      
+    print(clips_path)
+    diarization_info = diarization_audio[1] 
+    print(diarization_info)
+
+    print("Transcribiendo el video completo...")
+    full_transcription = transcribe_long_audio(audio_video)
+    print(full_transcription)
+
+    print("Transcribiendo los clips diartizados...")
+    for clip_path, clip_info in zip(clips_path, diarization_info):
+        clip_transcription = transcribe_short_audio(clip_path)
+        clip_info["transcription"] = clip_transcription
+
+    print("Calculando los embeddings para cada uno de los clips obtenidos y posteriormente identificar las voces...")
+    for clip_path, clip_info in zip(clips_path, diarization_info):
+        clip_speaker = identificar_veu(clip_path, voice_col)
+        clip_info["speaker"] = clip_speaker
+
+    return full_transcription, diarization_info
+
+def pipeline_video_analysis(video_path: str):
+    """
+    Pipeline que toma un video y realiza todo el preprocesamiento del video (visión y audio).
+    """
+
+    print(f"Procesando video completo: {video_path}")
+    
+    # Procesamiento de visión
+    info_escenas, info_images_per_second = pipeline_preprocessing_vision(video_path)
+    torch.cuda.empty_cache()
+
+    # Procesamiento de audio
+    full_transcription, info_clips = pipeline_preprocessing_audio(video_path)
+
+    info_clips = map_identities_per_second(info_images_per_second, info_clips)
+
+    generate_srt_from_segments(info_clips, transcripcion_inicial)
+
+    output_json_path = Path(informacion_json)
+    output_json_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_json_path.open("w", encoding="utf-8") as f:
+        json.dump({
+            "full_transcription": full_transcription,
+            "info_escenas": info_escenas,
+            "info_clips": info_clips
+        }, f, ensure_ascii=False, indent=4)
+
+    print("Pipeline de preprocesamiento completo finalizado.")
+    torch.cuda.empty_cache()
+    print(info_escenas)
+    print(info_images_per_second)
+    print(info_clips)
+
+    return {}