app.py

# app.py
import os
import cv2
import tempfile
import shutil
import numpy as np
# import easyocr
import gradio as gr
from ultralytics import YOLO
from openpyxl import Workbook
from openpyxl.utils import get_column_letter
import pandas as pd
import pytesseract
from PIL import Image

# =============================
# USER CONFIG (edit if needed)
# =============================
# Put model files in space root or change to full path
TABLE_MODEL_PATH = "models/Table_Detection.pt"
MODEL_PATH = "models/RoCoCe_best.pt"

# Device choices: 'cpu' or 'cuda'
USE_CUDA = False

# Detection thresholds
CONF_THRESHOLD = 0.25
IOU_THRESHOLD = 0.4

# OCR settings
OCR_LANGS = ["fr"]
# USE_GPU_FOR_OCR = False  # EasyOCR GPU usage (separate from YOLO device)

# Pipeline settings
MIN_COL_OVERLAP = 0.3
MERGE_SPANNING_IN_EXCEL = True

# Reading order tolerance (pixels)
ROW_TOLERANCE = 50

# Global cached models (lazy load)
_table_model = None
_structure_model = None
_reader = None

# Device strings
YOLO_DEVICE = "cuda" if USE_CUDA else "cpu"

# Maximum number of sheet tabs to show in the UI (increase if you expect more sheets)
MAX_SHEETS = 12

#  Add helper to read Excel into HTML tables
def excel_to_html_sheets(excel_path):
    """Return a dict {sheet_name: HTML table} from an Excel file."""
    xls = pd.ExcelFile(excel_path)
    sheet_html = {}
    for sheet in xls.sheet_names:
        df = pd.read_excel(excel_path, sheet_name=sheet, header=None)
        sheet_html[sheet] = df.to_html(index=False, header=False, escape=False)
    return sheet_html

# =============================
# Helper: load models lazily
# =============================
def load_models_if_needed():
    global _table_model, _structure_model, _reader
    if _table_model is None:
        if not os.path.exists(TABLE_MODEL_PATH):
            raise FileNotFoundError(f"Table model not found: {TABLE_MODEL_PATH}")
        print(f"[INFO] Loading table model from {TABLE_MODEL_PATH} to {YOLO_DEVICE} ...")
        _table_model = YOLO(TABLE_MODEL_PATH).to(YOLO_DEVICE)
    if _structure_model is None:
        if not os.path.exists(MODEL_PATH):
            raise FileNotFoundError(f"Structure model not found: {MODEL_PATH}")
        print(f"[INFO] Loading structure model from {MODEL_PATH} to {YOLO_DEVICE} ...")
        _structure_model = YOLO(MODEL_PATH).to(YOLO_DEVICE)
    # if _reader is None:
    #     print(f"[INFO] Initializing EasyOCR reader (langs={OCR_LANGS}, gpu={USE_GPU_FOR_OCR}) ...")
    #     _reader = easyocr.Reader(OCR_LANGS, gpu=USE_GPU_FOR_OCR)

# =============================
# Utility functions (your original logic)
# =============================
def run_detection(model, image_path, conf_thres=0.25, iou_thres=0.4):
    results = model.predict(source=image_path, conf=conf_thres, iou=iou_thres, device=YOLO_DEVICE, verbose=False)
    if not results:
        return []
    r = results[0]
    detections = []
    # handle case where boxes may be empty
    boxes = getattr(r.boxes, "xyxy", None)
    if boxes is None or len(r.boxes) == 0:
        return []
    for box, cls_id, conf in zip(r.boxes.xyxy.cpu().numpy(),
                                 r.boxes.cls.cpu().numpy(),
                                 r.boxes.conf.cpu().numpy()):
        detections.append({
            "bbox": tuple(map(int, box)),
            "cls": int(cls_id),
            "conf": float(conf)
        })
    return detections

def detect_and_crop_tables(table_model, image_path, temp_dir, save_debug=False):
    print(f"[INFO] Running table detector on {image_path} ...")
    results = table_model.predict(source=image_path, conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, verbose=False)
    if not results or len(results[0].boxes) == 0:
        print("[WARN] Table detector found no boxes.")
        return []

    r = results[0]
    xyxy = r.boxes.xyxy.cpu().numpy()
    confs = r.boxes.conf.cpu().numpy()

    img = cv2.imread(image_path)
    crops = []
    for idx, (x1, y1, x2, y2) in enumerate(xyxy):
        x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2)
        crop = img[y1:y2, x1:x2]
        if crop.size == 0:
            continue
        crop_path = os.path.join(temp_dir, f"temp_table_crop_{idx+1}.jpg")
        cv2.imwrite(crop_path, crop)
        crops.append((crop_path, (x1, y1, x2, y2), float(confs[idx])))

        if save_debug:
            # draw on debug image (we'll save debug later)
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

    if save_debug:
        debug_path = os.path.join(temp_dir, "table_detection_debug.jpg")
        cv2.imwrite(debug_path, img)
    print(f"[INFO] {len(crops)} table(s) cropped.")
    return crops

def assign_cells_to_columns(detections, min_overlap=0.3):
    columns = [d for d in detections if d["cls"] == 0]
    rows = [d for d in detections if d["cls"] == 1]
    cells = [d for d in detections if d["cls"] in [2, 3]]
    columns = sorted(columns, key=lambda c: c["bbox"][0])
    rows = sorted(rows, key=lambda r: r["bbox"][1])

    assigned_cells = []
    for cell in cells:
        c_x1, c_y1, c_x2, c_y2 = cell["bbox"]
        col_matches = []
        for idx, col in enumerate(columns):
            col_x1, col_y1, col_x2, col_y2 = col["bbox"]
            overlap_x = max(0, min(c_x2, col_x2) - max(c_x1, col_x1))
            min_width = min(col_x2 - col_x1, c_x2 - c_x1)
            if min_width <= 0:
                continue
            if (overlap_x / min_width) >= min_overlap:
                col_matches.append(idx)
        if col_matches:
            cell["columns"] = col_matches
            assigned_cells.append(cell)

    return columns, rows, assigned_cells

# def ocr_cells_on_image(img, cells, reader):
#     for cell in cells:
#         x1, y1, x2, y2 = cell["bbox"]
#         # safety bounds
#         x1, y1 = max(0, x1), max(0, y1)
#         x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2)
#         if x2 <= x1 or y2 <= y1:
#             cell["text"] = ""
#             continue
#         crop = img[y1:y2, x1:x2]
#         # EasyOCR wants RGB or grayscale but works with BGR too; convert to RGB optionally
#         try:
#             ocr_result = reader.readtext(crop)
#         except Exception as e:
#             print(f"[WARN] EasyOCR failed on crop: {e}")
#             ocr_result = []
#         text = " ".join([res[1] for res in ocr_result]) if ocr_result else ""
#         cell["text"] = text.strip()
#     return cells

def ocr_cells_on_image(img, cells, lang="fra"):
    """
    OCR each detected cell using Tesseract instead of EasyOCR.

    Args:
        img (numpy.ndarray): The full table image (BGR from cv2).
        cells (list[dict]): Each cell dict must have a "bbox" key: (x1, y1, x2, y2).
        lang (str): Tesseract language code(s), e.g. "fra" or "eng+fra".

    Returns:
        list[dict]: cells with added "text" keys.
    """
    for cell in cells:
        x1, y1, x2, y2 = map(int, cell["bbox"])
        # safety bounds
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2)
        if x2 <= x1 or y2 <= y1:
            cell["text"] = ""
            continue

        crop = img[y1:y2, x1:x2]

        # Convert OpenCV BGR → RGB → PIL
        pil_crop = Image.fromarray(crop[:, :, ::-1])

        # Optional: preprocess for better OCR accuracy
        pil_crop = pil_crop.convert("L")  # grayscale

        try:
            text = pytesseract.image_to_string(
                pil_crop,
                lang=lang,
                config="--psm 6"  # treat image as a block of text
            )
        except Exception as e:
            print(f"[WARN] Tesseract failed on crop: {e}")
            text = ""

        cell["text"] = text.strip()

    return cells


def group_cells_into_rows(columns, row_boxes, cells):
    row_boxes_sorted = sorted(row_boxes, key=lambda r: r["bbox"][1])
    rows_ordered = []
    for row in row_boxes_sorted:
        rows_ordered.append(row["bbox"])

    cells_grouped = []
    for row_bbox in rows_ordered:
        r_x1, r_y1, r_x2, r_y2 = row_bbox
        row_cells = []
        for cell in cells:
            c_x1, c_y1, c_x2, c_y2 = cell["bbox"]
            overlap_y = max(0, min(c_y2, r_y2) - max(c_y1, r_y1))
            min_height = min(r_y2 - r_y1, c_y2 - c_y1)
            if min_height <= 0:
                continue
            if (overlap_y / min_height) >= 0.5:
                row_cells.append(cell)
        cells_grouped.append(row_cells)
    return rows_ordered, cells_grouped

def build_table_matrix(columns, rows_ordered, cells_grouped, num_columns):
    table = [["" for _ in range(num_columns)] for _ in range(len(rows_ordered))]
    merges = []
    for r_idx, row_cells in enumerate(cells_grouped, start=1):
        for cell in row_cells:
            col_indices = cell.get("columns", [])
            if not col_indices:
                continue
            text = cell.get("text", "")
            c_start = min(col_indices)
            c_end = max(col_indices)
            table[r_idx-1][c_start] = text
            if c_end > c_start:
                merges.append((r_idx, c_start+1, c_end+1))
    return table, merges

def save_all_tables_to_excel(tables_data, output_path):
    wb = Workbook()
    # Remove default sheet
    if wb.active:
        wb.remove(wb.active)

    for table_matrix, merges, sheet_name in tables_data:
        ws = wb.create_sheet(title=sheet_name)
        for r_idx, row in enumerate(table_matrix, start=1):
            for c_idx, val in enumerate(row, start=1):
                ws.cell(row=r_idx, column=c_idx, value=val)
        applied = set()
        for m in merges or []:
            if len(m) == 4:
                r1, c1, r2, c2 = m
            elif len(m) == 3:
                r1, c1, c2 = m
                r2 = r1
            else:
                continue
            key = (r1, c1, r2, c2)
            if key in applied:
                continue
            if c2 > c1:
                ws.merge_cells(start_row=r1, start_column=c1, end_row=r2, end_column=c2)
                applied.add(key)
        for col in ws.columns:
            max_length = 0
            col_letter = get_column_letter(col[0].column)
            for cell in col:
                if cell.value:
                    max_length = max(max_length, len(str(cell.value)))
            ws.column_dimensions[col_letter].width = max_length + 2

    wb.save(output_path)
    print(f"[INFO] All tables saved to {output_path}")

# sort_tables_by_reading_order (keeps bbox & conf)
def sort_tables_by_reading_order(tables, row_tolerance=50):
    """
    Sort tables in natural reading order (top-to-bottom, left-to-right).
    tables: list of tuples (crop_path, bbox, conf)
    bbox format: (x1, y1, x2, y2)
    """
    if not tables:
        return []
    # First, sort by top edge then left edge
    tables_sorted = sorted(tables, key=lambda t: (t[1][1], t[1][0]))
    final_sorted = []
    current_band = []
    current_band_y = None

    for t in tables_sorted:
        _, bbox, _ = t
        x1, y1, _, _ = bbox
        if current_band_y is None or abs(y1 - current_band_y) <= row_tolerance:
            current_band.append(t)
            if current_band_y is None:
                current_band_y = y1
        else:
            # sort current band left-to-right
            current_band.sort(key=lambda tb: tb[1][0])
            final_sorted.extend(current_band)
            current_band = [t]
            current_band_y = y1

    if current_band:
        current_band.sort(key=lambda tb: tb[1][0])
        final_sorted.extend(current_band)

    return final_sorted

# =============================
# Gradio pipeline function
# =============================
def process_image_with_steps(image):
    """
    Input: image as numpy array (RGB) from Gradio
    Returns: logs (text), table_detection_image (path or None), crops list (list of paths),
             overlays list (list of paths), excel file path (or None), sheet_list_for_ui (list of (name, html))
    """
    # ensure models loaded
    try:
        load_models_if_needed()
    except Exception as e:
        msg = f"[ERROR] Failed loading models: {e}"
        print(msg)
        return msg, None, [], [], None, []

    log_lines = []
    def log_print(*args):
        msg = " ".join(str(a) for a in args)
        print(msg)
        log_lines.append(msg)

    # create temp dir for this run
    temp_dir = tempfile.mkdtemp(prefix="yolo_ocr_")
    log_print(f"[INFO] Temporary directory: {temp_dir}")

    # save input image (gradio passes RGB; convert to BGR for cv2)
    input_path = os.path.join(temp_dir, "input.jpg")
    cv2.imwrite(input_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
    log_print(f"[INFO] Saved uploaded image to {input_path}")

    # Step 1: Table detection & crop
    try:
        table_crops = detect_and_crop_tables(_table_model, input_path, temp_dir, save_debug=True)
    except Exception as e:
        log_print(f"[ERROR] Table detector failed: {e}")
        # cleanup on failure
        # keep temp dir for debugging when error occurs (do not delete)
        return "\n".join(log_lines), None, [], [], None, []

    if not table_crops:
        log_print("[WARN] No tables detected.")
        # cleanup and return
        shutil.rmtree(temp_dir)
        return "\n".join(log_lines), None, [], [], None, []

    # Save table detection debug image if present
    table_detection_debug = os.path.join(temp_dir, "table_detection_debug.jpg")
    if os.path.exists(table_detection_debug):
        table_det_img_path = table_detection_debug
    else:
        table_det_img_path = None

    # Sort tables by reading order (important!)
    log_print("[INFO] Sorting detected tables by reading order ...")
    table_crops_sorted = sort_tables_by_reading_order(table_crops, row_tolerance=ROW_TOLERANCE)
    log_print(f"[INFO] {len(table_crops_sorted)} table(s) after sorting.")

    # Build list of crop image paths (sorted)
    crop_paths_sorted = [t[0] for t in table_crops_sorted]

    # Step 2..N: For each crop, run structure detection, OCR and create overlays
    overlays = []
    tables_data = []
    for idx, (crop_path, bbox, tconf) in enumerate(table_crops_sorted, start=1):
        log_print(f"[INFO] Processing table #{idx} -> {crop_path} bbox={bbox} conf={tconf:.3f}")
        try:
            detections = run_detection(_structure_model, crop_path, conf_thres=CONF_THRESHOLD, iou_thres=IOU_THRESHOLD)
        except Exception as e:
            log_print(f"[ERROR] Structure detection failed on {crop_path}: {e}")
            continue

        if len(detections) == 0:
            log_print(f"[WARN] No structure detections inside table #{idx} - skipping.")
            continue

        crop_img = cv2.imread(crop_path)
        if crop_img is None:
            log_print(f"[WARN] Failed to read crop image {crop_path} - skipping.")
            continue

        columns, row_boxes, assigned_cells = assign_cells_to_columns(detections, min_overlap=MIN_COL_OVERLAP)
        if len(columns) == 0:
            log_print(f"[ERROR] No column detections found in table #{idx} - skipping.")
            continue
        log_print(f"[INFO] Table #{idx}: {len(columns)} columns, {len(assigned_cells)} candidate cells, {len(row_boxes)} row boxes")

        # OCR assigned cells
        # assigned_cells = ocr_cells_on_image(crop_img, assigned_cells, _reader)
        assigned_cells = ocr_cells_on_image(crop_img, assigned_cells, lang="fra")
        # Group into rows
        rows_ordered, cells_grouped = group_cells_into_rows(columns, row_boxes, assigned_cells)
        log_print(f"[INFO] Table #{idx}: Formed {len(cells_grouped)} rows")

        # Build matrix and merges
        num_columns = len(columns)
        table_matrix, merges = build_table_matrix(columns, rows_ordered, cells_grouped, num_columns)
        log_print(f"[INFO] Table #{idx}: Built table rows={len(table_matrix)}, cols={num_columns}, merges={len(merges)}")

        # Create overlay image showing structure boxes + OCR text
        overlay = crop_img.copy()
        # draw columns (green), rows (blue), cells (red), and put OCR text
        for col in columns:
            x1, y1, x2, y2 = col["bbox"]
            cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 255, 0), 1)
        for r in row_boxes:
            x1, y1, x2, y2 = r["bbox"]
            cv2.rectangle(overlay, (x1, y1), (x2, y2), (255, 0, 0), 1)
        for cell in assigned_cells:
            x1, y1, x2, y2 = cell["bbox"]
            cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 0, 255), 1)
            text = cell.get("text", "")
            if text:
                # label safely inside bounds
                tx, ty = x1 + 2, max(12, y1 + 12)
                # small font scale to avoid overflow; may wrap not implemented
                cv2.putText(overlay, text if len(text) < 80 else text[:80] + "...",
                            (tx, ty), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0,0,255), 1, cv2.LINE_AA)

        overlay_path = os.path.join(temp_dir, f"overlay_table_{idx}.jpg")
        cv2.imwrite(overlay_path, overlay)
        overlays.append(overlay_path)

        # Add to tables_data for Excel
        sheet_name = f"Table_{idx}"
        tables_data.append((table_matrix, merges if MERGE_SPANNING_IN_EXCEL else [], sheet_name))

    # Step: create Excel if any tables_data
    excel_path = None
    sheet_html_dict = {}
    if tables_data:
        excel_path = os.path.join(temp_dir, "output_tables.xlsx")
        save_all_tables_to_excel(tables_data, excel_path)
        log_print(f"[INFO] Excel created at: {excel_path}")
        # Convert all sheets to HTML for UI display
        try:
            sheet_html_dict = excel_to_html_sheets(excel_path)
        except Exception as e:
            log_print(f"[WARN] Failed to convert excel to html: {e}")
            sheet_html_dict = {}
    else:
        log_print("[WARN] No table matrices were produced; no Excel to save.")
        sheet_html_dict = {}

    # Prepare results for Gradio:
    # logs_text = "\n".join(log_lines)
    # Return: logs, table detection image, crops list, overlays list, excel file path, sheet_list_for_ui
    # If table_det_img_path is None, return the original input as fallback small image
    if table_det_img_path is None:
        table_det_img_path = input_path  # fallback

    # Convert dict to list of (sheet_name, html_content) for Gradio display
    sheet_list_for_ui = [(name, html) for name, html in sheet_html_dict.items()]
    # return logs_text, table_det_img_path, crop_paths_sorted, overlays, excel_path, sheet_list_for_ui
    return table_det_img_path, crop_paths_sorted, overlays, excel_path, sheet_list_for_ui

# -----------------------------
# Wrapper to expand sheets to fixed number of HTML outputs
# -----------------------------
def process_and_expand_sheets(image):
    # logs_text, table_det_img_path, crop_paths_sorted, overlays, excel_path, sheet_list_for_ui = process_image_with_steps(image)
    table_det_img_path, crop_paths_sorted, overlays, excel_path, sheet_list_for_ui = process_image_with_steps(image)
    # Build html list
    html_contents = [html for _, html in sheet_list_for_ui]
    # Pad with empty htmls so number of outputs is constant
    while len(html_contents) < MAX_SHEETS:
        html_contents.append("<div></div>")
    # If there are more sheets than MAX_SHEETS, truncate (or optionally handle differently)
    if len(html_contents) > MAX_SHEETS:
        html_contents = html_contents[:MAX_SHEETS]
    # Return expanded outputs: logs, table detection image, crops, overlays, excel file, then the HTMLs
    # return (logs_text, table_det_img_path, crop_paths_sorted, overlays, excel_path, *html_contents)
    return (table_det_img_path, crop_paths_sorted, overlays, excel_path, *html_contents)

# =============================
# Gradio UI layout
# =============================
# with gr.Blocks(title="YOLO + EasyOCR Table Extraction (with reading order)") as demo:
with gr.Blocks(title="YOLO + Tesseract Table Extraction (with reading order)") as demo:  
    gr.Markdown("## 📄 Dual-stage : Table Detection → Structure Detection → OCR → Excel")
    gr.Markdown("Upload an image that contains one or more tables. The app will show intermediate steps (detection, crops, overlays) and produce an Excel workbook. Each Excel sheet is shown in its own tab below.")

    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(type="numpy", label="Upload Image (JPG/PNG)")
            run_btn = gr.Button("Run Pipeline")
            # Quick settings (optional)
            with gr.Accordion("Advanced options (change before Run)", open=False):
                conf_in = gr.Slider(minimum=0.01, maximum=1.0, value=CONF_THRESHOLD, label="Confidence threshold", step=0.01)
                iou_in = gr.Slider(minimum=0.01, maximum=1.0, value=IOU_THRESHOLD, label="IOU threshold", step=0.01)
                tol_in = gr.Slider(minimum=0, maximum=300, value=ROW_TOLERANCE, label="Reading-order row tolerance (px)")
        with gr.Column(scale=1):
        #     # logs_out = gr.Textbox(label="Processing Log (debug prints)", lines=18)
        #     # input_image = gr.Image(label="Upload Image", type="numpy")
        #     example_images = [
        #                         ["examples/example1.png"],
        #                         ["examples/example2.jpg"]
        #     ]

            example_images = [
                                ["examples/example1.jpg"],
                                ["examples/example2.jpg"]
            ]
            
            gr.Examples(
                examples=example_images,
                inputs = [inp],
                label="Example Images"
            )


        
    gr.Markdown("### Step 1: Table detection visualization")
    table_detection_img = gr.Image(label="Table Detection (debug overlay)")

    gr.Markdown("### Step 2: Cropped tables (sorted by reading order)")
    crops_gallery = gr.Gallery(label="Table Crops (sorted)", columns=3)

    gr.Markdown("### Step 3: Structure + OCR overlay for each table")
    overlays_gallery = gr.Gallery(label="Structure + OCR Overlays", columns=3)

    gr.Markdown("### Result: Download Excel")
    excel_file = gr.File(label="Download Excel workbook (contains one sheet per table)")

    # Create MAX_SHEETS HTML outputs (each will be a tab)
    sheet_html_outputs = [gr.HTML(label=f"Sheet {i+1}") for i in range(MAX_SHEETS)]

    # Hook up the button
    run_btn.click(
        fn=process_and_expand_sheets,
        inputs=[inp],
        # outputs=[logs_out, table_detection_img, crops_gallery, overlays_gallery, excel_file] + sheet_html_outputs,
        outputs=[table_detection_img, crops_gallery, overlays_gallery, excel_file] + sheet_html_outputs,
    )

if __name__ == "__main__":
    # demo.launch(server_name="127.0.0.1", server_port=7860)
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, ssr_mode=False)