predict_fusion.py

# -*- coding: utf-8 -*-
"""predict_fusion.py

CLI inference script for the Hybrid AI Image Forensics system.
It loads the persisted XGBoost model and StandardScaler, extracts the nine
scalar forensic scores from a single image, and prints a polished report.
"""

import os
import argparse
import joblib
import pandas as pd
import torch
import cv2
import numpy as np
from sklearn.preprocessing import StandardScaler

# Core components – already in the repo
from core.alignment import GeometricAligner
from core.diffusion_latent import DiffusionErrorLoop
from core.statistical_extraction import StatisticalFeatureExtractor

# --------------------------------------------------------------
# Helper functions – same as used in the FastAPI service
# --------------------------------------------------------------
def high_freq_energy(freq_tensor: torch.Tensor) -> float:
    freq = freq_tensor.squeeze().cpu()
    h, w = freq.shape
    mask = torch.zeros_like(freq, dtype=torch.bool)
    margin_h = h // 4
    margin_w = w // 4
    mask[:margin_h, :] = True
    mask[-margin_h:, :] = True
    mask[:, :margin_w] = True
    mask[:, -margin_w:] = True
    return float(torch.sum(torch.abs(freq)[mask]))

def compute_entropy(img_rgb):
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    from skimage.measure import shannon_entropy
    return float(shannon_entropy(gray))

def compute_edge_density(img_rgb):
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 100, 200)
    return float(np.sum(edges > 0) / edges.size)

def compute_laplacian_variance(img_rgb):
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    lap = cv2.Laplacian(gray, cv2.CV_64F)
    return float(lap.var())

def compute_color_kurtosis(img_rgb):
    from scipy.stats import kurtosis
    k = [kurtosis(img_rgb[..., c].ravel()) for c in range(3)]
    return float(np.mean(k))

def compute_jpeg_consistency(img_rgb):
    # Approximate JPEG consistency by computing DCT of the whole grayscale image
    # and measuring variance of the high‑frequency coefficients (excluding low‑freq 8×8 corner).
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY).astype(np.float32)
    # Apply 2‑D DCT to the entire image
    dct_full = cv2.dct(gray)
    # Mask out low‑frequency top‑left 8×8 region
    mask = np.ones_like(dct_full, dtype=bool)
    mask[:8, :8] = False
    # Return variance of remaining coefficients as proxy
    return float(np.var(dct_full[mask]))

# --------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Run forensic inference on a single image")
    parser.add_argument("--image", "-i", required=True, help="Path to image file")
    parser.add_argument("--model", default=os.path.join('dataset', 'fusion_engine_best.json'), help="Path to saved XGBoost model JSON")
    parser.add_argument("--scaler", default=os.path.join('dataset', 'scaler.json'), help="Path to saved StandardScaler JSON")
    args = parser.parse_args()

    # Load model & scaler
    if not os.path.isfile(args.model):
        raise FileNotFoundError(f"Model file not found: {args.model}")
    if not os.path.isfile(args.scaler):
        raise FileNotFoundError(f"Scaler file not found: {args.scaler}")
        
    from xgboost import XGBClassifier
    import json
    
    model = XGBClassifier()
    model.load_model(args.model)
    
    scaler = StandardScaler()
    with open(args.scaler, 'r') as f:
        s_data = json.load(f)
    scaler.mean_ = np.array(s_data["mean"])
    scaler.var_ = np.array(s_data["var"])
    scaler.scale_ = np.array(s_data["scale"])
    scaler.n_features_in_ = s_data["n_features_in"]

    # -------------------- Load image & align --------------------
    bgr = cv2.imread(args.image)
    if bgr is None:
        raise ValueError(f"Could not read image: {args.image}")
    aligner = GeometricAligner(device='cpu')
    aligned = aligner.align_and_crop(bgr, return_tensor=True)
    if aligned is None:
        raise RuntimeError("Face detection failed – no face found in the image.")

    # -------------------- Compute original four scores --------------------
    spatial_score = float(torch.mean(torch.abs(aligned)).item())
    # Frequency score (exactly matching data_pipeline.py on the normalized aligned tensor)
    gray_tensor = 0.2989 * aligned[0:1, :, :] + 0.5870 * aligned[1:2, :, :] + 0.1140 * aligned[2:3, :, :]
    gray_tensor = gray_tensor.unsqueeze(0)  # Shape: [1, 1, 512, 512]
    freq_complex = torch.fft.fft2(gray_tensor)
    freq_shifted = torch.fft.fftshift(torch.abs(freq_complex), dim=(-2, -1))
    freq_tensor = torch.log(1 + freq_shifted)
    freq_score = high_freq_energy(freq_tensor)
    # Latent error (TAESD auto‑encoder)
    error_loop = DiffusionErrorLoop(device='cpu')
    latent_err = error_loop(aligned.unsqueeze(0))
    latent_score = float(torch.mean(torch.abs(latent_err)).item())
    # Statistical embedding (LBP / entropy tensor)
    stat_extractor = StatisticalFeatureExtractor()
    stat_tensor = stat_extractor(aligned.unsqueeze(0)).cpu()
    stat_score = float(torch.mean(stat_tensor).item())

    # -------------------- New lightweight descriptors --------------------
    # Denormalize: pixel = (normalized * std + mean) * 255
    mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)
    std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)
    aligned_np = aligned.cpu().numpy()
    unnorm = (aligned_np * std + mean) * 255.0
    rgb = np.clip(unnorm, 0, 255).transpose(1, 2, 0).astype(np.uint8)

    entropy_score = compute_entropy(rgb)
    edge_density_score = compute_edge_density(rgb)
    laplacian_var_score = compute_laplacian_variance(rgb)
    color_kurtosis_score = compute_color_kurtosis(rgb)
    jpeg_consistency_score = compute_jpeg_consistency(rgb)

    # -------------------- Assemble feature vector --------------------
    feature_dict = {
        "spatial_score": spatial_score,
        "freq_score": freq_score,
        "latent_score": latent_score,
        "stat_score": stat_score,
        "entropy": entropy_score,
        "edge_density": edge_density_score,
        "laplacian_variance": laplacian_var_score,
        "color_kurtosis": color_kurtosis_score,
        "jpeg_consistency": jpeg_consistency_score,
    }
    df_feat = pd.DataFrame([feature_dict])
    # Scale using the same scaler that was fit on training data
    df_scaled = scaler.transform(df_feat)

    # -------------------- Predict --------------------
    prob_fake = model.predict_proba(df_scaled)[0, 1]
    verdict = "AI Generated" if prob_fake >= 0.5 else "Real"
    confidence = prob_fake * 100 if prob_fake >= 0.5 else (1 - prob_fake) * 100

    # -------------------- Pretty report --------------------
    print("\n" + "=" * 40)
    print("          AI IMAGE FORENSIC REPORT")
    print("=" * 40 + "\n")
    print(f"Spatial Artifact Score   : {spatial_score:.2f}")
    print(f"Frequency Anomaly Score  : {freq_score:.2f}")
    print(f"Noise Residual Score     : {latent_score:.2f}")
    print(f"Embedding Consistency    : {stat_score:.2f}")
    print(f"Entropy Score            : {entropy_score:.2f}")
    print(f"Edge Density Score       : {edge_density_score:.2f}")
    print(f"Laplacian Variance Score : {laplacian_var_score:.2f}")
    print(f"Color Kurtosis Score     : {color_kurtosis_score:.2f}")
    print(f"JPEG Consistency Score   : {jpeg_consistency_score:.2f}")
    print("-" * 40)
    print("FINAL RESULT:")
    print(f"Likely {verdict}")
    print(f"Confidence: {confidence:.0f}%")
    print("=" * 40 + "\n")

if __name__ == "__main__":
    main()