features/heuristics.py

"""Heuristic image feature extraction using OpenCV/numpy/scipy."""

import cv2
import numpy as np
from PIL import Image
from scipy.stats import entropy as scipy_entropy
from typing import Dict


def compute_heuristic_features(img: Image.Image) -> Dict[str, float]:
    """
    Compute 9+ heuristic features from a PIL Image using OpenCV.
    
    Returns dict with keys:
        edge_density, contrast, whitespace_ratio, near_white_ratio,
        color_entropy, saturation_entropy, symmetry_lr, symmetry_tb,
        sharpness, center_weight
    """
    img_rgb = img.convert("RGB")
    arr = np.array(img_rgb)
    gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)
    H, W = gray.shape

    features = {}

    # 1. Edge density via Canny
    edges = cv2.Canny(gray, 100, 200)
    features["edge_density"] = float(edges.sum()) / (255 * H * W + 1e-8)

    # 2. Contrast (std of luminance / max possible)
    features["contrast"] = float(gray.std()) / 127.5

    # 3. Whitespace ratio (near-white pixels)
    features["whitespace_ratio"] = float((gray > 240).sum()) / (H * W + 1e-8)
    features["near_white_ratio"] = float((gray > 200).sum()) / (H * W + 1e-8)

    # 4. Color entropy (average per-channel)
    entropies = []
    for c in range(3):
        hist, _ = np.histogram(arr[:, :, c], bins=256, range=(0, 256))
        entropies.append(float(scipy_entropy(hist + 1)))
    features["color_entropy"] = float(np.mean(entropies))

    # 5. Saturation entropy (from HSV)
    hsv = cv2.cvtColor(arr, cv2.COLOR_RGB2HSV)
    sat_hist, _ = np.histogram(hsv[:, :, 1].ravel(), bins=256, range=(0, 256))
    features["saturation_entropy"] = float(scipy_entropy(sat_hist + 1))

    # 6. Symmetry (left-right)
    flipped_lr = np.fliplr(gray).astype(float)
    gray_f = gray.astype(float)
    denom_lr = np.sqrt((gray_f ** 2).sum() * (flipped_lr ** 2).sum()) + 1e-8
    features["symmetry_lr"] = float((gray_f * flipped_lr).sum() / denom_lr)

    # 7. Symmetry (top-bottom)
    flipped_tb = np.flipud(gray).astype(float)
    denom_tb = np.sqrt((gray_f ** 2).sum() * (flipped_tb ** 2).sum()) + 1e-8
    features["symmetry_tb"] = float((gray_f * flipped_tb).sum() / denom_tb)

    # 8. Sharpness (Laplacian variance, log-normalized)
    lap_var = cv2.Laplacian(gray, cv2.CV_64F).var()
    features["sharpness"] = float(np.log1p(lap_var) / 10.0)

    # 9. Center weight (rule of thirds center vs overall)
    h3, w3 = max(1, H // 3), max(1, W // 3)
    if h3 * 2 <= H and w3 * 2 <= W:
        thirds_zone = gray[h3:2*h3, w3:2*w3]
        center_mean = float(thirds_zone.mean())
    else:
        center_mean = float(gray.mean())
    features["center_weight"] = center_mean / (gray.mean() + 1e-8)

    return features