#!/usr/bin/env python3 import os import re import gc import json import logging import fitz import boto3 import base64 import time import asyncio import tempfile import requests from io import BytesIO from typing import List, Dict, Any import torch import cv2 import numpy as np from google import genai from google.genai import types from magic_pdf.data.dataset import PymuDocDataset from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze from magic_pdf.data.data_reader_writer.base import DataWriter from table_row_extraction import TableExtractor logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) file_handler = logging.FileHandler("topic_extraction.log") file_handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s - %(message)s")) logger.addHandler(file_handler) _GEMINI_CLIENT = None def unify_whitespace(text: str) -> str: return re.sub(r"\s+", " ", text).strip() def find_all_occurrences(pdf_bytes: bytes, search_text: str) -> List[int]: doc = fitz.open(stream=pdf_bytes, filetype="pdf") st_norm = unify_whitespace(search_text) found = [] for i in range(doc.page_count): raw = doc[i].get_text("raw") norm = unify_whitespace(raw) if st_norm in norm: found.append(i) doc.close() return sorted(found) def create_subset_pdf(original_pdf_bytes: bytes, page_indices: List[int]) -> bytes: if not page_indices: raise ValueError("No page indices provided for subset creation.") doc = fitz.open(stream=original_pdf_bytes, filetype="pdf") new_doc = fitz.open() for p in sorted(set(page_indices)): if 0 <= p < doc.page_count: new_doc.insert_pdf(doc, from_page=p, to_page=p) else: logger.error(f"Page index {p} out of range (0..{doc.page_count - 1}).") raise ValueError(f"Page index {p} out of range.") subset_bytes = new_doc.tobytes() new_doc.close() doc.close() return subset_bytes class s3Writer: def __init__(self, ak: str, sk: str, bucket: str, endpoint_url: str): self.bucket = bucket self.client = boto3.client( 's3', aws_access_key_id=ak, aws_secret_access_key=sk, endpoint_url=endpoint_url ) def write(self, path: str, data: bytes) -> None: try: file_obj = BytesIO(data) self.client.upload_fileobj( file_obj, self.bucket, path ) logger.info(f"Uploaded to S3: {path}") except Exception as e: logger.error(f"Failed to upload to S3: {str(e)}") raise def preprocess_image(image_data: bytes, max_dim: int = 600, quality: int = 60) -> bytes: arr = np.frombuffer(image_data, np.uint8) img = cv2.imdecode(arr, cv2.IMREAD_COLOR) if img is not None: h, w, _ = img.shape if max(h, w) > max_dim: scale = max_dim / float(max(h, w)) new_w = int(w * scale) new_h = int(h * scale) img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA) encode_params = [int(cv2.IMWRITE_JPEG_QUALITY), quality] success, enc = cv2.imencode(".jpg", img, encode_params) if success: return enc.tobytes() return image_data def call_gemini_for_table_classification(image_data: bytes, api_key: str, max_retries: int = 1) -> str: """ Existing Gemini call to classify an image as TWO_COLUMN, THREE_COLUMN, or NO_TABLE. """ for attempt in range(max_retries + 1): try: prompt = """You are given an image. Determine if it shows a table that has exactly 2 or 3 columns. The three-column 'table' image includes such key features: - Three columns header - Headers like 'Topics', 'Content', 'Guidelines' - Possibly sections (e.g. 8.4, 9.1) The two-column 'table' image includes such key features: - Two columns - Headers like 'Subject content' and 'Additional information' - Possibly sections (e.g. 2.1, 3.4) If the image is a relevant table with 2 columns, respond with 'TWO_COLUMN'. If the image is a relevant table with 3 columns, respond with 'THREE_COLUMN'. If the image does not show a table at all, respond with 'NO_TABLE'. Return only one of these exact labels. """ global _GEMINI_CLIENT if _GEMINI_CLIENT is None: _GEMINI_CLIENT = genai.Client(api_key=api_key) client = _GEMINI_CLIENT resp = client.models.generate_content( model="gemini-2.0-flash", contents=[ { "parts": [ {"text": prompt}, { "inline_data": { "mime_type": "image/jpeg", "data": base64.b64encode(image_data).decode('utf-8') } } ] } ], config=types.GenerateContentConfig(temperature=0.0) ) if resp and resp.text: classification = resp.text.strip().upper() if "THREE" in classification: return "THREE_COLUMN" elif "TWO" in classification: return "TWO_COLUMN" return "NO_TABLE" except Exception as e: logger.error(f"Gemini table classification error: {e}") if "503" in str(e): return "NO_TABLE" if attempt < max_retries: time.sleep(0.5) else: return "NO_TABLE" async def classify_image_async(image_data: bytes, api_key: str, max_retries: int = 1) -> str: loop = asyncio.get_event_loop() preprocessed = preprocess_image(image_data) return await loop.run_in_executor(None, call_gemini_for_table_classification, preprocessed, api_key, max_retries) def call_gemini_for_subtopic_identification_image(image_data: bytes, api_key: str, max_retries: int = 1) -> dict: for attempt in range(max_retries + 1): try: prompt = """ You are given an image from an educational curriculum specification. The image may contain either: 1) A main topic heading in the format: " ", for example "2 Algebra and functions continued". 2) A subtopic heading in the format ".", for example "2.5", "2.6", or "3.4". 3) Possibly no relevant text at all. Your task: 1. If the cell shows a main topic, extract the topic name (e.g. "2 Algebra and functions") and place it in the JSON key "title". 2. If the cell shows one or more subtopic numbers (e.g. "2.5", "2.6"), collect them in the JSON key "subtopics" as an array of strings. 3. If neither a main topic nor subtopic is detected, return empty values. Output only valid JSON in this exact structure, with no extra text or explanation: Output only valid JSON in this exact structure, with no extra text or explanation: { "title": "...", "subtopics": [...] } Where: - "title" is the recognized main topic (if any). Otherwise, an empty string. - "subtopics" is an array of recognized subtopic numbers (e.g. ["2.5", "2.6"]). Otherwise, an empty array. Examples: 1. If the image text is "2 Algebra and functions continued", return: { "title": "2 Algebra and functions continued", "subtopics": [] } 2. If the image text is "2.5 Solve linear and quadratic inequalities ...", return: { "title": "", "subtopics": ["2.5"] } 3. If the image text is "2.6 Manipulate polynomials algebraically ...", return: { "title": "", "subtopics": ["2.6"] } If you cannot recognize any text matching these patterns, or if nothing is found, return: { "title": "", "subtopics": [] } """ global _GEMINI_CLIENT if _GEMINI_CLIENT is None: _GEMINI_CLIENT = genai.Client(api_key=api_key) client = _GEMINI_CLIENT resp = client.models.generate_content( model="gemini-2.0-flash", contents=[ { "parts": [ {"text": prompt}, { "inline_data": { "mime_type": "image/jpeg", "data": base64.b64encode(image_data).decode("utf-8") } } ] } ], config=types.GenerateContentConfig(temperature=0.0) ) # logger.info(f"Gemini subtopic extraction raw response: {resp.text if resp and resp.text else 'None'}") if not resp or not resp.text: logger.warning("Gemini returned an empty response for subtopic extraction.") return {"title": "", "subtopics": []} raw = resp.text.strip() raw = raw.replace("```json", "").replace("```", "").strip() data = json.loads(raw) title = data.get("title", "") subtopics = data.get("subtopics", []) if not isinstance(subtopics, list): subtopics = [] return {"title": title, "subtopics": subtopics} except Exception as e: logger.error(f"Gemini subtopic identification error on attempt {attempt}: {e}") if attempt < max_retries: time.sleep(0.5) else: return {"title": "", "subtopics": []} return {"title": "", "subtopics": []} class S3ImageWriter(DataWriter): def __init__(self, s3_writer: s3Writer, base_path: str, gemini_api_key: str): self.s3_writer = s3_writer self.base_path = base_path if base_path.endswith("/") else base_path + "/" self.gemini_api_key = gemini_api_key self.descriptions = {} self._img_count = 0 self.extracted_tables = {} self.extracted_subtopics = {} def write(self, path: str, data: bytes) -> None: self._img_count += 1 unique_id = f"img_{self._img_count}.jpg" s3_key = f"{self.base_path}{unique_id}" self.s3_writer.write(s3_key, data) self.descriptions[path] = { "data": data, "s3_path": s3_key, "table_classification": "NO_TABLE", "final_alt": "" } async def post_process_async(self, key: str, md_content: str) -> str: logger.info("Classifying images to detect tables.") tasks = { p: asyncio.create_task(classify_image_async(info["data"], self.gemini_api_key)) for p, info in self.descriptions.items() } results = await asyncio.gather(*tasks.values(), return_exceptions=True) for p, result in zip(tasks.keys(), results): if isinstance(result, Exception): logger.error(f"Table classification error for {p}: {result}") self.descriptions[p]['table_classification'] = "NO_TABLE" else: self.descriptions[p]['table_classification'] = result for p, info in self.descriptions.items(): cls = info['table_classification'] if cls == "TWO_COLUMN": info['final_alt'] = "HAS TO BE PROCESSED - two column table" elif cls == "THREE_COLUMN": info['final_alt'] = "HAS TO BE PROCESSED - three column table" else: info['final_alt'] = "NO_TABLE image" md_content = md_content.replace(f"![]({key}{p})", f"![{info['final_alt']}]({info['s3_path']})") md_content = await self._process_table_images_in_markdown(key, md_content) # Filter final lines to keep only lines with images final_lines = [ line.strip() for line in md_content.split("\n") if re.match(r"^\!\[.*\]\(.*\)", line.strip()) ] return "\n".join(final_lines) async def _process_table_images_in_markdown(self, key: str, md_content: str) -> str: pat = r"!\[HAS TO BE PROCESSED - (two|three) column table\]\(([^)]+)\)" matches = re.findall(pat, md_content, flags=re.IGNORECASE) if not matches: return md_content for (col_type, s3_key) in matches: logger.info(f"Processing table image: {s3_key}, columns={col_type}") img_data = None for desc in self.descriptions.values(): if desc.get("s3_path") == s3_key: img_data = desc.get("data") break if img_data is None: logger.warning(f"No image data found for S3 key {s3_key}. Skipping.") continue # Write temporary file for processing. with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: temp_file.write(img_data) temp_path = temp_file.name try: if col_type.lower() == 'two': extractor = TableExtractor( skip_header=True, merge_two_col_rows=True, enable_subtopic_merge=True, subtopic_threshold=0.2 ) else: extractor = TableExtractor( skip_header=True, merge_two_col_rows=False, enable_subtopic_merge=False, subtopic_threshold=0.2 ) row_boxes = extractor.process_image(temp_path) # logger.info(f"Extracted {len(row_boxes)} rows from {temp_path}") # for i, row in enumerate(row_boxes): # logger.info(f"Row {i} has {len(row)} cells") out_folder = temp_path + "_rows" os.makedirs(out_folder, exist_ok=True) # out_folder = os.path.join(os.path.dirname(temp_path), os.path.basename(temp_path) + "_rows") # os.makedirs(out_folder, exist_ok=True) extractor.save_extracted_cells(temp_path, row_boxes, out_folder) #just to print structure how cells are saved and named for each table image # logger.info(f"Files in {out_folder}:") # for root, dirs, files in os.walk(out_folder): # logger.info(f"{root}: {files}") recognized_main_topic = "" main_topic_image_key = None recognized_subtopics = [] # Loop over each cell image. for i, row in enumerate(row_boxes): row_dir = os.path.join(out_folder, f"row_{i}") for j, _ in enumerate(row): cell_path = os.path.join(row_dir, f"col_{j}.png") if not os.path.isfile(cell_path): alternative_path = os.path.join(row_dir, f"col_{j}.jpg") if os.path.isfile(alternative_path): cell_path = alternative_path else: logger.warning(f"Cell image not found: {cell_path}") continue with open(cell_path, "rb") as cf: cell_image_data = cf.read() # Save cell image to S3. cell_key = f"{self.base_path}cells/{os.path.basename(s3_key)}_r{i}_c{j}.png" self.s3_writer.write(cell_key, cell_image_data) info = call_gemini_for_subtopic_identification_image(cell_image_data, self.gemini_api_key) # logger.info(f"Gemini subtopic extraction result for cell {cell_path}: {info}") if info["title"] and not recognized_main_topic: recognized_main_topic = info["title"] main_topic_image_key = cell_key for st in info["subtopics"]: recognized_subtopics.append({ "title": st, "contents": [{"type": "image", "key": cell_key}], "children": [] }) final_json = { "title": recognized_main_topic, "contents": [], "children": recognized_subtopics } if main_topic_image_key: final_json["contents"].append({"type": "image", "key": main_topic_image_key}) # Save the final JSON. self.extracted_subtopics[s3_key] = final_json # Optionally, create a snippet to replace the markdown line. snippet = ["**Extracted table cells:**"] for i, row in enumerate(row_boxes): for j, _ in enumerate(row): snippet.append(f"![Row {i} Col {j}]({self.base_path}cells/{os.path.basename(s3_key)}_r{i}_c{j}.jpg)") new_snip = "\n".join(snippet) old_line = f"![HAS TO BE PROCESSED - {col_type} column table]({s3_key})" md_content = md_content.replace(old_line, new_snip) except Exception as e: logger.error(f"Error processing table image {s3_key}: {e}") finally: os.remove(temp_path) return md_content def post_process(self, key: str, md_content: str) -> str: return asyncio.run(self.post_process_async(key, md_content)) class LocalImageWriter(DataWriter): def __init__(self, output_folder: str, gemini_api_key: str): self.output_folder = output_folder os.makedirs(self.output_folder, exist_ok=True) self.descriptions = {} self._img_count = 0 self.gemini_api_key = gemini_api_key self.extracted_tables = {} def write(self, path: str, data: bytes) -> None: self._img_count += 1 unique_id = f"img_{self._img_count}.jpg" self.descriptions[path] = { "data": data, "relative_path": unique_id, "table_classification": "NO_TABLE", "final_alt": "" } image_path = os.path.join(self.output_folder, unique_id) with open(image_path, "wb") as f: f.write(data) async def post_process_async(self, key: str, md_content: str) -> str: logger.info("Classifying images to detect tables.") tasks = [] for p, info in self.descriptions.items(): tasks.append((p, classify_image_async(info["data"], self.gemini_api_key))) for p, task in tasks: try: classification = await task self.descriptions[p]['table_classification'] = classification except Exception as e: logger.error(f"Table classification error: {e}") self.descriptions[p]['table_classification'] = "NO_TABLE" for p, info in self.descriptions.items(): cls = info['table_classification'] if cls == "TWO_COLUMN": info['final_alt'] = "HAS TO BE PROCESSED - two column table" elif cls == "THREE_COLUMN": info['final_alt'] = "HAS TO BE PROCESSED - three column table" else: info['final_alt'] = "NO_TABLE image" md_content = md_content.replace(f"![]({key}{p})", f"![{info['final_alt']}]({info['relative_path']})") md_content = self._process_table_images_in_markdown(md_content) final_lines = [] for line in md_content.split("\n"): if re.match(r"^\!\[.*\]\(.*\)", line.strip()): final_lines.append(line.strip()) return "\n".join(final_lines) def _process_table_images_in_markdown(self, md_content: str) -> str: pat = r"!\[HAS TO BE PROCESSED - (two|three) column table\]\(([^)]+)\)" matches = re.findall(pat, md_content, flags=re.IGNORECASE) if not matches: return md_content for (col_type, image_id) in matches: logger.info(f"Processing table image => {image_id}, columns={col_type}") temp_path = os.path.join(self.output_folder, image_id) desc_item = None for k, val in self.descriptions.items(): if val["relative_path"] == image_id: desc_item = val break if not desc_item: logger.warning(f"No matching image data for {image_id}, skipping extraction.") continue if not os.path.exists(temp_path): with open(temp_path, "wb") as f: f.write(desc_item["data"]) try: if col_type.lower() == 'two': #check for table_row_extr script for more details extractor = TableExtractor( skip_header=True, merge_two_col_rows=True, enable_subtopic_merge=True, subtopic_threshold=0.2 ) else: extractor = TableExtractor( skip_header=True, merge_two_col_rows=False, enable_subtopic_merge=False, subtopic_threshold=0.2 ) row_boxes = extractor.process_image(temp_path) out_folder = temp_path + "_rows" os.makedirs(out_folder, exist_ok=True) extractor.save_extracted_cells(temp_path, row_boxes, out_folder) # List all extracted cell images relative to the output folder. extracted_cells = [] for root, dirs, files in os.walk(out_folder): for file in files: rel_path = os.path.relpath(os.path.join(root, file), self.output_folder) extracted_cells.append(rel_path) # Save mapping for testing. self.extracted_tables[image_id] = extracted_cells snippet = ["**Extracted table cells:**"] for i, row in enumerate(row_boxes): row_dir = os.path.join(out_folder, f"row_{i}") for j, _ in enumerate(row): cell_file = f"col_{j}.jpg" cell_path = os.path.join(row_dir, cell_file) relp = os.path.relpath(cell_path, self.output_folder) snippet.append(f"![Row {i} Col {j}]({relp})") new_snip = "\n".join(snippet) old_line = f"![HAS TO BE PROCESSED - {col_type} column table]({image_id})" md_content = md_content.replace(old_line, new_snip) except Exception as e: logger.error(f"Error processing table image {image_id}: {e}") finally: if os.path.exists(temp_path): os.remove(temp_path) return md_content def post_process(self, key: str, md_content: str) -> str: return asyncio.run(self.post_process_async(key, md_content)) class GeminiTopicExtractor: def __init__(self, api_key: str = None, num_pages: int = 14): self.api_key = api_key or os.getenv("GEMINI_API_KEY", "") self.num_pages = num_pages def extract_subtopics(self, pdf_path: str) -> Dict[str, List[int]]: first_pages_text = self._read_first_pages_raw(pdf_path, self.num_pages) if not first_pages_text.strip(): logger.error("No text from first pages => cannot extract subtopics.") return {} prompt = f""" You have the first pages of a PDF specification, including a table of contents. Instructions: 1. Identify the 'Contents' section listing all topics, subtopics, and their corresponding pages. 2. Identify the major academic subtopics (common desired topic names "Paper X", "Theme X", "Content of X", "AS Unit X", "A2 Unit X", or similar headings). 3. For each subtopic, give the range of pages [start_page, end_page] (1-based) from the table of contents. 4. Output only valid JSON of the form: {{ "Subtopic A": [start_page, end_page], "Subtopic B": [start_page, end_page] }} 5. If you can't find any subtopics, return an empty JSON. Important notes: - The correct "end_page" must be the page number of the next topic or subtopic minus 1. - The final output must be valid JSON only, with no extra text or code blocks. Examples: 1. Given this table of contents: 1 Introduction – 2 Why choose Edexcel A Level Mathematics? - 2 Supporting you in planning and implementing this qualification - 3 Qualification at a glance - 5 2 Subject content and assessment information – 7 Paper 1 and Paper 2: Pure Mathematics - 11 Paper 3: Statistics and Mechanics - 30 Assessment Objectives - 40 3 Administration and general information – 42 Entries - 42 Access arrangements, reasonable adjustments, special consideration and malpractice - 42 Student recruitment and progression - 45 Appendix 1: Formulae – 49 Appendix 2: Notation – 53 Appendix 3: Use of calculators – 59 Appendix 4: Assessment Objectives – 60 Appendix 5: The context for the development of this qualification – 62 Appendix 6: Transferable skills – 64 Appendix 7: Level 3 Extended Project qualification – 65 Appendix 8: Codes – 67 The correct output should be: {{ "Paper 1 and Paper 2: Pure Mathematics": [11, 29], "Paper 3: Statistics and Mechanics": [30, 42] }} 2. Given this table of contents: Qualification at a glance – 1 Assessment Objectives and weightings - 4 Knowledge, skills and understanding – 5 Theme 1: Introduction to markets and market failure - 5 Theme 2: The UK economy – performance and policies - 11 Theme 3: Business behaviour and the labour market - 21 Theme 4: A global perspective - 29 Assessment – 39 Assessment summary - 39 Assessment objectives - 41 Assessment overview - 42 Breakdown of assessment objectives - 42 Synoptic assessment - 43 Discount code and performance tables - 43 Access arrangements, reasonable adjustments and special consideration - 44 Malpractice - 45 Equality Act 2010 and Pearson equality policy - 45 Synoptic assessment - 46 Awarding and reporting - 47 Other information – 49 Student recruitment -49 Prior learning and other requirements -49 Progression - 49 Appendix 1: Transferable skills – 53 Appendix 2: Level 3 Extended Project qualification – 55 Appendix 3: Quantitative skills – 59 Appendix 4: Codes – 61 Appendix 5: Index – 63 The correct output should be: {{ "Theme 1: Introduction to markets and market failure": [5, 10], "Theme 2: The UK economy – performance and policies": [11, 20], "Theme 3: Business behaviour and the labour market": [21, 28], "Theme 4: A global perspective": [29, 38] }} 3. You might also see sections like: 2.1 AS Unit 1 11 2.2 AS Unit 2 18 2.3 A2 Unit 3 24 2.4 A2 Unit 4 31 In that scenario, your output might look like: {{ "2.1 AS Unit 1": [11, 17], "2.2 AS Unit 2": [18, 23], "2.3 A2 Unit 3": [24, 30], "2.4 A2 Unit 4": [31, 35] }} 4. Another example might list subtopics: 3.1 Overarching themes 11 3.2 A: Proof 12 3.3 B: Algebra and functions 13 3.4 C: Coordinate geometry in the ( x , y ) plane 14 3.5 D: Sequences and series 15 3.6 E: Trigonometry 16 3.7 F: Exponentials and logarithms 17 3.8 G: Differentiation 18 3.9 H: Integration 19 3.10 I: Numerical methods 20 3.11 J: Vectors 20 3.12 K: Statistical sampling 21 3.13 L: Data presentation and interpretation 21 3.14 M: Probability 22 3.15 N: Statistical distributions 23 3.16 O: Statistical hypothesis testing 23 3.17 P: Quantities and units in mechanics 24 3.18 Q: Kinematics 24 3.19 R: Forces and Newton’s laws 24 3.20 S: Moments 25 3.21 Use of data in statistics 26 Here the correct output might look like: {{ "A: Proof": [12, 12], "B: Algebra and functions": [13, 13], ... }} Now, extract topics from this text: {first_pages_text} """ global _GEMINI_CLIENT if _GEMINI_CLIENT is None: _GEMINI_CLIENT = genai.Client(api_key=self.api_key) client = _GEMINI_CLIENT try: response = client.models.generate_content( model="gemini-2.0-flash", contents=[prompt], config=types.GenerateContentConfig(temperature=0.0) ) if not response or not response.text: logger.warning("No text from LLM => returning empty subtopics.") return {} raw_json = response.text.strip() cleaned = raw_json.replace("```json", "").replace("```", "") try: data = json.loads(cleaned) except Exception as json_err: logger.error(f"JSON parsing error: {json_err}") return {} final_dict = {} found_sub_dict = None for k, v in data.items(): if isinstance(v, dict): found_sub_dict = v break if found_sub_dict is not None: for subk, rng in found_sub_dict.items(): if isinstance(rng, list) and len(rng) == 2: final_dict[subk] = rng else: for subk, rng in data.items(): if isinstance(rng, list) and len(rng) == 2: final_dict[subk] = rng return final_dict except Exception as e: logger.error(f"Gemini subtopic extraction error: {e}") return {} def _read_first_pages_raw(self, pdf_path: str, num_pages: int) -> str: text_parts = [] try: if pdf_path.startswith("http://") or pdf_path.startswith("https://"): response = requests.get(pdf_path) if response.status_code != 200: logger.error("Failed to download PDF from %s. Status code: %d", pdf_path, response.status_code) return "" pdf_bytes = response.content else: with open(pdf_path, "rb") as f: pdf_bytes = f.read() doc = fitz.open(stream=pdf_bytes, filetype="pdf") pages_to_read = min(num_pages, doc.page_count) for i in range(pages_to_read): raw_text = doc[i].get_text("raw") text_parts.append(raw_text) doc.close() except Exception as e: logger.error(f"Could not open PDF: {e}") return "\n".join(text_parts) class MineruNoTextProcessor: def __init__(self, output_folder: str, gemini_api_key: str): self.output_folder = output_folder os.makedirs(self.output_folder, exist_ok=True) self.layout_model = "doclayout_yolo" self.formula_enable = True self.table_enable = False self.language = "en" self.subtopic_extractor = GeminiTopicExtractor(api_key=gemini_api_key, num_pages=20) self.gemini_api_key = gemini_api_key or os.getenv("GEMINI_API_KEY", "") self.use_s3 = True self.s3_writer = s3Writer( ak=os.getenv("S3_ACCESS_KEY"), sk=os.getenv("S3_SECRET_KEY"), bucket="quextro-resources", endpoint_url=os.getenv("S3_ENDPOINT") ) def cleanup_gpu(self): try: gc.collect() torch.cuda.empty_cache() logger.info("GPU memory cleaned up.") except Exception as e: logger.error(f"Error during GPU cleanup: {e}") def process(self, pdf_path: str) -> Dict[str, Any]: logger.info(f"Processing PDF: {pdf_path}") try: # Possibly call subtopic_extractor on first pages to find subtopics in the PDF as a whole subtopics = self.subtopic_extractor.extract_subtopics(pdf_path) logger.info(f"Gemini returned subtopics: {subtopics}") if pdf_path.startswith("http://") or pdf_path.startswith("https://"): response = requests.get(pdf_path) if response.status_code != 200: logger.error("Failed to download PDF from %s. Status code: %d", pdf_path, response.status_code) raise Exception(f"Failed to download PDF: {pdf_path}") pdf_bytes = response.content logger.info("Downloaded %d bytes for pdf_url='%s'", len(pdf_bytes), pdf_path) else: with open(pdf_path, "rb") as f: pdf_bytes = f.read() logger.info("Loaded %d bytes from local file '%s'", len(pdf_bytes), pdf_path) doc = fitz.open(stream=pdf_bytes, filetype="pdf") total_pages = doc.page_count doc.close() # Decide which pages to process final_pages = set() if not subtopics: # fallback final_pages = set(range(total_pages)) else: offset_candidates = [] for subname, rng in subtopics.items(): start_p, _ = rng occs = find_all_occurrences(pdf_bytes, subname) for p in occs: candidate = p - (start_p - 1) if candidate > 0: offset_candidates.append(candidate) if offset_candidates: try: from statistics import mode global_offset = mode(offset_candidates) except: from statistics import median global_offset = int(median(offset_candidates)) else: global_offset = 0 logger.info(f"Computed global offset: {global_offset}") for subname, rng in subtopics.items(): if not (isinstance(rng, list) and len(rng) == 2): continue start_p, end_p = rng if start_p > end_p: continue s0 = (start_p - 1) + global_offset e0 = (end_p - 1) + global_offset for pp in range(s0, e0 + 1): final_pages.add(pp) if not final_pages: final_pages = set(range(total_pages)) logger.info(f"Processing pages (0-based): {sorted(final_pages)}") subset_pdf_bytes = create_subset_pdf(pdf_bytes, sorted(final_pages)) # 4) Analyze and produce markdown dataset = PymuDocDataset(subset_pdf_bytes) inference = doc_analyze( dataset, ocr=True, lang=self.language, layout_model=self.layout_model, formula_enable=self.formula_enable, table_enable=self.table_enable ) #S3 writer = S3ImageWriter(self.s3_writer, "/topic-extraction", self.gemini_api_key) #local # writer = LocalImageWriter(self.output_folder, self.gemini_api_key) md_prefix = "/topic-extraction/" pipe_result = inference.pipe_ocr_mode(writer, lang=self.language) md_content = pipe_result.get_markdown(md_prefix) final_markdown = writer.post_process(md_prefix, md_content) subtopic_list = list(writer.extracted_subtopics.values()) out_path = os.path.join(self.output_folder, "final_subtopics.json") with open(out_path, "w", encoding="utf-8") as f: json.dump(subtopic_list, f, indent=2) logger.info(f"Final subtopics JSON saved locally at {out_path}") return { "final_markdown": final_markdown, "subtopics_extracted": subtopic_list } finally: self.cleanup_gpu() if __name__ == "__main__": input_pdf = "/home/user/app/input_output/a-level-pearson-mathematics-specification.pdf" output_dir = "/home/user/app/pearson_json" gemini_key = os.getenv("GEMINI_API_KEY", "AIzaSyDtoakpXa2pjJwcQB6TJ5QaXHNSA5JxcrU") try: processor = MineruNoTextProcessor(output_folder=output_dir, gemini_api_key=gemini_key) result = processor.process(input_pdf) logger.info("Processing completed successfully.") except Exception as e: logger.error(f"Processing failed: {e}")