features.py

"""Poster feature extraction shared by training and the Gradio app.

Training and inference must call the same function, otherwise the model
sees differently computed columns at serving time (train/serve skew).
"""

import os

import numpy as np
import pytesseract
from PIL import Image, ImageDraw, ImageFilter, ImageOps

# Allow a custom Tesseract path on Windows without editing code
if os.environ.get("TESSERACT_CMD"):
    pytesseract.pytesseract.tesseract_cmd = os.environ["TESSERACT_CMD"]

# Column order is the model contract; training must use this exact list
FEATURE_COLUMNS = [
    "img_width_px",
    "img_height_px",
    "aspect_ratio",
    "word_count",
    "char_count",
    "text_block_count",
    "text_area_frac",
    "graphic_area_frac",
    "background_area_frac",
    "text_share",
    "text_to_graphic_ratio",
    "mean_text_height_frac",
    "max_text_height_frac",
    "text_height_std_frac",
]

MAX_SIDE_PX = 1600  # Normalize scale so OCR speed and noise stay consistent
MIN_SIDE_PX = 200  # Below this OCR output is unreliable, so reject
MIN_WORD_CONF = 60  # Tesseract confidence floor for keeping a word
BOX_PAD_PX = 3  # Small pad so anti-aliased glyph edges count as text
BG_DIST_THRESHOLD = 40.0  # RGB distance separating content from background
BORDER_FRAC = 0.03  # Border strip width used to estimate background color
RATIO_EPS = 1e-3  # Caps text_to_graphic_ratio when no graphics are found


def prepare_image(image):
    """Apply EXIF rotation, force RGB, and downscale large posters."""
    image = ImageOps.exif_transpose(image).convert("RGB")  # Phone photos often carry rotation tags
    scale = MAX_SIDE_PX / max(image.size)  # Only shrink, never upscale small images
    if scale < 1:
        new_size = (round(image.width * scale), round(image.height * scale))
        image = image.resize(new_size, Image.LANCZOS)  # Lanczos keeps small text edges sharp
    return image


def _ocr_words(image):
    """Run Tesseract and keep confident, non-empty word boxes."""
    # PSM 11 finds sparse text anywhere, which suits poster layouts
    gray = ImageOps.grayscale(image)  # Color shapes confuse Tesseract layout analysis
    data = pytesseract.image_to_data(
        gray, config="--psm 11", output_type=pytesseract.Output.DICT
    )
    words = []
    for i, text in enumerate(data["text"]):
        text = text.strip()
        conf = float(data["conf"][i])  # Tesseract returns -1 for non-word rows
        n_alnum = sum(c.isalnum() for c in text)
        if conf < MIN_WORD_CONF or (n_alnum < 2 and not text.isdigit()):
            continue  # Single glyphs on shapes are the most common false positive
        words.append({
            "text": text,
            "left": data["left"][i],
            "top": data["top"][i],
            "width": data["width"][i],
            "height": data["height"][i],
            "block": (data["page_num"][i], data["block_num"][i]),  # Unique block id
        })
    return words


def _background_color(rgb):
    """Estimate background as the median color of the border strip."""
    h, w, _ = rgb.shape
    b = max(1, int(min(h, w) * BORDER_FRAC))  # Strip width scales with poster size
    border = np.concatenate([
        rgb[:b].reshape(-1, 3), rgb[-b:].reshape(-1, 3),
        rgb[:, :b].reshape(-1, 3), rgb[:, -b:].reshape(-1, 3),
    ])
    return np.median(border, axis=0)  # Median resists logos touching the edge


def extract_features(image):
    """Return (feature dict, overlay image) for one poster.

    The overlay shades detected graphics blue and outlines text in orange
    so a user can check what the numbers were computed from.
    """
    original_w, original_h = image.size  # Report true size, not the resized one
    img = prepare_image(image)
    w, h = img.size

    words = _ocr_words(img)

    # Text mask: union of padded word boxes
    text_mask = np.zeros((h, w), dtype=bool)
    for wd in words:
        x0 = max(0, wd["left"] - BOX_PAD_PX)
        y0 = max(0, wd["top"] - BOX_PAD_PX)
        x1 = min(w, wd["left"] + wd["width"] + BOX_PAD_PX)
        y1 = min(h, wd["top"] + wd["height"] + BOX_PAD_PX)
        text_mask[y0:y1, x0:x1] = True

    # Content mask: pixels far from the background color after denoising
    smooth = np.asarray(img.filter(ImageFilter.MedianFilter(5)), dtype=np.float32)
    bg = _background_color(smooth)
    dist = np.linalg.norm(smooth - bg, axis=2)  # Per-pixel RGB distance to background
    content_mask = dist > BG_DIST_THRESHOLD

    # Graphics are content that is not inside a text box
    graphic_mask = content_mask & ~text_mask

    text_frac = float(text_mask.mean())
    graphic_frac = float(graphic_mask.mean())
    heights = np.array([wd["height"] / h for wd in words]) if words else np.zeros(1)

    row = {
        "img_width_px": original_w,
        "img_height_px": original_h,
        "aspect_ratio": original_w / original_h,
        "word_count": len(words),
        "char_count": sum(len(wd["text"]) for wd in words),
        "text_block_count": len({wd["block"] for wd in words}),
        "text_area_frac": text_frac,
        "graphic_area_frac": graphic_frac,
        "background_area_frac": max(0.0, 1.0 - text_frac - graphic_frac),
        # Bounded 0..1 share is easier for models than an unbounded ratio
        "text_share": text_frac / (text_frac + graphic_frac) if (text_frac + graphic_frac) > 0 else 0.0,
        "text_to_graphic_ratio": text_frac / max(graphic_frac, RATIO_EPS),
        "mean_text_height_frac": float(heights.mean()),
        "max_text_height_frac": float(heights.max()),
        "text_height_std_frac": float(heights.std()),
    }
    assert list(row) == FEATURE_COLUMNS, "Feature dict drifted from FEATURE_COLUMNS."
    return row, _draw_overlay(img, graphic_mask, words)


def _draw_overlay(img, graphic_mask, words):
    """Tint graphic pixels blue and outline each OCR word box."""
    base = np.asarray(img, dtype=np.float32)
    tint = np.array([40, 110, 255], dtype=np.float32)
    base[graphic_mask] = 0.55 * base[graphic_mask] + 0.45 * tint  # Semi-transparent blend
    overlay = Image.fromarray(base.astype(np.uint8))
    draw = ImageDraw.Draw(overlay)
    for wd in words:
        box = [wd["left"], wd["top"], wd["left"] + wd["width"], wd["top"] + wd["height"]]
        draw.rectangle(box, outline=(255, 140, 0), width=3)  # Orange reads on most palettes
    return overlay