# 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 # ============================= # 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 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) # 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("
") # 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: 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)