super-document-processor 0.1.3__py3-none-any.whl

super_document_processor/utils/pdf_image_utils.py

@@ -0,0 +1,90 @@
+import hashlib
+import base64
+import logging
+import threading
+import concurrent.futures
+
+logger = logging.getLogger("super_document_processor.pdf_images")
+
+try:
+    from super_document_processor.utils.vlm import generate_image_description
+except ImportError:
+    from utils.vlm import generate_image_description
+
+
+def process_pdf_full_page_image(page, page_num, descriptions, model=None):
+    """Helper function to process the entire page as an image."""
+    pix = page.get_pixmap(dpi=200)
+    full_page_bytes = pix.tobytes("png")
+    raw_description = generate_image_description(full_page_bytes, model=model)
+    labeled_description = f"[Image description for full page {page_num + 1}]:\n{raw_description}"
+    descriptions.append((0, labeled_description))
+    return descriptions, True
+
+def process_image(seen_image_hashes, image_data_store, image_page_map, page, page_num, doc, lock=None, model=None):
+    if lock is None:
+        lock = threading.Lock()
+    descriptions = []
+    need_full_page_ocr = False
+
+    text = page.get_text().strip()
+    page_text_detected = bool(text)
+
+    # 1st Condition: If vector drawings exist and no text detected
+    drawings = page.get_drawings()
+    logger.debug(f"Vector paths found on page {page_num + 1}: {len(drawings)}")
+    
+    if len(drawings) > 0 and page_text_detected == False:
+        logger.debug(f"Processing full page image for description on page {page_num + 1}")
+        return process_pdf_full_page_image(page, page_num, descriptions, model=model)
+
+    # 2nd Condition: Otherwise, process individual images on the page
+    images = page.get_images(full=True)
+    logger.debug(f"Found {len(images)} images on page {page_num + 1}")
+
+    if len(images) > 5:
+        logger.debug(f"More than 5 images on page {page_num + 1}, processing full page image")
+        return process_pdf_full_page_image(page, page_num, descriptions, model=model)
+
+    # 3rd Condition: Otherwise, process each image individually (in parallel)
+    def process_single_image(img_index, img_info):
+        xref = img_info[0]
+        base_image = doc.extract_image(xref)
+        img_bytes = base_image["image"]
+
+        # Hash image bytes
+        img_hash = hashlib.sha256(img_bytes).hexdigest()
+
+        with lock:
+            if (page_num + 1) not in image_page_map[img_hash]:
+                image_page_map[img_hash].append(page_num + 1)
+            if img_hash not in image_data_store:
+                image_data_store[img_hash] = base64.b64encode(img_bytes).decode("utf-8")
+            description = seen_image_hashes.get(img_hash)
+
+        if description is None:
+            logger.debug(f"Generating LLM description for image {img_index + 1} on page {page_num + 1}")
+            description = generate_image_description(img_bytes, model=model)
+            with lock:
+                seen_image_hashes[img_hash] = description
+        else:
+            logger.debug(f"Reusing description for duplicate image on page {page_num + 1}")
+
+        labeled_description = f"[Image description for image {img_index + 1} on page {page_num + 1}]:\n{description}\n[Image description {img_index+1} end]"
+
+        image_rects = page.get_image_rects(img_info)
+        if not image_rects:
+            logger.debug(f"No bbox for image {img_index + 1} on page {page_num + 1}, skipping")
+            return None
+        img_bbox = image_rects[0]
+        return (img_bbox.y0, labeled_description)
+
+    workers = int(os.environ.get("MAX_WORKERS", 3))
+    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+        futures = [executor.submit(process_single_image, img_index, img_info) for img_index, img_info in enumerate(images)]
+        for future in concurrent.futures.as_completed(futures):
+            result = future.result()
+            if result:
+                descriptions.append(result)
+
+    return descriptions, need_full_page_ocr