Upload src/poster_sentry/features.py with huggingface_hub

src/poster_sentry/features.py

@@ -0,0 +1,217 @@
+"""
+Feature extractors for PosterSentry.
+
+Two feature channels:
+    1. Visual features — image-level statistics (color, edges, FFT, whitespace)
+    2. PDF structural features — page geometry, text blocks, font diversity
+
+Both are cheap to compute (no GPU needed), providing strong priors that
+complement the text embedding from model2vec.
+"""
+
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# ── Visual Feature Extractor ────────────────────────────────────
+
+VISUAL_FEATURE_NAMES = [
+    "img_width",
+    "img_height",
+    "img_aspect_ratio",
+    "mean_r", "mean_g", "mean_b",
+    "std_r", "std_g", "std_b",
+    "local_contrast",
+    "color_diversity",
+    "edge_density",
+    "spatial_complexity",
+    "white_space_ratio",
+    "high_contrast_ratio",
+]
+
+N_VISUAL_FEATURES = len(VISUAL_FEATURE_NAMES)
+
+
+class VisualFeatureExtractor:
+    """Extract visual features from rendered PDF pages."""
+
+    FEATURE_NAMES = VISUAL_FEATURE_NAMES
+
+    def __init__(self, target_size: Tuple[int, int] = (256, 256)):
+        self.target_size = target_size
+
+    def pdf_to_image(self, pdf_path: str, dpi: int = 72) -> Optional[np.ndarray]:
+        """Render first page of PDF to RGB numpy array."""
+        try:
+            import fitz
+            doc = fitz.open(pdf_path)
+            if len(doc) == 0:
+                doc.close()
+                return None
+            page = doc[0]
+            mat = fitz.Matrix(dpi / 72, dpi / 72)
+            pix = page.get_pixmap(matrix=mat)
+            img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
+            if pix.n == 4:
+                img = img[:, :, :3]
+            elif pix.n == 1:
+                img = np.stack([img[:, :, 0]] * 3, axis=-1)
+            doc.close()
+            return img
+        except Exception as e:
+            logger.debug(f"PDF to image failed: {e}")
+            return None
+
+    def extract(self, image: np.ndarray) -> Dict[str, float]:
+        """Extract 15 visual features from an RGB image."""
+        feats = {n: 0.0 for n in self.FEATURE_NAMES}
+        try:
+            from PIL import Image as PILImage
+
+            h, w = image.shape[:2]
+            feats["img_width"] = float(w)
+            feats["img_height"] = float(h)
+            feats["img_aspect_ratio"] = w / h if h > 0 else 0.0
+
+            pil = PILImage.fromarray(image).resize(self.target_size, PILImage.Resampling.BILINEAR)
+            resized = np.array(pil)
+
+            for i, ch in enumerate(["r", "g", "b"]):
+                feats[f"mean_{ch}"] = float(np.mean(resized[:, :, i]))
+                feats[f"std_{ch}"] = float(np.std(resized[:, :, i]))
+
+            gray = np.mean(resized, axis=2)
+            feats["local_contrast"] = float(np.std(gray))
+
+            # Color diversity (unique quantized colors in 32x32 thumbnail)
+            small = np.array(pil.resize((32, 32)))
+            quantized = (small // 32).astype(np.uint8)
+            unique_colors = len(np.unique(quantized.reshape(-1, 3), axis=0))
+            feats["color_diversity"] = unique_colors / 512.0
+
+            # Edge density
+            gy = np.abs(np.diff(gray, axis=0))
+            gx = np.abs(np.diff(gray, axis=1))
+            feats["edge_density"] = float(np.mean(gy) + np.mean(gx)) / 255.0
+
+            # Spatial complexity (high-freq ratio via FFT)
+            fft = np.fft.fft2(gray)
+            fft_shift = np.fft.fftshift(fft)
+            mag = np.abs(fft_shift)
+            ch, cw = mag.shape[0] // 2, mag.shape[1] // 2
+            radius = min(mag.shape) // 4
+            y, x = np.ogrid[:mag.shape[0], :mag.shape[1]]
+            center_mask = ((y - ch) ** 2 + (x - cw) ** 2) <= radius ** 2
+            total_e = np.sum(mag ** 2)
+            low_e = np.sum(mag[center_mask] ** 2)
+            feats["spatial_complexity"] = 1.0 - (low_e / total_e) if total_e > 0 else 0.0
+
+            # White space ratio
+            white_px = np.sum(np.all(resized > 240, axis=2))
+            feats["white_space_ratio"] = white_px / (self.target_size[0] * self.target_size[1])
+
+            # High contrast ratio (very dark + very bright pixels)
+            feats["high_contrast_ratio"] = float(np.sum(gray < 50) + np.sum(gray > 240)) / gray.size
+
+        except Exception as e:
+            logger.debug(f"Visual feature extraction failed: {e}")
+        return feats
+
+    def to_vector(self, feats: Dict[str, float]) -> np.ndarray:
+        return np.array([feats.get(n, 0.0) for n in self.FEATURE_NAMES], dtype="float32")
+
+
+# ── PDF Structural Feature Extractor ────────────────────────────
+
+STRUCTURAL_FEATURE_NAMES = [
+    "page_count",
+    "page_width_pt",
+    "page_height_pt",
+    "page_aspect_ratio",
+    "page_area_sqin",
+    "is_landscape",
+    "text_block_count",
+    "font_count",
+    "avg_font_size",
+    "font_size_variance",
+    "title_score",
+    "text_density",
+    "line_count",
+    "file_size_kb",
+    "size_per_page_kb",
+]
+
+N_STRUCTURAL_FEATURES = len(STRUCTURAL_FEATURE_NAMES)
+
+
+class PDFStructuralExtractor:
+    """Extract structural features from PDF layout."""
+
+    FEATURE_NAMES = STRUCTURAL_FEATURE_NAMES
+
+    def extract(self, pdf_path: str) -> Dict[str, float]:
+        """Extract 15 structural features from a PDF."""
+        feats = {n: 0.0 for n in self.FEATURE_NAMES}
+        try:
+            import fitz
+            path = Path(pdf_path)
+            doc = fitz.open(str(path))
+            if len(doc) == 0:
+                doc.close()
+                return feats
+
+            feats["page_count"] = float(len(doc))
+            feats["file_size_kb"] = path.stat().st_size / 1024.0
+            feats["size_per_page_kb"] = feats["file_size_kb"] / max(len(doc), 1)
+
+            page = doc[0]
+            rect = page.rect
+            feats["page_width_pt"] = rect.width
+            feats["page_height_pt"] = rect.height
+            feats["page_aspect_ratio"] = rect.width / rect.height if rect.height > 0 else 0.0
+            feats["page_area_sqin"] = (rect.width / 72.0) * (rect.height / 72.0)
+            feats["is_landscape"] = float(rect.width > rect.height)
+
+            # Text blocks
+            blocks = page.get_text("dict")["blocks"]
+            text_blocks = [b for b in blocks if b.get("type") == 0]
+            feats["text_block_count"] = float(len(text_blocks))
+
+            if text_blocks:
+                heights = [b["bbox"][3] - b["bbox"][1] for b in text_blocks]
+                widths = [b["bbox"][2] - b["bbox"][0] for b in text_blocks]
+                total_area = sum(h * w for h, w in zip(heights, widths))
+                page_area = rect.width * rect.height
+                feats["text_density"] = total_area / page_area if page_area > 0 else 0.0
+
+            # Font statistics
+            fonts = set()
+            font_sizes = []
+            line_count = 0
+            for block in text_blocks:
+                for line in block.get("lines", []):
+                    line_count += 1
+                    for span in line.get("spans", []):
+                        fonts.add(span.get("font", ""))
+                        sz = span.get("size", 0)
+                        if sz > 0:
+                            font_sizes.append(sz)
+
+            feats["font_count"] = float(len(fonts))
+            feats["line_count"] = float(line_count)
+            if font_sizes:
+                feats["avg_font_size"] = float(np.mean(font_sizes))
+                feats["font_size_variance"] = float(np.var(font_sizes)) if len(font_sizes) > 1 else 0.0
+                feats["title_score"] = max(font_sizes) / (np.mean(font_sizes) + 1.0)
+
+            doc.close()
+        except Exception as e:
+            logger.debug(f"PDF structural extraction failed: {e}")
+        return feats
+
+    def to_vector(self, feats: Dict[str, float]) -> np.ndarray:
+        return np.array([feats.get(n, 0.0) for n in self.FEATURE_NAMES], dtype="float32")