vision_tools.py

# vision_tools.py
# -----------------------------------------------------------------------------
# Veureu — VISION utilities (self-contained)
#  - Image processing and analysis
#  - Object detection and recognition
#  - Face detection and recognition
#  - Scene description
#  - Montage sequence analysis
# -----------------------------------------------------------------------------
from __future__ import annotations


import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import json
import logging
import math
import os
import shlex
import subprocess

import numpy as np
import torch
import torchaudio
import torchaudio.transforms as T
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from pyannote.audio import Pipeline as PyannotePipeline
from speechbrain.pretrained import SpeakerRecognition
from pydub import AudioSegment
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from scenedetect import VideoManager, SceneManager
from scenedetect.detectors import ContentDetector

import os, base64, requests, subprocess, contextlib, time

from transformers import AutoProcessor, LlavaForConditionalGeneration
from PIL import Image

from libs.audio_tools_ana_2 import process_audio_for_video

import cv2

try:
    import face_recognition  # type: ignore
except Exception:
    face_recognition = None  # type: ignore

try:
    # Utility wrapper for DeepFace
    from libs.face_utils import FaceRecognizer as DFRecognizer
except Exception:
    try:
        from face_utils import FaceRecognizer as DFRecognizer
    except Exception:
        DFRecognizer = None  # type: ignore

try:
    from deepface import DeepFace
except ImportError:
    DeepFace = None

import easyocr

# -------------------------------- Logging ------------------------------------
log = logging.getLogger("audio_tools")
if not log.handlers:
    h = logging.StreamHandler()
    h.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
    log.addHandler(h)
log.setLevel(logging.INFO)

# ============================ UTILS ===========================================
def load_config(path: str = "configs/config_veureu.yaml") -> Dict[str, Any]:
    p = Path(path)
    if not p.exists():
        log.warning("Config file not found: %s (using defaults)", path)
        return {}
    try:
        import yaml
        cfg = yaml.safe_load(p.read_text(encoding="utf-8")) or {}
        cfg["__path__"] = str(p)
        return cfg
    except Exception as e:
        log.error("Failed to read YAML config: %s", e)
        return {}

# ---------------------------- IMAGE EMBEDDING ----------------------------------
class FaceOfImageEmbedding:
    """Preferred backend: `face_recognition`; fallback: DeepFace via libs.face_utils."""
    def __init__(self, deepface_model: str = 'Facenet512'):
        self.use_fr = face_recognition is not None
        self.df = None
        if not self.use_fr and DFRecognizer is not None:
            try:
                self.df = DFRecognizer(model_name=deepface_model)
                log.info("Using DeepFace (%s) as face embedding backend.", deepface_model)
            except Exception as e:
                log.warning("Failed to initialize DeepFace: %s", e)
        elif self.use_fr:
            log.info("Using face_recognition as face embedding backend.")
        else:
            log.error("No face embedding backend available.")

    def encode_image(self, image_path: Path) -> Optional[List[float]]:
        import numpy as np
        try:
            if self.use_fr:
                img = face_recognition.load_image_file(str(image_path))  # type: ignore
                encs = face_recognition.face_encodings(img)
                if encs:
                    # Normalizar cada embedding a norma 1
                    embeddings = [(e / np.linalg.norm(e)).astype(float).tolist() for e in encs]
                    return embeddings
                return None

            if self.df is not None:
                emb = self.df.get_face_embedding_from_path(str(image_path))
                if emb is None:
                    return None
                # Convertir a numpy array y normalizar
                emb = np.array(emb, dtype=float)
                emb = emb / np.linalg.norm(emb)
                return emb.tolist()

        except Exception as e:
            log.debug("Fallo embedding cara %s: %s", image_path, e)

        return None

class FaceAnalyzer:
    """Wrapper sencillo para DeepFace que obtiene edad y género de una imagen."""
    def __init__(self, actions=None):
        if actions is None:
            actions = ["age", "gender"]
        self.actions = actions

    def analyze_image(self, img_path: str) -> Optional[Dict[str, Any]]:
        try:
            result = DeepFace.analyze(img_path=img_path, actions=self.actions)

            # Si DeepFace devuelve una lista (varias caras), tomamos la primera
            if isinstance(result, list) and len(result) > 0:
                result = result[0]

            # Ahora sí podemos acceder a 'age' y 'dominant_gender'
            return {
                "age": result.get("age", "unknown"),
                "gender": result.get("dominant_gender", "unknown")
            }

        except Exception as e:
            log.warning("No se pudo analizar la imagen %s: %s", img_path, e)
            return None

# ----------------------------------- FUNCTIONS -------------------------------------
def map_identities_per_second(frames_per_second, intervals):
    for seg in intervals:
        seg_start = seg["start"]
        seg_end = seg["end"]

        # recolectar identidades de los frames en el rango del segmento
        identities = []
        for f in frames_per_second:
            if seg_start <= f["start"] <= seg_end:
                for face in f.get("faces", []):
                    identities.append(face)

        # contar apariciones
        seg["counts"] = dict(Counter(identities))

    return intervals

def _split_montage(img: np.ndarray, n: int, cfg: Dict[str, Any]) -> List[np.ndarray]:
    vd = cfg.get('vision_describer', {})
    montage_cfg = vd.get('montage', {})
    mode = montage_cfg.get('split_mode', 'horizontal')  # 'horizontal'|'vertical'|'grid'

    h, w = img.shape[:2]
    tiles: List[np.ndarray] = []

    if mode == 'vertical':
        tile_h = h // n
        for i in range(n):
            y0 = i * tile_h; y1 = h if i == n-1 else (i+1) * tile_h
            tiles.append(img[y0:y1, 0:w])
        return tiles

    if mode == 'grid':
        rows = int(montage_cfg.get('rows', 1) or 1)
        cols = int(montage_cfg.get('cols', n) or n)
        assert rows * cols >= n, "grid rows*cols must be >= n"
        tile_h = h // rows; tile_w = w // cols
        k = 0
        for r in range(rows):
            for c in range(cols):
                if k >= n: break
                y0, y1 = r*tile_h, h if (r==rows-1) else (r+1)*tile_h
                x0, x1 = c*tile_w, w if (c==cols-1) else (c+1)*tile_w
                tiles.append(img[y0:y1, x0:x1]); k += 1
        return tiles

    tile_w = w // n
    for i in range(n):
        x0 = i * tile_w; x1 = w if i == n-1 else (i+1) * tile_w
        tiles.append(img[0:h, x0:x1])
    return tiles

def generar_montage(frame_paths: List[str], output_dir: str) -> None:
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    montage_path = ""

    if frame_paths:
        imgs = [cv2.imread(kf) for kf in frame_paths if os.path.exists(kf)]
        imgs = [img for img in imgs if img is not None]
        print(f"Se encontraron {len(imgs)} imágenes para el montaje.")

        if imgs:
            h = max(img.shape[0] for img in imgs)  # altura máxima
            imgs_resized = [cv2.resize(img, (int(img.shape[1]*h/img.shape[0]), h)) for img in imgs]
            montage = cv2.hconcat(imgs_resized)
            montage_path = os.path.join(output_dir, "keyframes_montage.jpg")
            print(f"Guardando montaje en: {montage_path}")
            cv2.imwrite(montage_path, montage)
            print("Montaje guardado.")
        else:
            print("No se encontraron imágenes válidas para el montaje.")

    return montage_path

def describe_montage_sequence(
    montage_path: str,
    n: int,
    informacion,
    face_identities,
    *,
    config_path: str = 'configs/config_veureu.yaml'
) -> Dict[str, Any]:
    """Describe each sub-image of a montage.

    Inputs
    ------
    montage_path: str
        Path to a composite image made of n sub-ismages placed sequentially.
    n: int
        Number of sub-images to split and describe.
    config_path: str
        Path to YAML with 'vision_describer' configuration (provider and params).
    
    Returns
    -------
    list []: with the descripcion of each image
    """
    
    path_model = "BSC-LT/salamandra-7b-vision"

    processor = AutoProcessor.from_pretrained(path_model)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    model = LlavaForConditionalGeneration.from_pretrained(
        path_model,
        torch_dtype=dtype,
        low_cpu_mem_usage=True
    ).to(device)

    img = cv2.imread(montage_path, cv2.IMREAD_COLOR)
    if img is None:
        raise RuntimeError(f"No se puede leer la imagen: {montage_path}")

    cfg = load_config(config_path)
    tiles = _split_montage(img, n, cfg)
    if len(tiles) < n:
        raise RuntimeError(f"Se produjeron {len(tiles)} tiles, se esperaban {n}")

    # Convertir cada tile a PIL Image
    tile_images = [Image.fromarray(cv2.cvtColor(t, cv2.COLOR_BGR2RGB)) for t in tiles]

    sys_prompt = (
        "Ets un expert en narrativa visual. "
        "Descriu la imatge de manera molt breu i senzilla, en català, "
        "explicant només l’acció principal que s’hi veu. "
        "Respon amb una sola frase curta (10–20 paraules com a màxim), "
        "sense afegir detalls innecessaris ni descriure l’entorn."
    )

    all_results = []

    for i in range(len(tile_images)):
        batch = [tile_images[i]]  # lista con un solo tile

        conversation = [
            {"role": "system", "content": sys_prompt},
            {"role": "user", "content": [
                {"type": "image", "image": batch[0]},
                {"type": "text", "text": (
                    f"Descriu la imatge de manera molt breu i senzilla, en català. ")}
            ]}
        ]
        prompt_batch = processor.apply_chat_template(conversation, add_generation_prompt=True)
        
        inputs = processor(images=batch, text=prompt_batch, return_tensors="pt")
        for k, v in inputs.items():
            if v.dtype.is_floating_point:
                inputs[k] = v.to(device, dtype)
            else:
                inputs[k] = v.to(device)
        
        output = model.generate(**inputs, max_new_tokens=1024)
        text = processor.decode(output[0], skip_special_tokens=True)
        lines = text.split("\n")

        desc = ""
        for i, line in enumerate(lines):
            if line.lower().startswith(" assistant"):
                desc = "\n".join(lines[i+1:]).strip()
                break

        all_results.append(desc)
        torch.cuda.empty_cache()

    return all_results

# --------------------------- IMAGES EXTRACTION -----------------------------
def keyframe_conditional_extraction_ana(
    video_path, 
    output_dir, 
    threshold=30.0, 
    offset_frames=10
):
    """
    Detecta cambios de escena en un vídeo, guarda un fotograma por cada cambio,
    devuelve intervalos con start y end basados en los tiempos de los keyframes
    y genera un montaje con todos los keyframes.
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    video_manager = VideoManager([video_path])
    scene_manager = SceneManager()
    scene_manager.add_detector(ContentDetector(threshold=threshold))

    video_manager.start()
    scene_manager.detect_scenes(video_manager)

    scene_list = scene_manager.get_scene_list()

    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    video_duration = total_frames / fps

    keyframes = []
    for i, (start_time, end_time) in enumerate(scene_list):
        frame_number = int(start_time.get_frames()) + offset_frames
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, frame = cap.read()
        if ret:
            ts = frame_number / fps
            frame_path = os.path.join(output_dir, f"scene_{i+1:03d}.jpg")
            cv2.imwrite(frame_path, frame)
            keyframes.append({
                "index": i+1,
                "time": round(ts, 2),
                "path": frame_path
            })

    cap.release()
    video_manager.release()

    # Construimos intervalos con start y end
    intervals = []
    for i, kf in enumerate(keyframes):
        start = kf["time"]
        if i < len(keyframes) - 1:
            end = keyframes[i+1]["time"]
        else:
            end = video_duration  # última escena hasta el final
        intervals.append({
            "index": kf["index"],
            "start": start,
            "end": round(end, 2),
            "path": kf["path"]
        })

    return intervals

def keyframe_every_second(
    video_path: str,
    output_dir: str = ".",
    max_frames: Optional[int] = 10000,
) -> List[dict]:
    """
    Extrae un fotograma por cada segundo del video.

    Returns:
        List[dict]: Cada elemento es {"index", "start", "end", "path"}
    """
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    
    cap = cv2.VideoCapture(str(video_path))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps

    frames: List[dict] = []
    idx = 0
    sec = 0.0

    while sec <= duration:
        frame_number = int(sec * fps)
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, frame = cap.read()
        if not ret:
            break

        timestamp = frame_number / fps  
        frame_path = out_dir / f"frame_per_second{idx:03d}.jpg"
        cv2.imwrite(str(frame_path), frame)

        frames.append({
            "index": idx + 1,
            "start": round(timestamp, 2),
            "end": None,   # lo completamos después
            "path": str(frame_path),
        })
        
        idx += 1
        sec += 1.0  
        
        if max_frames and idx >= max_frames:
            break

    cap.release()

    # Completar los "end" con el inicio del siguiente frame
    for i in range(len(frames)):
        if i < len(frames) - 1:
            frames[i]["end"] = frames[i+1]["start"]
        else:
            frames[i]["end"] = round(duration, 2)

    return frames

from collections import Counter, defaultdict

# --------------------------- FRAMES PROCESSING -----------------------------
def process_frames(
    frames: List[dict],  # cada elemento es {"index", "start", "end", "path"}
    config: dict,
    face_col=None,
    embedding_model=None,
) -> Tuple[List[dict], List[int]]:
    """
    Procesa keyframes:
    - Detecta caras
    - Genera embeddings con FaceEmbedding
    - Opcionalmente compara con face_col (KNN top-3)
    - Opcionalmente ejecuta OCR
    """

    frame_results = []

    # Crear embedding_model si no se pasa
    if embedding_model is None:
        from libs.vision_tools import FaceOfImageEmbedding
        embedding_model = FaceOfImageEmbedding()

    for idx, frame in enumerate(frames):
        frame_path = frame["path"]

        try:
            raw_faces = embedding_model.encode_image(Path(frame_path))
        except Exception as e:
            print(f"Error procesando {frame_path}: {e}")
            raw_faces = None

        faces = []
        if raw_faces is not None:
            if isinstance(raw_faces[0], list):  # múltiples
                for e in raw_faces:
                    faces.append({"embedding": e})
            else:  # uno solo
                faces.append({"embedding": raw_faces})

        faces_detected = []
        for f in faces:
            embedding = f.get("embedding")
            identity = "Unknown"
            knn = []

            if face_col is not None and embedding is not None:
                try:
                    num_embeddings = face_col.count()
                    if num_embeddings < 1:
                        knn = []
                        identity = "Unknown"

                    else:
                        n_results = min(3, num_embeddings)
                        q = face_col.query(
                            query_embeddings=[embedding],
                            n_results=n_results,
                            include=["metadatas", "distances"]
                        )

                        knn = []
                        metas = q.get("metadatas", [[]])[0]
                        dists = q.get("distances", [[]])[0]
                        for meta, dist in zip(metas, dists):
                            person_id = meta.get("identity", "Unknown") if isinstance(meta, dict) else "Unknown"
                            knn.append({"identity": person_id, "distance": float(dist)})

                        if knn and knn[0]["distance"] < 0.6:
                            identity = knn[0]["identity"]
                        else:
                            identity = "Unknown"

                except Exception as e:
                    print(f"Face KNN failed: {e}")
                    knn = []
                    identity = "Unknown"

            faces_detected.append(identity)

        use_easyocr = True
        if use_easyocr:
            try:
                reader = easyocr.Reader(['en', 'es'], gpu=True)  # Cambiar gpu=False si no hay GPU
                results = reader.readtext(frame_path)
                ocr_text_easyocr = " ".join([text for _, text, _ in results]).strip()

            except Exception as e:
                print(f"OCR error: {e}")

        frame_results.append({
            "id": frame["index"],
            "start": frame["start"],
            "end": frame["end"],
            "image_path": frame_path,
            "faces": faces_detected,
            "ocr": ocr_text_easyocr,
        })

    return frame_results

if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser(description="Veureu — Audio tools (self-contained)")
    ap.add_argument("--video", required=True)
    ap.add_argument("--out", default="results")
    ap.add_argument("--config", default="configs/config_veureu.yaml")
    args = ap.parse_args()

    # Lightweight config loader (only for sample run)
    import yaml
    cfg = {}
    p = Path(args.config)
    if p.exists():
        cfg = yaml.safe_load(p.read_text(encoding="utf-8")) or {}

    out_dir = Path(args.out) / Path(args.video).stem
    out_dir.mkdir(parents=True, exist_ok=True)

    segs, srt = process_audio_for_video(args.video, out_dir, cfg, voice_collection=None)
    print(json.dumps({
        "segments": len(segs),
        "srt": srt
    }, indent=2, ensure_ascii=False))