app.py

"""
PDF Annotation Extractor — Flask app
pikepdf + pypdfium2 + Tesseract OCR (Bengali + English)

TYPE column shows the comment/label text linked to each annotation (e.g. "FAC").
If no comment is linked, TYPE is blank.

Label detection uses four methods (in priority order):
  1. /Contents directly on the markup annotation (Acrobat inline comment)
  2. /IRT reverse-lookup (sticky note replies to markup via In Reply To)
  3. Spatial proximity (sticky note rect within 80 pt of markup rect)
  4. Array adjacency (sticky note immediately before/after markup in /Annots)
"""

import os, io, re, csv, json
import numpy as np
import pikepdf
import pypdfium2 as pdfium
from PIL import Image, ImageEnhance
import pytesseract
from flask import (Flask, request, render_template,
                   send_file, redirect, url_for, flash)
from werkzeug.utils import secure_filename

# ── Paths ──────────────────────────────────────────────────────────────────────
BASE_DIR   = os.path.dirname(os.path.abspath(__file__))
UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")
CSV_PATH   = os.path.join(OUTPUT_DIR, "annotations.csv")
JSON_PATH  = os.path.join(OUTPUT_DIR, "annotations.json")
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ── Flask ──────────────────────────────────────────────────────────────────────
app = Flask(__name__)
app.secret_key = "pdf-annotator-secret-key"
app.config["MAX_CONTENT_LENGTH"] = 50 * 1024 * 1024
ALLOWED_EXTENSIONS = {"pdf"}

def allowed_file(f: str) -> bool:
    return "." in f and f.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS

# ── Constants ──────────────────────────────────────────────────────────────────
MARKUP_SUBTYPES = {"/Highlight", "/Underline", "/Squiggly", "/StrikeOut", "/FreeText"}

RENDER_DPI   = 300
RENDER_SCALE = RENDER_DPI / 72.0
V_PAD = 4    # vertical padding in PDF points
H_PAD = 2    # horizontal padding in PDF points (minimal — prevents adjacent text bleed)

# Max distance (in PDF points) between a sticky note and a markup rect
# for them to be considered linked by proximity
PROXIMITY_MARGIN = 80


# ══════════════════════════════════════════════════════════════════════════════
# IMAGE PREPROCESSING
# ══════════════════════════════════════════════════════════════════════════════

def preprocess(img: Image.Image) -> Image.Image:
    if img.mode == "RGBA":
        bg = Image.new("RGB", img.size, (255, 255, 255))
        bg.paste(img, mask=img.split()[3])
        img = bg
    if img.mode != "L":
        img = img.convert("L")
    arr = np.array(img, dtype=np.float32)
    if arr.mean() < 128:
        arr = 255.0 - arr
    lo, hi = arr.min(), arr.max()
    if hi > lo:
        arr = (arr - lo) / (hi - lo) * 255.0
    img = Image.fromarray(arr.astype(np.uint8))
    img = ImageEnhance.Contrast(img).enhance(3.0)
    img = ImageEnhance.Sharpness(img).enhance(2.5)
    if img.height < 80:
        f = max(2, 80 // max(img.height, 1))
        img = img.resize((img.width * f, img.height * f), Image.LANCZOS)
    return img


# ══════════════════════════════════════════════════════════════════════════════
# CROPPING
# ══════════════════════════════════════════════════════════════════════════════

def pdf_to_img_box(x0, y0_pdf, x1, y1_pdf, page_height,
                   h_pad=H_PAD, v_pad=V_PAD):
    top    = (page_height - y1_pdf) * RENDER_SCALE - v_pad * RENDER_SCALE
    bottom = (page_height - y0_pdf) * RENDER_SCALE + v_pad * RENDER_SCALE
    left   =  x0 * RENDER_SCALE    - h_pad * RENDER_SCALE
    right  =  x1 * RENDER_SCALE    + h_pad * RENDER_SCALE
    if right <= left or bottom <= top:
        return None
    return (int(max(0, left)), int(max(0, top)), int(right), int(bottom))


def do_crop(full_img, box):
    if box is None:
        return None
    W, H = full_img.size
    b = (box[0], box[1], min(W, box[2]), min(H, box[3]))
    if b[2] <= b[0] or b[3] <= b[1]:
        return None
    return full_img.crop(b)


# ══════════════════════════════════════════════════════════════════════════════
# OCR
# ══════════════════════════════════════════════════════════════════════════════

def clean_text(text: str) -> str:
    if not text:
        return ""
    bengali = sum(1 for c in text if "\u0980" <= c <= "\u09FF")
    alpha   = sum(1 for c in text if c.isalpha())
    if alpha > 0 and bengali / alpha > 0.35:
        text = re.sub(r"^[ংঃঁ।\s\-_|\"\'.,;:!?0-9]+", "", text)
        text = re.sub(r"(?<![a-zA-Z])[a-zA-Z]{1,3}(?![a-zA-Z])", "", text)
    text = re.sub(r"[\"\']+$", "", text)
    text = re.sub(r"[|\\`~\[\]{}@#$%^&*_+=<>]", "", text)
    text = re.sub(r"\n+", " ", text)
    text = re.sub(r" {2,}", " ", text)
    return text.strip()


def ocr_image(img: Image.Image, single_line: bool = True) -> str:
    processed = preprocess(img)
    psm_list  = [7, 13, 6] if single_line else [6, 7, 13]
    candidates = []
    for psm in psm_list:
        for lang in ("ben", "ben+eng"):
            try:
                raw = pytesseract.image_to_string(
                    processed, lang=lang, config=f"--psm {psm} --oem 1"
                )
                t = clean_text(re.sub(r"\n+", " ", raw).strip())
                if t:
                    candidates.append(t)
            except Exception:
                continue
    if not candidates:
        return ""
    return max(candidates, key=len)


# ══════════════════════════════════════════════════════════════════════════════
# ANNOTATION LABEL DETECTION
# ══════════════════════════════════════════════════════════════════════════════

def get_rect(annot: pikepdf.Dictionary) -> list | None:
    """Return [x0, y0, x1, y1] from annotation /Rect, or None."""
    try:
        r = annot.get("/Rect")
        if r:
            return [float(v) for v in r]
    except Exception:
        pass
    return None


def rects_are_near(rect_a: list, rect_b: list, margin: float) -> bool:
    """
    Return True if rect_a and rect_b overlap or are within `margin` PDF points
    of each other (in both X and Y).
    """
    ax0, ay0, ax1, ay1 = rect_a
    bx0, by0, bx1, by1 = rect_b
    # Expand rect_b by margin
    bx0 -= margin; by0 -= margin
    bx1 += margin; by1 += margin
    return not (ax1 < bx0 or ax0 > bx1 or ay1 < by0 or ay0 > by1)


def resolve_obj(ref):
    """Dereference a pikepdf indirect reference to a Dictionary."""
    try:
        if isinstance(ref, pikepdf.Dictionary):
            return ref
        obj = ref.get_object()
        return obj if isinstance(obj, pikepdf.Dictionary) else None
    except Exception:
        return None


def build_label_maps(annots_raw: list):
    """
    Pre-process the full /Annots array to build two label-lookup structures:

    irt_map   : {markup_objgen → label_text}
                from /Text annotations that have /IRT pointing to a markup.

    text_pool : list of {"rect", "contents", "idx"} for /Text annotations
                that are NOT linked via /IRT (standalone sticky notes).
                Used for proximity and adjacency matching.

    Also returns a set of indices occupied by linked /Text annotations
    so they are not matched again by proximity.
    """
    irt_map   = {}
    text_pool = []
    linked_text_indices = set()

    for idx, ref in enumerate(annots_raw):
        a = resolve_obj(ref)
        if a is None:
            continue
        if str(a.get("/Subtype", "")) != "/Text":
            continue

        contents = str(a.get("/Contents", "")).strip()
        if not contents:
            continue

        irt_ref = a.get("/IRT")
        if irt_ref is not None:
            try:
                target_id = irt_ref.objgen if hasattr(irt_ref, "objgen") else None
                if target_id:
                    irt_map[target_id] = contents
                    linked_text_indices.add(idx)
                    continue
            except Exception:
                pass

        # Not linked via IRT — add to pool for proximity/adjacency matching
        rect = get_rect(a)
        if rect:
            text_pool.append({"rect": rect, "contents": contents, "idx": idx})

    return irt_map, text_pool, linked_text_indices


def get_label(annot_ref, annot: pikepdf.Dictionary,
              markup_idx: int,
              irt_map: dict,
              text_pool: list,
              used_text_indices: set) -> str:
    """
    Find the comment/label text for a markup annotation using four methods:

    1. /Contents on the markup itself          (Acrobat inline comment)
    2. /IRT reverse-lookup                     (reply thread)
    3. Spatial proximity from text_pool        (sticky note near markup)
    4. Array adjacency from text_pool          (sticky note adjacent in array)

    Returns the label string, or "" if nothing is found.
    Marks matched pool entries as used so they aren't reused.
    """

    # ── Method 1: /Contents directly on the markup ─────────────────────────
    contents = str(annot.get("/Contents", "")).strip()
    if contents:
        return contents

    # ── Method 2: /IRT reverse-lookup ─────────────────────────────────────
    try:
        obj_id = annot_ref.objgen if hasattr(annot_ref, "objgen") else None
        if obj_id and obj_id in irt_map:
            return irt_map[obj_id]
    except Exception:
        pass

    # ── Methods 3 & 4: use text_pool (unlinked sticky notes) ──────────────
    markup_rect = get_rect(annot)
    if markup_rect is None or not text_pool:
        return ""

    # Method 3: spatial proximity
    for entry in text_pool:
        if entry["idx"] in used_text_indices:
            continue
        if rects_are_near(entry["rect"], markup_rect, PROXIMITY_MARGIN):
            used_text_indices.add(entry["idx"])
            return entry["contents"]

    # Method 4: array adjacency (sticky note within ±2 positions in /Annots)
    for entry in text_pool:
        if entry["idx"] in used_text_indices:
            continue
        if abs(entry["idx"] - markup_idx) <= 2:
            used_text_indices.add(entry["idx"])
            return entry["contents"]

    return ""


# ══════════════════════════════════════════════════════════════════════════════
# MARKUP TEXT EXTRACTION
# ══════════════════════════════════════════════════════════════════════════════

def quads_to_rects(qp: list) -> list:
    rects = []
    for i in range(0, len(qp) - 7, 8):
        xs = qp[i : i + 8 : 2]
        ys = qp[i + 1 : i + 8 : 2]
        rects.append((min(xs), min(ys), max(xs), max(ys)))
    return rects


def extract_markup_text(full_img, annot: pikepdf.Dictionary,
                         page_height: float) -> str:
    rect_obj = annot.get("/Rect")
    annot_h  = float(rect_obj[3]) - float(rect_obj[1]) if rect_obj else 0

    # Strategy 1: QuadPoints — strict X bounds, prevents adjacent text bleed
    quad_obj = annot.get("/QuadPoints")
    if quad_obj and len(quad_obj) >= 8:
        qp     = [float(v) for v in quad_obj]
        rects  = quads_to_rects(qp)
        chunks = []
        for (x0, y0, x1, y1) in rects:
            region = do_crop(full_img, pdf_to_img_box(x0, y0, x1, y1, page_height))
            if region:
                t = ocr_image(region, single_line=True)
                if t:
                    chunks.append(t)
        if chunks:
            return " ".join(chunks)

    # Strategy 2: /Rect fallback
    if rect_obj:
        x0, y0, x1, y1 = [float(v) for v in rect_obj]
        region = do_crop(full_img, pdf_to_img_box(x0, y0, x1, y1, page_height))
        if region:
            t = ocr_image(region, single_line=annot_h <= 25)
            if t:
                return t

    return ""


# ══════════════════════════════════════════════════════════════════════════════
# MAIN EXTRACTION
# ══════════════════════════════════════════════════════════════════════════════

def extract_annotations(pdf_path: str) -> list:
    """
    Extract all markup annotations from the PDF.
    Returns list of dicts: {page, label, text}
      label = comment/label text (e.g. "FAC") or "" if none attached
      text  = OCR'd content of the highlighted/underlined region
    """
    annotations = []
    doc_pk     = pikepdf.open(pdf_path)
    doc_pdfium = pdfium.PdfDocument(pdf_path)

    for page_idx in range(len(doc_pk.pages)):
        page_pk     = doc_pk.pages[page_idx]
        page_pdfium = doc_pdfium[page_idx]

        if "/Annots" not in page_pk:
            continue

        page_height = float(page_pk.mediabox[3])
        annots_raw  = list(page_pk["/Annots"])

        # Build label lookup structures for this page
        irt_map, text_pool, _ = build_label_maps(annots_raw)
        used_text_indices = set()

        # Render page image once — reused for all annotations
        bitmap   = page_pdfium.render(scale=RENDER_SCALE)
        full_img = bitmap.to_pil()

        for markup_idx, annot_ref in enumerate(annots_raw):
            annot = resolve_obj(annot_ref)
            if annot is None:
                continue
            try:
                subtype = str(annot.get("/Subtype", ""))

                # Skip sticky notes — they are labels, not content rows
                if subtype == "/Text":
                    continue

                # Skip unsupported types
                if subtype not in MARKUP_SUBTYPES:
                    continue

                # Get label (comment text) for this annotation
                label = get_label(
                    annot_ref, annot,
                    markup_idx,
                    irt_map, text_pool,
                    used_text_indices,
                )

                # OCR the annotated region
                text = extract_markup_text(full_img, annot, page_height)

                annotations.append({
                    "page" : page_idx + 1,
                    "label": label,
                    "text" : text,
                })

            except Exception:
                continue

    doc_pdfium.close()
    doc_pk.close()
    return annotations


# ══════════════════════════════════════════════════════════════════════════════
# FILE OUTPUT
# ══════════════════════════════════════════════════════════════════════════════

def save_files(annotations: list) -> None:
    with open(CSV_PATH, "w", newline="", encoding="utf-8-sig") as f:
        w = csv.DictWriter(f, fieldnames=["page", "label", "text"])
        w.writeheader()
        w.writerows(annotations)
    with open(JSON_PATH, "w", encoding="utf-8") as f:
        json.dump(annotations, f, ensure_ascii=False, indent=2)


# ══════════════════════════════════════════════════════════════════════════════
# ROUTES
# ══════════════════════════════════════════════════════════════════════════════

@app.route("/", methods=["GET", "POST"])
def index():
    annotations = None
    error = None

    if request.method == "POST":
        if "pdf_file" not in request.files:
            return render_template("index.html", error="No file attached.")
        file = request.files["pdf_file"]
        if not file or file.filename == "":
            return render_template("index.html", error="No file selected.")
        if not allowed_file(file.filename):
            return render_template("index.html", error="Only PDF files are supported.")

        filename = secure_filename(file.filename)
        pdf_path = os.path.join(UPLOAD_DIR, filename)
        file.save(pdf_path)

        try:
            annotations = extract_annotations(pdf_path)
            save_files(annotations)
        except Exception as exc:
            return render_template("index.html", error=f"Processing error: {exc}")

        return render_template("index.html",
                               annotations=annotations,
                               filename=filename,
                               count=len(annotations))

    return render_template("index.html")


@app.route("/download/csv")
def download_csv():
    if not os.path.exists(CSV_PATH):
        flash("No CSV yet — upload a PDF first.")
        return redirect(url_for("index"))
    return send_file(CSV_PATH, as_attachment=True, download_name="annotations.csv")


@app.route("/download/json")
def download_json():
    if not os.path.exists(JSON_PATH):
        flash("No JSON yet — upload a PDF first.")
        return redirect(url_for("index"))
    return send_file(JSON_PATH, as_attachment=True, download_name="annotations.json")


if __name__ == "__main__":
    app.run(debug=True, port=5000)