# import base64 # from PIL import Image # import re # import fitz # PyMuPDF # import numpy as np # import cv2 # import torch # import torch.serialization # import os # import time # from typing import Optional, Tuple, List, Dict, Any # from ultralytics import YOLO # import logging # import gradio as gr # import shutil # import tempfile # import io # # ============================================================================ # # --- Global Patches and Setup --- # # ============================================================================ # # Patch torch.load to prevent weights_only error with older models # _original_torch_load = torch.load # def patched_torch_load(*args, **kwargs): # kwargs["weights_only"] = False # return _original_torch_load(*args, **kwargs) # torch.load = patched_torch_load # logging.basicConfig(level=logging.WARNING) # # ============================================================================ # # --- CONFIGURATION AND CONSTANTS --- # # ============================================================================ # WEIGHTS_PATH = 'best.pt' # SCALE_FACTOR = 2.0 # # OUTPUT_DIR = "yolo_extracted_regions" # # OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "yolo_extracted_regions") # from transformers import TrOCRProcessor # from optimum.onnxruntime import ORTModelForVision2Seq # MODEL_NAME = 'breezedeus/pix2text-mfr-1.5' # processor = TrOCRProcessor.from_pretrained(MODEL_NAME) # ort_model = ORTModelForVision2Seq.from_pretrained(MODEL_NAME, use_cache=False) # # Detection parameters # CONF_THRESHOLD = 0.2 # TARGET_CLASSES = ['figure', 'equation'] # IOU_MERGE_THRESHOLD = 0.4 # IOA_SUPPRESSION_THRESHOLD = 0.7 # # Global counters (Reset per run) # GLOBAL_FIGURE_COUNT = 0 # GLOBAL_EQUATION_COUNT = 0 # # ============================================================================ # # --- BOX COMBINATION LOGIC (Retained for detection accuracy) --- # # ============================================================================ # def calculate_iou(box1, box2): # x1_a, y1_a, x2_a, y2_a = box1 # x1_b, y1_b, x2_b, y2_b = box2 # x_left = max(x1_a, x1_b) # y_top = max(y1_a, y1_b) # x_right = min(x2_a, x2_b) # y_bottom = min(y2_a, y2_b) # intersection_area = max(0, x_right - x_left) * max(0, y_bottom - y_top) # box_a_area = (x2_a - x1_a) * (y2_a - y1_a) # box_b_area = (x2_b - x1_b) * (y2_b - y1_b) # union_area = float(box_a_area + box_b_area - intersection_area) # return intersection_area / union_area if union_area > 0 else 0 # def filter_nested_boxes(detections, ioa_threshold=0.80): # if not detections: return [] # for d in detections: # x1, y1, x2, y2 = d['coords'] # d['area'] = (x2 - x1) * (y2 - y1) # detections.sort(key=lambda x: x['area'], reverse=True) # keep_indices = [] # is_suppressed = [False] * len(detections) # for i in range(len(detections)): # if is_suppressed[i]: continue # keep_indices.append(i) # box_a = detections[i]['coords'] # for j in range(i + 1, len(detections)): # if is_suppressed[j]: continue # box_b = detections[j]['coords'] # x_left = max(box_a[0], box_b[0]) # y_top = max(box_a[1], box_b[1]) # x_right = min(box_a[2], box_b[2]) # y_bottom = min(box_a[3], box_b[3]) # intersection = max(0, x_right - x_left) * max(0, y_bottom - y_top) # area_b = detections[j]['area'] # if area_b > 0 and intersection / area_b > ioa_threshold: # is_suppressed[j] = True # return [detections[i] for i in keep_indices] # def merge_overlapping_boxes(detections, iou_threshold): # if not detections: return [] # detections.sort(key=lambda d: d['conf'], reverse=True) # merged_detections = [] # is_merged = [False] * len(detections) # for i in range(len(detections)): # if is_merged[i]: continue # current_box = detections[i]['coords'] # current_class = detections[i]['class'] # merged_x1, merged_y1, merged_x2, merged_y2 = current_box # for j in range(i + 1, len(detections)): # if is_merged[j] or detections[j]['class'] != current_class: continue # other_box = detections[j]['coords'] # iou = calculate_iou(current_box, other_box) # if iou > iou_threshold: # merged_x1 = min(merged_x1, other_box[0]) # merged_y1 = min(merged_y1, other_box[1]) # merged_x2 = max(merged_x2, other_box[2]) # merged_y2 = max(merged_y2, other_box[3]) # is_merged[j] = True # merged_detections.append({ # 'coords': (merged_x1, merged_y1, merged_x2, merged_y2), # 'y1': merged_y1, 'class': current_class, 'conf': detections[i]['conf'] # }) # return merged_detections # # ============================================================================ # # --- UTILITY FUNCTIONS --- # # ============================================================================ # def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray: # """Converts a PyMuPDF Pixmap to a NumPy array for OpenCV/YOLO.""" # img = np.frombuffer(pix.samples, dtype=np.uint8).reshape( # (pix.h, pix.w, pix.n) # ) # if pix.n == 4: # img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) # elif pix.n == 1: # img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # return img # def run_yolo_detection_and_count( # image: np.ndarray, model: YOLO, page_num: int # ) -> Tuple[int, int, List[Dict[str, str]]]: # global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT # yolo_detections = [] # page_equations = 0 # page_figures = 0 # detected_items = [] # try: # results = model.predict(image, conf=CONF_THRESHOLD, verbose=False) # if results and results[0].boxes: # for box in results[0].boxes.data.tolist(): # x1, y1, x2, y2, conf, cls_id = box # cls_name = model.names[int(cls_id)] # if cls_name in TARGET_CLASSES: # yolo_detections.append({ # 'coords': (x1, y1, x2, y2), # 'class': cls_name, # 'conf': conf # }) # except Exception as e: # logging.error(f"YOLO inference failed on page {page_num}: {e}") # return 0, 0, [] # merged_detections = merge_overlapping_boxes(yolo_detections, IOU_MERGE_THRESHOLD) # final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD) # for det in final_detections: # bbox = det["coords"] # if det["class"] == "equation": # GLOBAL_EQUATION_COUNT += 1 # page_equations += 1 # b64 = crop_and_convert_to_base64(image, bbox) # detected_items.append({ # "type": "equation", # "id": f"EQUATION{GLOBAL_EQUATION_COUNT}", # "base64": b64 # }) # elif det["class"] == "figure": # GLOBAL_FIGURE_COUNT += 1 # page_figures += 1 # b64 = crop_and_convert_to_base64(image, bbox) # detected_items.append({ # "type": "figure", # "id": f"FIGURE{GLOBAL_FIGURE_COUNT}", # "base64": b64 # }) # logging.warning(f" -> Page {page_num}: EQs={page_equations}, Figs={page_figures}") # return page_equations, page_figures, detected_items # def get_latex_from_base64(base64_string: str) -> str: # if ort_model is None or processor is None: # return "[MODEL_ERROR: Model not initialized]" # try: # image_data = base64.b64decode(base64_string) # image = Image.open(io.BytesIO(image_data)).convert('RGB') # pixel_values = processor(images=image, return_tensors="pt").pixel_values # generated_ids = ort_model.generate(pixel_values) # raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True) # if not raw_text: # return "[OCR_WARNING: No formula found]" # latex = raw_text[0] # latex = re.sub(r'[\r\n]+', '', latex) # return latex # except Exception as e: # return f"[TR_OCR_ERROR: {e}]" # def extract_images_from_page_in_memory(page) -> Dict[str, str]: # """ # Extract images from a page and return: # { "EQUATION1": base64_string, "FIGURE1": base64_string } # """ # image_map = {} # image_list = page.get_images(full=True) # for idx, img in enumerate(image_list, start=1): # xref = img[0] # base = page.parent.extract_image(xref) # image_bytes = base["image"] # base64_img = base64.b64encode(image_bytes).decode("utf-8") # # Convention: first image = FIGURE1, second image = EQUATION1 etc # # You can tune this if needed # image_map[f"FIGURE{idx}"] = base64_img # return image_map # def embed_images_as_base64_in_memory(structured_data, detected_items): # tag_regex = re.compile(r'(figure|equation)(\d+)', re.IGNORECASE) # item_lookup = {d["id"]: d for d in detected_items} # final_data = [] # for item in structured_data: # text_fields = [ # item.get('question', ''), # item.get('passage', ''), # item.get('new_passage', '') # ] # if 'options' in item: # text_fields.extend(item['options'].values()) # used_tags = set() # for text in text_fields: # for m in tag_regex.finditer(text or ""): # used_tags.add(m.group(0).upper()) # for tag in used_tags: # base_key = tag.lower().replace(" ", "") # if tag not in item_lookup: # item[base_key] = "[MISSING_IMAGE]" # continue # entry = item_lookup[tag] # if entry["type"] == "equation": # item[base_key] = get_latex_from_base64(entry["base64"]) # else: # item[base_key] = entry["base64"] # final_data.append(item) # return final_data # def crop_and_convert_to_base64(image: np.ndarray, bbox: Tuple[float, float, float, float]) -> str: # x1, y1, x2, y2 = map(int, bbox) # h, w, _ = image.shape # x1 = max(0, x1) # y1 = max(0, y1) # x2 = min(w, x2) # y2 = min(h, y2) # crop = image[y1:y2, x1:x2] # _, buffer = cv2.imencode(".png", crop) # return base64.b64encode(buffer).decode("utf-8") # # ============================================================================ # # --- MAIN DOCUMENT PROCESSING FUNCTION (Fixed for JSON serialization) --- # # ============================================================================ # # NOTE: The return signature now uses Dict[str, int] for the equation counts # def run_single_pdf_preprocessing(pdf_path: str) -> Tuple[int, int, int, str, float, Dict[str, int], List[str]]: # """ # Runs the pipeline, returns counts, report, total time, page counts dict (str keys), and empty list. # """ # global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT # start_time = time.time() # log_messages = [] # all_saved_images = [] # all_base64_images: List[str] = [] # # Dictionary to store {page_number (int): equation_count (int)} # equation_counts_per_page: Dict[int, int] = {} # # Reset globals # GLOBAL_FIGURE_COUNT = 0 # GLOBAL_EQUATION_COUNT = 0 # # if os.path.exists(OUTPUT_DIR): # # shutil.rmtree(OUTPUT_DIR) # # os.makedirs(OUTPUT_DIR, exist_ok=True) # # 1. Validation and Model Loading # t0 = time.time() # if not os.path.exists(pdf_path): # report = f"❌ FATAL ERROR: Input PDF not found at {pdf_path}." # return 0, 0, 0, report, time.time() - start_time, {}, [] # try: # model = YOLO(WEIGHTS_PATH) # logging.warning(f"✅ Loaded YOLO model from: {WEIGHTS_PATH}") # except Exception as e: # report = f"❌ ERROR loading YOLO model: {e}\n(Ensure 'best.pt' is available and valid.)" # return 0, 0, 0, report, time.time() - start_time, {}, [] # t1 = time.time() # log_messages.append(f"Model Loading Time: {t1-t0:.4f}s") # # 2. PDF Loading # t2 = time.time() # try: # doc = fitz.open(pdf_path) # total_pages = doc.page_count # logging.warning(f"✅ Opened PDF with {doc.page_count} pages") # except Exception as e: # report = f"❌ ERROR loading PDF file: {e}" # return 0, 0, 0, report, time.time() - start_time, {}, [] # t3 = time.time() # log_messages.append(f"PDF Initialization Time: {t3-t2:.4f}s") # mat = fitz.Matrix(SCALE_FACTOR, SCALE_FACTOR) # # 3. Page Processing and Detection Loop # t4 = time.time() # for page_num_0_based in range(doc.page_count): # page_start_time = time.time() # fitz_page = doc.load_page(page_num_0_based) # page_num = page_num_0_based + 1 # # Render page to image for YOLO # try: # pix_start = time.time() # pix = fitz_page.get_pixmap(matrix=mat) # original_img = pixmap_to_numpy(pix) # pix_time = time.time() - pix_start # except Exception as e: # logging.error(f"Error converting page {page_num} to image: {e}. Skipping.") # continue # # Core Detection # detect_start = time.time() # # page_equations, _ = run_yolo_detection_and_count(original_img, model, page_num) # page_equations, _, page_images = run_yolo_detection_and_count(original_img, model, page_num) # all_saved_images.extend(page_images) # detect_time = time.time() - detect_start # # Store the count in the dictionary (INT keys) # equation_counts_per_page[page_num] = page_equations # page_total_time = time.time() - page_start_time # log_messages.append(f"Page {page_num} Time: Total={page_total_time:.4f}s (Render={pix_time:.4f}s, Detect={detect_time:.4f}s)") # doc.close() # t5 = time.time() # detection_loop_time = t5 - t4 # log_messages.append(f"Total Detection Loop Time ({total_pages} pages): {detection_loop_time:.4f}s") # # FIX APPLIED HERE: Convert integer keys to string keys for JSON serialization # equation_counts_per_page_str_keys: Dict[str, int] = { # str(k): v for k, v in equation_counts_per_page.items() # } # # 4. Final Report Generation # total_execution_time = t5 - start_time # report = ( # f"✅ **YOLO Counting Complete!**\n\n" # f"**1) Total Pages Detected in PDF:** **{total_pages}**\n" # f"**2) Total Equations Detected:** **{GLOBAL_EQUATION_COUNT}**\n" # f"**3) Total Figures Detected:** **{GLOBAL_FIGURE_COUNT}**\n" # f"---\n" # f"**4) Total Execution Time:** **{total_execution_time:.4f}s**\n" # f"### Detailed Step Timing\n" # f"```\n" # + "\n".join(log_messages) + # f"\n```" # ) # # Return the dictionary with string keys # # return total_pages, GLOBAL_EQUATION_COUNT, GLOBAL_FIGURE_COUNT, report, total_execution_time, equation_counts_per_page_str_keys, [] # return total_pages, GLOBAL_EQUATION_COUNT, GLOBAL_FIGURE_COUNT, report, total_execution_time, equation_counts_per_page_str_keys, all_saved_images # # ============================================================================ # # --- GRADIO INTERFACE FUNCTION (Updated) --- # # ============================================================================ # def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, Dict[str, int], List[str]]: # """ # Gradio wrapper function to handle file upload and return results. # """ # if pdf_file is None: # # Return an empty dict with string keys # return "N/A", "N/A", "N/A", "Please upload a PDF file.", {}, [] # pdf_path = pdf_file.name # try: # # Unpack the new return value: equation_counts_per_page (with string keys) # # num_pages, num_equations, num_figures, report, total_time, equation_counts_per_page, _ = run_single_pdf_preprocessing( # # pdf_path # # ) # # num_pages, num_equations, num_figures, report, total_time, equation_counts_per_page, images = run_single_pdf_preprocessing(pdf_path) # num_pages, num_equations, num_figures, report, total_time, equation_counts_per_page, images = run_single_pdf_preprocessing(pdf_path) # # Return results (6 items now) # # return str(num_pages), str(num_equations), str(num_figures), report, equation_counts_per_page, [] # return str(num_pages), str(num_equations), str(num_figures), report, equation_counts_per_page, images # except Exception as e: # error_msg = f"An unexpected error occurred: {e}" # logging.error(error_msg, exc_info=True) # # Return an empty dict on error # return "Error", "Error", "Error", error_msg, {}, [] # # ============================================================================ # # --- GRADIO INTERFACE DEFINITION (Updated) --- # # ============================================================================ # if __name__ == "__main__": # if not os.path.exists(WEIGHTS_PATH): # logging.error(f"❌ FATAL ERROR: YOLO weight file '{WEIGHTS_PATH}' not found. Cannot run live inference.") # input_file = gr.File(label="Upload PDF Document", type="filepath", file_types=[".pdf"]) # # Outputs # output_pages = gr.Textbox(label="Total Pages in PDF", interactive=False) # output_equations = gr.Textbox(label="Total Equations Detected", interactive=False) # output_figures = gr.Textbox(label="Total Figures Detected", interactive=False) # output_report = gr.Markdown(label="Processing Summary and Timing") # # NEW OUTPUT: JSON component for structured data # output_page_counts = gr.JSON(label="Equation Count Per Page (Dictionary)") # # Gradio Gallery is retained but will receive an empty list [] # output_gallery = gr.Gallery( # label="Detected Equations (Disabled for Speed)", # columns=5, # height="auto", # object_fit="contain", # allow_preview=False # ) # interface = gr.Interface( # fn=gradio_process_pdf, # inputs=input_file, # # Outputs list remains the same, but the JSON component now receives string keys. # outputs=[ # output_pages, # output_equations, # output_figures, # output_report, # output_page_counts, # output_gallery # ], # title="📊 YOLO Counting with Per-Page Data & Timing", # description=( # "Upload a PDF to run YOLO detection. The results include total counts, a breakdown of " # "equation counts per page (in JSON format), and detailed timing." # ), # ) # print("\nStarting Gradio application...") # # interface.launch(inbrowser=True) # interface.launch( # inbrowser=True, # # allowed_paths=[OUTPUT_DIR] # ) import base64 from PIL import Image import re import fitz # PyMuPDF import numpy as np import cv2 import torch import torch.serialization import os import time from typing import Optional, Tuple, List, Dict, Any, Union from ultralytics import YOLO import logging import gradio as gr import shutil import tempfile import io # ============================================================================ # --- Global Patches and Setup --- # ============================================================================ # Patch torch.load to prevent weights_only error with older models _original_torch_load = torch.load def patched_torch_load(*args, **kwargs): kwargs["weights_only"] = False return _original_torch_load(*args, **kwargs) torch.load = patched_torch_load logging.basicConfig(level=logging.WARNING) # ============================================================================ # --- CONFIGURATION AND CONSTANTS --- # ============================================================================ WEIGHTS_PATH = 'best.pt' SCALE_FACTOR = 2.0 # --- OCR Model Initialization (Retained but not used in the main loop for counting) --- from transformers import TrOCRProcessor from optimum.onnxruntime import ORTModelForVision2Seq MODEL_NAME = 'breezedeus/pix2text-mfr-1.5' # Note: These models are kept global but unused in the main flow, # as the user did not explicitly ask to remove the heavy OCR dependency yet. try: processor = TrOCRProcessor.from_pretrained(MODEL_NAME) ort_model = ORTModelForVision2Seq.from_pretrained(MODEL_NAME, use_cache=False) except Exception as e: logging.warning(f"OCR model loading failed (expected if dependencies are missing): {e}") processor = None ort_model = None # Detection parameters CONF_THRESHOLD = 0.2 TARGET_CLASSES = ['figure', 'equation'] IOU_MERGE_THRESHOLD = 0.4 IOA_SUPPRESSION_THRESHOLD = 0.7 # --- REMOVED GLOBAL COUNTERS --- # GLOBAL_FIGURE_COUNT = 0 # GLOBAL_EQUATION_COUNT = 0 # ============================================================================ # --- BOX COMBINATION LOGIC (Retained) --- # ============================================================================ def calculate_iou(box1, box2): x1_a, y1_a, x2_a, y2_a = box1 x1_b, y1_b, x2_b, y2_b = box2 x_left = max(x1_a, x1_b) y_top = max(y1_a, y1_b) x_right = min(x2_a, x2_b) y_bottom = min(y2_a, y2_b) intersection_area = max(0, x_right - x_left) * max(0, y_bottom - y_top) box_a_area = (x2_a - x1_a) * (y2_a - y1_a) box_b_area = (x2_b - x1_b) * (y2_b - y1_b) union_area = float(box_a_area + box_b_area - intersection_area) return intersection_area / union_area if union_area > 0 else 0 def filter_nested_boxes(detections, ioa_threshold=0.80): if not detections: return [] for d in detections: x1, y1, x2, y2 = d['coords'] d['area'] = (x2 - x1) * (y2 - y1) detections.sort(key=lambda x: x['area'], reverse=True) keep_indices = [] is_suppressed = [False] * len(detections) for i in range(len(detections)): if is_suppressed[i]: continue keep_indices.append(i) box_a = detections[i]['coords'] for j in range(i + 1, len(detections)): if is_suppressed[j]: continue box_b = detections[j]['coords'] x_left = max(box_a[0], box_b[0]) y_top = max(box_a[1], box_b[1]) x_right = min(box_a[2], box_b[2]) y_bottom = min(box_a[3], box_b[3]) intersection = max(0, x_right - x_left) * max(0, y_bottom - y_top) area_b = detections[j]['area'] if area_b > 0 and intersection / area_b > ioa_threshold: is_suppressed[j] = True return [detections[i] for i in keep_indices] def merge_overlapping_boxes(detections, iou_threshold): if not detections: return [] detections.sort(key=lambda d: d['conf'], reverse=True) merged_detections = [] is_merged = [False] * len(detections) for i in range(len(detections)): if is_merged[i]: continue current_box = detections[i]['coords'] current_class = detections[i]['class'] merged_x1, merged_y1, merged_x2, merged_y2 = current_box for j in range(i + 1, len(detections)): if is_merged[j] or detections[j]['class'] != current_class: continue other_box = detections[j]['coords'] iou = calculate_iou(current_box, other_box) if iou > iou_threshold: merged_x1 = min(merged_x1, other_box[0]) merged_y1 = min(merged_y1, other_box[1]) merged_x2 = max(merged_x2, other_box[2]) merged_y2 = max(other_box[3], other_box[3]) is_merged[j] = True merged_detections.append({ 'coords': (merged_x1, merged_y1, merged_x2, merged_y2), 'y1': merged_y1, 'class': current_class, 'conf': detections[i]['conf'] }) return merged_detections # ============================================================================ # --- UTILITY FUNCTIONS --- # ============================================================================ def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray: """Converts a PyMuPDF Pixmap to a NumPy array for OpenCV/YOLO.""" img = np.frombuffer(pix.samples, dtype=np.uint8).reshape( (pix.h, pix.w, pix.n) ) if pix.n == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) elif pix.n == 1: img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) return img def crop_and_convert_to_base64(image: np.ndarray, bbox: Tuple[float, float, float, float]) -> str: x1, y1, x2, y2 = map(int, bbox) h, w, _ = image.shape x1 = max(0, x1) y1 = max(0, y1) x2 = min(w, x2) y2 = min(h, y2) crop = image[y1:y2, x1:x2] _, buffer = cv2.imencode(".png", crop) return base64.b64encode(buffer).decode("utf-8") # --- NEW: Function to format base64 for Gradio Gallery --- def base64_to_gradio_gallery_tuple(base64_str: str, label: str) -> Tuple[str, str]: """Converts raw base64 to a data URI tuple for Gradio Gallery.""" # Format: ('data:image/png;base64,...', 'label') return (f"data:image/png;base64,{base64_str}", label) # --- UPDATED: run_yolo_detection_and_count to use passed counters --- def run_yolo_detection_and_count( image: np.ndarray, model: YOLO, page_num: int, current_eq_count: int, current_fig_count: int ) -> Tuple[int, int, List[Dict[str, str]], int, int]: """ Performs YOLO detection and returns page counts, detected items, and the updated global counters. """ # Use the passed counters as starting points for this page eq_counter = current_eq_count fig_counter = current_fig_count page_equations = 0 page_figures = 0 detected_items = [] yolo_detections = [] try: results = model.predict(image, conf=CONF_THRESHOLD, verbose=False) if results and results[0].boxes: for box in results[0].boxes.data.tolist(): x1, y1, x2, y2, conf, cls_id = box cls_name = model.names[int(cls_id)] if cls_name in TARGET_CLASSES: yolo_detections.append({ 'coords': (x1, y1, x2, y2), 'class': cls_name, 'conf': conf }) except Exception as e: logging.error(f"YOLO inference failed on page {page_num}: {e}") return 0, 0, [], eq_counter, fig_counter merged_detections = merge_overlapping_boxes(yolo_detections, IOU_MERGE_THRESHOLD) final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD) for det in final_detections: bbox = det["coords"] if det["class"] == "equation": eq_counter += 1 page_equations += 1 b64 = crop_and_convert_to_base64(image, bbox) detected_items.append({ "type": "equation", "id": f"EQUATION{eq_counter}", "base64": b64 }) elif det["class"] == "figure": fig_counter += 1 page_figures += 1 b64 = crop_and_convert_to_base64(image, bbox) detected_items.append({ "type": "figure", "id": f"FIGURE{fig_counter}", "base64": b64 }) logging.warning(f" -> Page {page_num}: EQs={page_equations}, Figs={page_figures}") # Return page counts, detected items, and the UPDATED total counters return page_equations, page_figures, detected_items, eq_counter, fig_counter # --- Other unused functions (get_latex_from_base64, etc.) are kept but not modified as # the focus is on the concurrency and Gradio Gallery fix. --- def get_latex_from_base64(base64_string: str) -> str: if ort_model is None or processor is None: return "[MODEL_ERROR: Model not initialized]" try: image_data = base64.b64decode(base64_string) image = Image.open(io.BytesIO(image_data)).convert('RGB') pixel_values = processor(images=image, return_tensors="pt").pixel_values generated_ids = ort_model.generate(pixel_values) raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True) if not raw_text: return "[OCR_WARNING: No formula found]" latex = raw_text[0] latex = re.sub(r'[\r\n]+', '', latex) return latex except Exception as e: return f"[TR_OCR_ERROR: {e}]" def embed_images_as_base64_in_memory(structured_data, detected_items): tag_regex = re.compile(r'(figure|equation)(\d+)', re.IGNORECASE) item_lookup = {d["id"]: d for d in detected_items} final_data = [] for item in structured_data: text_fields = [ item.get('question', ''), item.get('passage', ''), item.get('new_passage', '') ] if 'options' in item: text_fields.extend(item['options'].values()) used_tags = set() for text in text_fields: for m in tag_regex.finditer(text or ""): used_tags.add(m.group(0).upper()) for tag in used_tags: base_key = tag.lower().replace(" ", "") if tag not in item_lookup: item[base_key] = "[MISSING_IMAGE]" continue entry = item_lookup[tag] if entry["type"] == "equation": item[base_key] = get_latex_from_base64(entry["base64"]) else: item[base_key] = entry["base64"] final_data.append(item) return final_data # ============================================================================ # --- MAIN DOCUMENT PROCESSING FUNCTION (Fixed for concurrency) --- # ============================================================================ # --- UPDATED return type for clarity --- def run_single_pdf_preprocessing( pdf_path: str ) -> Tuple[int, int, int, str, float, Dict[str, int], List[Tuple[str, str]]]: """ Runs the pipeline, returns counts, report, total time, page counts dict (str keys), and a list of (image_data_uri, label) for the Gradio gallery. """ # --- INITIALIZE LOCAL COUNTERS --- start_time = time.time() log_messages = [] # This list now holds (data_uri, label) tuples for Gradio all_gradio_gallery_items: List[Tuple[str, str]] = [] # Dictionary to store {page_number (int): equation_count (int)} equation_counts_per_page: Dict[int, int] = {} # --- USE LOCAL COUNTERS FOR THREAD SAFETY --- total_figure_count = 0 total_equation_count = 0 # 1. Validation and Model Loading t0 = time.time() if not os.path.exists(pdf_path): report = f"❌ FATAL ERROR: Input PDF not found at {pdf_path}." # Return empty list of tuples for gallery on error return 0, 0, 0, report, time.time() - start_time, {}, [] try: model = YOLO(WEIGHTS_PATH) logging.warning(f"✅ Loaded YOLO model from: {WEIGHTS_PATH}") except Exception as e: report = f"❌ ERROR loading YOLO model: {e}\n(Ensure 'best.pt' is available and valid.)" return 0, 0, 0, report, time.time() - start_time, {}, [] t1 = time.time() log_messages.append(f"Model Loading Time: {t1-t0:.4f}s") # 2. PDF Loading t2 = time.time() try: doc = fitz.open(pdf_path) total_pages = doc.page_count logging.warning(f"✅ Opened PDF with {doc.page_count} pages") except Exception as e: report = f"❌ ERROR loading PDF file: {e}" return 0, 0, 0, report, time.time() - start_time, {}, [] t3 = time.time() log_messages.append(f"PDF Initialization Time: {t3-t2:.4f}s") mat = fitz.Matrix(SCALE_FACTOR, SCALE_FACTOR) # 3. Page Processing and Detection Loop t4 = time.time() for page_num_0_based in range(doc.page_count): page_start_time = time.time() fitz_page = doc.load_page(page_num_0_based) page_num = page_num_0_based + 1 # Render page to image for YOLO try: pix_start = time.time() pix = fitz_page.get_pixmap(matrix=mat) original_img = pixmap_to_numpy(pix) pix_time = time.time() - pix_start except Exception as e: logging.error(f"Error converting page {page_num} to image: {e}. Skipping.") continue # Core Detection detect_start = time.time() # --- PASSING AND RECEIVING THE COUNTERS HERE (Concurrency Fix) --- ( page_equations, page_figures, page_images_dicts, total_equation_count, total_figure_count ) = run_yolo_detection_and_count( original_img, model, page_num, total_equation_count, total_figure_count ) # --- FORMATTING FOR GRADIO GALLERY (Gradio Format Fix) --- for item in page_images_dicts: gradio_tuple = base64_to_gradio_gallery_tuple(item["base64"], item["id"]) all_gradio_gallery_items.append(gradio_tuple) detect_time = time.time() - detect_start # Store the count in the dictionary (INT keys) equation_counts_per_page[page_num] = page_equations page_total_time = time.time() - page_start_time log_messages.append(f"Page {page_num} Time: Total={page_total_time:.4f}s (Render={pix_time:.4f}s, Detect={detect_time:.4f}s)") doc.close() t5 = time.time() detection_loop_time = t5 - t4 log_messages.append(f"Total Detection Loop Time ({total_pages} pages): {detection_loop_time:.4f}s") # Convert integer keys to string keys for JSON serialization equation_counts_per_page_str_keys: Dict[str, int] = { str(k): v for k, v in equation_counts_per_page.items() } # 4. Final Report Generation total_execution_time = t5 - start_time report = ( f"✅ **YOLO Counting Complete!**\n\n" f"**1) Total Pages Detected in PDF:** **{total_pages}**\n" f"**2) Total Equations Detected:** **{total_equation_count}**\n" # Uses local final count f"**3) Total Figures Detected:** **{total_figure_count}**\n" # Uses local final count f"---\n" f"**4) Total Execution Time:** **{total_execution_time:.4f}s**\n" f"### Detailed Step Timing\n" f"```\n" + "\n".join(log_messages) + f"\n```" ) # Return the dictionary with string keys and the properly formatted gallery items return total_pages, total_equation_count, total_figure_count, report, total_execution_time, equation_counts_per_page_str_keys, all_gradio_gallery_items # ============================================================================ # --- GRADIO INTERFACE FUNCTION (Updated) --- # ============================================================================ # --- UPDATED return type for clarity --- def gradio_process_pdf(pdf_file) -> Tuple[str, str, str, str, Dict[str, int], List[Tuple[str, str]]]: """ Gradio wrapper function to handle file upload and return results. """ if pdf_file is None: # Return empty list of tuples for gallery on error return "N/A", "N/A", "N/A", "Please upload a PDF file.", {}, [] pdf_path = pdf_file.name try: # Unpack the new return value: equation_counts_per_page (with string keys) ( num_pages, num_equations, num_figures, report, total_time, equation_counts_per_page, gallery_items # Now correctly formatted list of tuples ) = run_single_pdf_preprocessing(pdf_path) # Return results (6 items now) return str(num_pages), str(num_equations), str(num_figures), report, equation_counts_per_page, gallery_items except Exception as e: error_msg = f"An unexpected error occurred: {e}" logging.error(error_msg, exc_info=True) # Return empty list of tuples for gallery on error return "Error", "Error", "Error", error_msg, {}, [] # ============================================================================ # --- GRADIO INTERFACE DEFINITION (Updated) --- # ============================================================================ if __name__ == "__main__": if not os.path.exists(WEIGHTS_PATH): logging.error(f"❌ FATAL ERROR: YOLO weight file '{WEIGHTS_PATH}' not found. Cannot run live inference.") input_file = gr.File(label="Upload PDF Document", type="filepath", file_types=[".pdf"]) # Outputs output_pages = gr.Textbox(label="Total Pages in PDF", interactive=False) output_equations = gr.Textbox(label="Total Equations Detected", interactive=False) output_figures = gr.Textbox(label="Total Figures Detected", interactive=False) output_report = gr.Markdown(label="Processing Summary and Timing") # NEW OUTPUT: JSON component for structured data output_page_counts = gr.JSON(label="Equation Count Per Page (Dictionary)") # Gradio Gallery is retained and now receives the correctly formatted list of tuples output_gallery = gr.Gallery( label="Detected Items (Gallery Format Fix Applied)", columns=5, height="auto", object_fit="contain", allow_preview=False ) interface = gr.Interface( fn=gradio_process_pdf, inputs=input_file, # Outputs list remains the same, but the gallery now works outputs=[ output_pages, output_equations, output_figures, output_report, output_page_counts, output_gallery ], title="📊 YOLO Counting with Per-Page Data & Timing (Concurrency Fix)", description=( "Upload a PDF to run YOLO detection. The concurrency bug and Gradio Gallery display error have been fixed." ), ) print("\nStarting Gradio application...") interface.launch(inbrowser=True)