ml/predict.py

"""
Inferência com o modelo treinado — integração com o pipeline principal.

Uso standalone:
  python ml/predict.py video.mp4

Uso programático:
  from ml.predict import load_custom_model, predict_video
  model = load_custom_model("ml/model.joblib")
  result = predict_video(model, video_bytes, "video.mp4")
"""

import sys
import tempfile
from pathlib import Path

import numpy as np
import joblib

from ml.extract_features import extract_random_frames, compute_features


def load_custom_model(model_path: str = "ml/model.joblib") -> dict:
    """Carrega o modelo treinado."""
    data = joblib.load(model_path)
    print(f"[ML] Modelo carregado: {data['model_name']} ({len(data['feature_columns'])} features)")
    return data


def predict_frames(model_data: dict, frames: list[np.ndarray]) -> dict:
    """
    Classifica uma lista de frames.

    Returns:
        {
            "prediction": str,           # "Real", "IA", "CGI"
            "prediction_label": int,     # 0, 1, 2
            "confidence": float,         # 0-100
            "class_probabilities": dict, # {"Real": 0.85, "IA": 0.10, "CGI": 0.05}
            "per_frame": list[dict],     # predição por frame
        }
    """
    model = model_data["model"]
    feature_cols = model_data["feature_columns"]
    label_map = model_data["label_map"]

    # Extrai features de cada frame
    X_rows = []
    for frame in frames:
        features = compute_features(frame)
        row = [features[col] for col in feature_cols]
        X_rows.append(row)

    X = np.array(X_rows, dtype=np.float32)
    X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)

    # Predição por frame
    proba = model.predict_proba(X)  # shape: (n_frames, n_classes)
    preds = model.predict(X)

    per_frame = []
    for i in range(len(frames)):
        frame_probs = {label_map[j]: round(float(proba[i][j]) * 100, 2) for j in range(proba.shape[1])}
        per_frame.append({
            "frame_index": i,
            "prediction": label_map[int(preds[i])],
            "probabilities": frame_probs,
        })

    # Agregação: média das probabilidades
    avg_proba = np.mean(proba, axis=0)
    best_class = int(np.argmax(avg_proba))

    return {
        "prediction": label_map[best_class],
        "prediction_label": best_class,
        "confidence": round(float(avg_proba[best_class]) * 100, 2),
        "class_probabilities": {
            label_map[i]: round(float(avg_proba[i]) * 100, 2)
            for i in range(len(avg_proba))
        },
        "per_frame": per_frame,
    }


def predict_video(model_data: dict, video_bytes: bytes, filename: str, n_frames: int = 10) -> dict:
    """Pipeline completo: bytes → frames → features → predição."""
    suffix = Path(filename).suffix or ".mp4"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(video_bytes)
        tmp_path = tmp.name

    try:
        frames = extract_random_frames(tmp_path, n_frames)
        if not frames:
            raise ValueError("Nenhum frame extraído do vídeo.")
        return predict_frames(model_data, frames)
    finally:
        import os
        os.unlink(tmp_path)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Uso: python ml/predict.py <video.mp4> [model.joblib]")
        sys.exit(1)

    video_path = sys.argv[1]
    model_path = sys.argv[2] if len(sys.argv) > 2 else "ml/model.joblib"

    model_data = load_custom_model(model_path)
    frames = extract_random_frames(video_path, 10)
    result = predict_frames(model_data, frames)

    print(f"\nResultado: {result['prediction']} ({result['confidence']}%)")
    print(f"Probabilidades: {result['class_probabilities']}")