Upload ml/07_detect_redactions.py with huggingface_hub

ml/07_detect_redactions.py

@@ -0,0 +1,300 @@
+#!/usr/bin/env python3
+"""
+Phase 4: Redaction Detection
+
+Detects blacked-out/redacted regions in scanned document pages using
+OpenCV contour analysis on rendered PDF pages.
+
+Algorithm:
+1. Render page at 150 DPI via PyMuPDF
+2. Convert to grayscale → binary threshold (< 30 = black)
+3. Find contours via OpenCV
+4. Filter: rectangular (> 0.85), min area (500px), aspect ratio (0.1-10)
+5. Store redaction count, area percentage, and bounding boxes
+
+Stores results in page_features (per page) and document_features (aggregate).
+Processes priority collections first (CIA, JFK, DOJ, Lincoln).
+
+Runs on: Hetzner CPU with joblib parallelism
+"""
+
+import json
+import logging
+import os
+import sys
+from multiprocessing import cpu_count
+
+import cv2
+import fitz  # PyMuPDF
+import numpy as np
+import psycopg2
+import psycopg2.extras
+
+from db import get_conn
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)-8s  %(message)s")
+log = logging.getLogger(__name__)
+
+DPI = 150
+BLACK_THRESHOLD = 30
+MIN_AREA = 500
+MIN_RECTANGULARITY = 0.85
+MIN_ASPECT = 0.1
+MAX_ASPECT = 10.0
+BATCH_SIZE = 100  # pages per DB flush
+WORKERS = max(1, cpu_count() - 2)  # leave 2 cores free
+
+# Process these collections first — most likely to have redactions
+PRIORITY_SECTIONS = [
+    'cia_declassified', 'cia_mkultra', 'cia_stargate',
+    'jfk_assassination', 'doj_disclosures', 'lincoln_archives',
+    'house_resolutions',
+]
+
+
+def get_pending_documents(conn, section, limit=500):
+    """Get documents that haven't been analyzed for redactions yet."""
+    with conn.cursor() as cur:
+        cur.execute("""
+            SELECT d.id, d.file_path
+            FROM documents d
+            WHERE d.source_section = %s
+              AND d.id NOT IN (
+                  SELECT DISTINCT df.document_id FROM document_features df
+                  WHERE df.feature_name = 'redaction_summary'
+              )
+            ORDER BY d.id
+            LIMIT %s
+        """, (section, limit))
+        return cur.fetchall()
+
+
+def get_pages_for_doc(conn, doc_id):
+    """Get page IDs and numbers for a document."""
+    with conn.cursor() as cur:
+        cur.execute("""
+            SELECT id, page_number FROM pages
+            WHERE document_id = %s
+              AND id NOT IN (
+                  SELECT page_id FROM page_features WHERE feature_name = 'redaction'
+              )
+            ORDER BY page_number
+        """, (doc_id,))
+        return cur.fetchall()
+
+
+def detect_redactions_page(pdf_path, page_num):
+    """Detect redacted regions on a single PDF page."""
+    try:
+        doc = fitz.open(pdf_path)
+        if page_num - 1 >= len(doc):
+            doc.close()
+            return None
+
+        page = doc[page_num - 1]  # 0-indexed
+        # Render at DPI
+        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)
+        doc.close()
+
+        # Convert to grayscale
+        if img.shape[2] == 4:  # RGBA
+            gray = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY)
+        elif img.shape[2] == 3:  # RGB
+            gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+        else:
+            gray = img[:, :, 0]
+
+        # Threshold: black regions
+        _, binary = cv2.threshold(gray, BLACK_THRESHOLD, 255, cv2.THRESH_BINARY_INV)
+
+        # Morphological close to merge nearby black regions
+        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
+        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
+
+        # Find contours
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        page_area = pix.width * pix.height
+        redactions = []
+
+        for cnt in contours:
+            area = cv2.contourArea(cnt)
+            if area < MIN_AREA:
+                continue
+
+            x, y, w, h = cv2.boundingRect(cnt)
+            rect_area = w * h
+            if rect_area == 0:
+                continue
+
+            rectangularity = area / rect_area
+            aspect = w / h if h > 0 else 0
+
+            if rectangularity >= MIN_RECTANGULARITY and MIN_ASPECT <= aspect <= MAX_ASPECT:
+                # Additional filter: must be at least 0.1% of page
+                if area / page_area >= 0.001:
+                    redactions.append({
+                        'x': int(x), 'y': int(y),
+                        'w': int(w), 'h': int(h),
+                        'area': int(area),
+                        'area_pct': round(area / page_area * 100, 2),
+                    })
+
+        total_redacted_area = sum(r['area'] for r in redactions)
+        return {
+            'count': len(redactions),
+            'total_area_pct': round(total_redacted_area / page_area * 100, 2),
+            'bboxes': redactions,
+        }
+
+    except Exception as e:
+        log.debug(f"Error on {pdf_path} p{page_num}: {e}")
+        return None
+
+
+def process_document(conn, doc_id, file_path):
+    """Process all pages of a document for redactions."""
+    # Resolve PDF path
+    pdf_path = None
+    for base in ['/data/raw/', '/data/ocr_output/']:
+        candidate = os.path.join(base, file_path)
+        if os.path.exists(candidate):
+            pdf_path = candidate
+            break
+
+    # Also try the file_path directly
+    if pdf_path is None and os.path.exists(file_path):
+        pdf_path = file_path
+
+    # Try finding the original PDF from file_path pattern
+    if pdf_path is None:
+        # file_path might be relative, try common patterns
+        for base in ['/data/raw/', '/data/']:
+            candidate = os.path.join(base, file_path)
+            if os.path.exists(candidate):
+                pdf_path = candidate
+                break
+
+    if pdf_path is None:
+        return 0, 0, []  # Can't find PDF
+
+    pages = get_pages_for_doc(conn, doc_id)
+    if not pages:
+        return 0, 0, []
+
+    page_results = []
+    total_redactions = 0
+    max_area_pct = 0
+
+    for page_id, page_num in pages:
+        result = detect_redactions_page(pdf_path, page_num)
+        if result is None:
+            continue
+
+        page_results.append((
+            page_id,
+            'redaction',
+            result['count'],
+            json.dumps(result),
+        ))
+
+        total_redactions += result['count']
+        if result['total_area_pct'] > max_area_pct:
+            max_area_pct = result['total_area_pct']
+
+    return total_redactions, max_area_pct, page_results
+
+
+def flush_page_features(conn, rows):
+    if not rows:
+        return
+    with conn.cursor() as cur:
+        psycopg2.extras.execute_batch(
+            cur,
+            """INSERT INTO page_features (page_id, feature_name, feature_value, feature_json)
+               VALUES (%s, %s, %s, %s::jsonb)
+               ON CONFLICT (page_id, feature_name) DO UPDATE SET
+                 feature_value = EXCLUDED.feature_value,
+                 feature_json = EXCLUDED.feature_json,
+                 created_at = NOW()""",
+            rows,
+            page_size=500,
+        )
+    conn.commit()
+
+
+def flush_doc_features(conn, rows):
+    if not rows:
+        return
+    with conn.cursor() as cur:
+        psycopg2.extras.execute_batch(
+            cur,
+            """INSERT INTO document_features (document_id, feature_name, feature_value, feature_json)
+               VALUES (%s, %s, %s, %s::jsonb)
+               ON CONFLICT (document_id, feature_name) DO UPDATE SET
+                 feature_value = EXCLUDED.feature_value,
+                 feature_json = EXCLUDED.feature_json,
+                 created_at = NOW()""",
+            rows,
+            page_size=500,
+        )
+    conn.commit()
+
+
+def main():
+    conn = get_conn()
+    grand_total = 0
+    docs_with_redactions = 0
+
+    for section in PRIORITY_SECTIONS:
+        log.info(f"=== Processing section: {section} ===")
+        batch_num = 0
+
+        while True:
+            docs = get_pending_documents(conn, section, limit=100)
+            if not docs:
+                break
+
+            page_feature_buffer = []
+            doc_feature_buffer = []
+
+            for doc_id, file_path in docs:
+                total_redactions, max_area_pct, page_results = process_document(conn, doc_id, file_path)
+
+                page_feature_buffer.extend(page_results)
+
+                # Store document-level summary
+                summary = {
+                    'total_redactions': total_redactions,
+                    'max_page_area_pct': max_area_pct,
+                    'pages_analyzed': len(page_results),
+                }
+                doc_feature_buffer.append((
+                    doc_id,
+                    'redaction_summary',
+                    float(total_redactions),
+                    json.dumps(summary),
+                ))
+
+                if total_redactions > 0:
+                    docs_with_redactions += 1
+
+                grand_total += total_redactions
+
+            flush_page_features(conn, page_feature_buffer)
+            flush_doc_features(conn, doc_feature_buffer)
+
+            batch_num += 1
+            log.info(f"  {section} batch {batch_num}: {len(docs)} docs, "
+                     f"{len(page_feature_buffer)} page results, "
+                     f"{grand_total} total redactions found, "
+                     f"{docs_with_redactions} docs with redactions")
+
+    conn.close()
+    log.info(f"Done. {grand_total} redactions across {docs_with_redactions} documents.")
+
+
+if __name__ == "__main__":
+    main()