Spaces:
Sleeping
Sleeping
File size: 5,639 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """PDF page layout segmentation via PyMuPDF (fitz).
Locates non-text regions a student can't otherwise select -> figures, plots,
diagrams, tables -> and returns pixel-accurate normalized bounding boxes plus a
PNG crop of each. Vision understanding (caption / LaTeX / table extraction) is a
separate step (SensesAgent.describe_region), so this module stays pure geometry
and is unit-testable without any model call.
Born-digital PDFs get exact geometry here. Scanned/image-only pages have no
embedded structure; `page_has_text()` lets the caller fall back to whole-page
vision instead.
"""
from __future__ import annotations
import base64
from typing import Any, Dict, List, Optional, Tuple
# Skip regions smaller than this fraction of the page (noise / tiny glyph runs).
_MIN_AREA_FRAC = 0.003
# Cap regions per page to bound downstream vision calls.
_MAX_REGIONS = 8
_CROP_ZOOM = 2.0
def _open(pdf_bytes: Optional[bytes], file_path: Optional[str]):
import pymupdf as fitz
if pdf_bytes is not None:
return fitz.open(stream=pdf_bytes, filetype="pdf")
return fitz.open(file_path)
def page_has_text(pdf_bytes: Optional[bytes], file_path: Optional[str], page_number: int) -> bool:
doc = _open(pdf_bytes, file_path)
try:
return bool(doc[page_number].get_text("text").strip())
finally:
doc.close()
def _norm(rect, pw: float, ph: float) -> Dict[str, float]:
return {
"x": max(0.0, rect[0] / pw),
"y": max(0.0, rect[1] / ph),
"w": min(1.0, (rect[2] - rect[0]) / pw),
"h": min(1.0, (rect[3] - rect[1]) / ph),
}
def crop_page_region(
pdf_bytes: Optional[bytes] = None,
file_path: Optional[str] = None,
page_number: int = 0,
bbox_norm: Dict[str, float] = None,
) -> str:
"""Extract a base64 PNG crop of the given normalized bounding box."""
import pymupdf as fitz
doc = _open(pdf_bytes, file_path)
try:
page = doc[page_number]
pw, ph = page.rect.width, page.rect.height
# Convert normalized bbox to absolute coordinates
x0 = bbox_norm["x"] * pw
y0 = bbox_norm["y"] * ph
x1 = (bbox_norm["x"] + bbox_norm["w"]) * pw
y1 = (bbox_norm["y"] + bbox_norm["h"]) * ph
rect = fitz.Rect(x0, y0, x1, y1)
pix = page.get_pixmap(clip=rect, matrix=fitz.Matrix(_CROP_ZOOM, _CROP_ZOOM))
return base64.b64encode(pix.tobytes("png")).decode("ascii")
finally:
doc.close()
def segment_page(
pdf_bytes: Optional[bytes] = None,
file_path: Optional[str] = None,
page_number: int = 0,
) -> Tuple[List[Dict[str, Any]], Tuple[float, float]]:
"""Return (regions, (page_width, page_height)).
Each region: {id, type, bbox_norm{x,y,w,h}, crop_base64}.
"""
import pymupdf as fitz
doc = _open(pdf_bytes, file_path)
try:
page = doc[page_number]
pr = page.rect
pw, ph = pr.width, pr.height
page_area = max(1.0, pw * ph)
candidates: List[Tuple[Any, str]] = []
# 1. Embedded raster images (figures/plots)
try:
for info in page.get_image_info():
bbox = info.get("bbox")
if bbox:
candidates.append((fitz.Rect(bbox), "figure"))
except Exception:
pass
# 2. Vector-drawing clusters (line plots, schematic diagrams)
try:
try:
rects = page.cluster_drawings(x_tolerance=20.0, y_tolerance=20.0)
except TypeError:
rects = page.cluster_drawings()
for rect in rects:
candidates.append((fitz.Rect(rect), "diagram"))
except Exception:
pass
regions: List[Dict[str, Any]] = []
kept_rects: List[Tuple[Any, str]] = []
# Sort all candidates by area (largest first)
candidates.sort(key=lambda c: c[0].get_area(), reverse=True)
for rect, rtype in candidates:
if rect.is_empty or rect.width <= 0 or rect.height <= 0:
continue
if rect.get_area() < _MIN_AREA_FRAC * page_area:
continue
overlap = False
for k_rect, k_type in kept_rects:
inter = rect & k_rect
# If same type, dedupe heavy overlap
if rtype == k_type:
if not inter.is_empty and inter.get_area() > 0.6 * rect.get_area():
overlap = True
break
# If diagram heavily overlaps a figure, skip it!
elif rtype == "diagram" and k_type == "figure":
if not inter.is_empty and inter.get_area() > 0.6 * rect.get_area():
overlap = True
break
if overlap:
continue
kept_rects.append((rect, rtype))
try:
mat = fitz.Matrix(_CROP_ZOOM, _CROP_ZOOM)
pix = page.get_pixmap(clip=rect, matrix=mat)
crop_b64 = base64.b64encode(pix.tobytes("png")).decode()
except Exception:
continue # skip regions that fail to crop
regions.append({
"id": f"r{len(regions)}",
"type": rtype,
"bbox_norm": _norm(rect, pw, ph),
"crop_base64": crop_b64,
})
# Removed _MAX_REGIONS cap to allow the page to be fully divided into sections
return regions, (pw, ph)
finally:
doc.close()
|