FLUJOS/POCS/BACK_BACK/IMAGENES/image_analyzer.py

"""
image_analyzer.py
-----------------
Analiza imágenes usando un VLM via ollama y extrae:
  - tema, subtema
  - keywords (lista de palabras clave)
  - descripción
  - entidades (personas, organizaciones, lugares)

Modelos recomendados en ollama para visión:
  - llava:13b          (bueno, ligero)
  - qwen2-vl:7b        (muy bueno para keywords)
  - minicpm-v:8b       (rápido y preciso)
  - llava-llama3:8b    (balance velocidad/calidad)

Para usar Qwen3.5 GGUF (texto), primero genera caption con un VLM
y luego pasa el texto a Qwen para enriquecer keywords — ver pipeline_pruebas.py

Uso:
    analyzer = ImageAnalyzer(model="qwen2-vl:7b")
    result = analyzer.analyze("foto.jpg")
    results = analyzer.analyze_folder("./mis_imagenes/")
"""

import base64
import json
import os
import re
from datetime import datetime
from pathlib import Path

import requests

# ── Configuración ──────────────────────────────────────────────────────────────

OLLAMA_URL  = os.getenv("OLLAMA_URL", "http://localhost:11434")
DEFAULT_MODEL = os.getenv("VISION_MODEL", "qwen2-vl:7b")

SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}

KEYWORD_PROMPT = """Analiza esta imagen en detalle.
Devuelve ÚNICAMENTE un objeto JSON válido con esta estructura exacta, sin texto adicional:

{
  "tema": "tema principal de la imagen (1-3 palabras en español)",
  "subtema": "subtema específico (1-4 palabras en español)",
  "keywords": ["palabra1", "palabra2", "palabra3"],
  "descripcion": "descripción breve y objetiva de lo que muestra la imagen (1-2 frases)",
  "entidades": ["nombre_propio1", "organizacion1", "lugar1"],
  "idioma_detectado": "es/en/fr/..."
}

Requisitos:
- keywords: entre 8 y 15 palabras clave relevantes, en minúsculas
- entidades: solo si son claramente visibles/identificables, puede estar vacío []
- todo el contenido en español salvo entidades propias
- SOLO el JSON, sin markdown ni explicaciones"""


# ── Clase principal ─────────────────────────────────────────────────────────────

class ImageAnalyzer:
    def __init__(self, model: str = DEFAULT_MODEL, ollama_url: str = OLLAMA_URL):
        self.model = model
        self.ollama_url = ollama_url.rstrip("/")

    # ── Helpers ────────────────────────────────────────────────────────────────

    @staticmethod
    def encode_image(image_path: str) -> str:
        with open(image_path, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")

    def _check_ollama(self) -> bool:
        try:
            r = requests.get(f"{self.ollama_url}/api/tags", timeout=5)
            return r.status_code == 200
        except Exception:
            return False

    def _parse_json_response(self, raw: str) -> dict:
        """Extrae JSON del response aunque venga con texto alrededor."""
        raw = raw.strip()
        # Buscar bloque JSON entre llaves
        match = re.search(r'\{[\s\S]*\}', raw)
        if match:
            return json.loads(match.group())
        raise ValueError(f"No se encontró JSON válido en el response:\n{raw[:300]}")

    # ── Análisis de una imagen ─────────────────────────────────────────────────

    def analyze(self, image_path: str, extra_context: str = "") -> dict:
        """
        Analiza una imagen y devuelve dict con keywords y metadata.

        Args:
            image_path: ruta a la imagen
            extra_context: contexto adicional para el prompt (ej: "esta imagen es de un periódico sobre guerras")

        Returns:
            dict con: archivo, tema, subtema, texto, keywords, entidades,
                      source_type, fecha, image_path
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Imagen no encontrada: {image_path}")

        prompt = KEYWORD_PROMPT
        if extra_context:
            prompt = f"Contexto adicional: {extra_context}\n\n" + prompt

        img_b64 = self.encode_image(image_path)

        payload = {
            "model": self.model,
            "prompt": prompt,
            "images": [img_b64],
            "stream": False,
            "format": "json",
            "options": {
                "temperature": 0.1,   # baja temperatura = más determinista
                "num_predict": 512,
            }
        }

        print(f"  → Enviando a ollama ({self.model}): {Path(image_path).name}")
        response = requests.post(
            f"{self.ollama_url}/api/generate",
            json=payload,
            timeout=180
        )
        response.raise_for_status()

        raw_response = response.json().get("response", "")
        parsed = self._parse_json_response(raw_response)

        return {
            "archivo":       Path(image_path).name,
            "image_path":    str(Path(image_path).resolve()),
            "tema":          parsed.get("tema", "sin_clasificar").lower(),
            "subtema":       parsed.get("subtema", "").lower(),
            "texto":         parsed.get("descripcion", ""),
            "keywords":      [k.lower().strip() for k in parsed.get("keywords", [])],
            "entidades":     parsed.get("entidades", []),
            "idioma":        parsed.get("idioma_detectado", "es"),
            "source_type":   "imagen",
            "fecha":         datetime.now().strftime("%Y-%m-%d"),
            "modelo_usado":  self.model,
        }

    # ── Análisis de una carpeta ────────────────────────────────────────────────

    def analyze_folder(self, folder_path: str, extra_context: str = "") -> list[dict]:
        """
        Analiza todas las imágenes de una carpeta.

        Returns:
            Lista de resultados, incluye errores como dicts con campo 'error'
        """
        folder = Path(folder_path)
        if not folder.exists():
            raise FileNotFoundError(f"Carpeta no encontrada: {folder_path}")

        images = [p for p in folder.iterdir() if p.suffix.lower() in SUPPORTED_EXTENSIONS]
        print(f"\n[ImageAnalyzer] {len(images)} imágenes encontradas en {folder_path}")

        if not self._check_ollama():
            print(f"  ⚠  ollama no responde en {self.ollama_url}")
            print("     Instala ollama: https://ollama.ai")
            print(f"     Luego: ollama pull {self.model}")
            return []

        results = []
        for i, img_path in enumerate(images, 1):
            print(f"  [{i}/{len(images)}] {img_path.name}")
            try:
                result = self.analyze(str(img_path), extra_context)
                results.append(result)
                print(f"       tema={result['tema']} | keywords={result['keywords'][:4]}...")
            except Exception as e:
                print(f"       ERROR: {e}")
                results.append({
                    "archivo": img_path.name,
                    "error": str(e),
                    "source_type": "imagen",
                    "fecha": datetime.now().strftime("%Y-%m-%d"),
                })

        print(f"\n[ImageAnalyzer] Completado: {len([r for r in results if 'error' not in r])}/{len(images)} OK\n")
        return results

    # ── Guardar resultados a JSON ──────────────────────────────────────────────

    @staticmethod
    def save_json(results: list[dict], output_path: str):
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"[ImageAnalyzer] Guardado: {output_path} ({len(results)} registros)")