diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index 94ba4be324..90c0d20316 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -33,9 +33,14 @@ from loggers import get_logger logger = get_logger(__name__) _PROMPT = ( - "Describe this figure in <=60 words. Focus on factual content " - "(axes, labels, captions, visible text, main objects). " - "Do not speculate beyond what is visible." + "This image is a region cropped from a PDF page that contains a " + "single figure (schematic, chart, diagram, table, photo, or " + "their combination). Describe the figure's structure and content " + "in <=80 words. Focus on factual visible content: axes, labels, " + "arrow labels, box labels, legends, visible text in the figure, " + "and what entities are connected to what. Do not speculate beyond " + "what is visible and do not describe the page header/footer or " + "body paragraphs." ) _MAX_NEW_TOKENS = 200 # Downscale large images so the base64 payload stays manageable; the chat diff --git a/studio/backend/core/rag/parsers/pdf.py b/studio/backend/core/rag/parsers/pdf.py index b9e6696a15..0f04def09b 100644 --- a/studio/backend/core/rag/parsers/pdf.py +++ b/studio/backend/core/rag/parsers/pdf.py @@ -58,45 +58,76 @@ def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult: doc.close() +# Pages smaller than this (in PDF points) are ignored as figure regions — +# bigger than a typical icon/glyph, smaller than a banner. +_MIN_FIGURE_PT = 60 +# 2× scale renders at 144 dpi (PDF default is 72 dpi). Enough resolution for +# the captioner to read axis labels, arrow text, and inset photos. +_RENDER_SCALE = 2.0 +# Expand the union bbox a few points so caption baselines / borders survive. +_FIGURE_MARGIN_PT = 8.0 + + def _extract_images_pymupdf(doc, pages: list[ParsedPage]) -> list[ParsedImage]: + """Render each page's figure region (vector drawings + raster sub-images) + as a single PNG. Vector schematics like Figure 1 (no embedded raster) + are visible to the captioner only via rendering — ``page.get_images`` + misses them entirely. We union all non-text geometry on a page into + one bbox; for academic papers this typically maps 1:1 to "the figure + on this page". + """ + import pymupdf + captions_by_page: dict[int, str] = { p.page_number: p.text for p in pages if p.page_number } out: list[ParsedImage] = [] for page_index in range(len(doc)): + page = doc[page_index] page_number = page_index + 1 try: - image_list = doc[page_index].get_images(full = True) + rects: list[pymupdf.Rect] = [] + for drawing in page.get_drawings() or []: + rect = drawing.get("rect") + if rect is not None: + rects.append(pymupdf.Rect(rect)) + for info in page.get_image_info(xrefs = True) or []: + bbox = info.get("bbox") + if bbox is not None: + rects.append(pymupdf.Rect(bbox)) except Exception: continue - for img_info in image_list: - xref = img_info[0] - try: - extracted = doc.extract_image(xref) - except Exception: - continue - image_bytes = extracted.get("image") - ext = (extracted.get("ext") or "png").lower() - mime = { - "png": "image/png", - "jpg": "image/jpeg", - "jpeg": "image/jpeg", - "gif": "image/gif", - "webp": "image/webp", - "bmp": "image/bmp", - "tiff": "image/tiff", - }.get(ext, f"image/{ext}") - if not image_bytes: - continue - caption = (captions_by_page.get(page_number, "") or "")[:1500] - out.append( - ParsedImage( - image_bytes = image_bytes, - mime_type = mime, - page_number = page_number, - nearest_caption = caption, - ) + if not rects: + continue + union = rects[0] + for r in rects[1:]: + union |= r + if union.width < _MIN_FIGURE_PT or union.height < _MIN_FIGURE_PT: + continue + # Expand and clip to page rect so we don't render past page edges. + union = pymupdf.Rect( + union.x0 - _FIGURE_MARGIN_PT, + union.y0 - _FIGURE_MARGIN_PT, + union.x1 + _FIGURE_MARGIN_PT, + union.y1 + _FIGURE_MARGIN_PT, + ) & page.rect + try: + matrix = pymupdf.Matrix(_RENDER_SCALE, _RENDER_SCALE) + pix = page.get_pixmap(clip = union, matrix = matrix, alpha = False) + png_bytes = pix.tobytes("png") + except Exception: + continue + if not png_bytes: + continue + caption = (captions_by_page.get(page_number, "") or "")[:1500] + out.append( + ParsedImage( + image_bytes = png_bytes, + mime_type = "image/png", + page_number = page_number, + nearest_caption = caption, ) + ) return out