app.py

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 

# 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]:
    """
    Runs YOLO inference, applies NMS/filtering, and updates global counters.
    Returns page counts only.
    """
    global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
    
    yolo_detections = []
    page_equations = 0
    page_figures = 0
    
    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

    # Apply NMS/Merging/Filtering
    merged_detections = merge_overlapping_boxes(yolo_detections, IOU_MERGE_THRESHOLD)
    final_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD)

    # Update Global Counters
    for det in final_detections:
        if det['class'] == 'figure':
            GLOBAL_FIGURE_COUNT += 1
            page_figures += 1
        elif det['class'] == 'equation':
            GLOBAL_EQUATION_COUNT += 1
            page_equations += 1
            
    logging.warning(f"  -> Page {page_num}: EQs={page_equations}, Figs={page_figures}")
    return page_equations, page_figures


# ============================================================================
# --- 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 = []
    
    # 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

    # 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)
        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, []


# ============================================================================
# --- 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
        )
        
        # Return results (6 items now)
        return str(num_pages), str(num_equations), str(num_figures), report, equation_counts_per_page, []
            
    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)