app.py

import gradio as gr
import torch
from transformers import TableTransformerForObjectDetection
import matplotlib.pyplot as plt
from transformers import DetrFeatureExtractor
import pandas as pd
import uuid
from surya.ocr import run_ocr
# from surya.model.detection.segformer import load_model as load_det_model, load_processor as load_det_processor
from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
from surya.model.recognition.model import load_model as load_rec_model
from surya.model.recognition.processor import load_processor as load_rec_processor
from PIL import ImageDraw, Image
import os
from pdf2image import convert_from_path
import tempfile
from ultralyticsplus import YOLO, render_result
import cv2
import numpy as np
from fpdf import FPDF

def convert_pdf_images(pdf_path):
    # Convert PDF to images
    images = convert_from_path(pdf_path)

    # Save each page as a temporary image and collect file paths
    temp_file_paths = []
    for i, page in enumerate(images):
        # Create a temporary file with a unique name
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        page.save(temp_file.name, 'PNG')  # Save the image to the temporary file
        temp_file_paths.append(temp_file.name)  # Add file path to the list

    return temp_file_paths[0]  # Return the list of temporary file paths


# Load model
model_yolo = YOLO('keremberke/yolov8m-table-extraction')

# Set model parameters
model_yolo.overrides['conf'] = 0.25  # NMS confidence threshold
model_yolo.overrides['iou'] = 0.45  # NMS IoU threshold
model_yolo.overrides['agnostic_nms'] = False  # NMS class-agnostic
model_yolo.overrides['max_det'] = 1000  # maximum number of detections per image
def resize_image(image, max_dimension=4200, min_dimension=50):
    width, height = image.size
    # Check if the dimensions are within range
    if width > max_dimension or height > max_dimension or width < min_dimension or height < min_dimension:
        scaling_factor = min(max_dimension / max(width, height), min_dimension / min(width, height))
        new_width = int(width * scaling_factor)
        new_height = int(height * scaling_factor)
        return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    return image
def crop_table(filename):
    # Set image
    image_path = filename
    image = Image.open(image_path)
    image_np = np.array(image)

    # Perform inference
    results = model_yolo.predict(image_path)

    # Extract the first bounding box (assuming there's only one table)
    bbox = results[0].boxes[0]
    x1, y1, x2, y2 = map(int, bbox.xyxy[0])  # Get the bounding box coordinates

    # Crop the image using the bounding box coordinates
    cropped_image = image_np[y1:y2, x1:x2]

    # Convert the cropped image to RGB (if it's not already in RGB)
    cropped_image_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)

    # Save the cropped image as a PDF
    cropped_image_pil = Image.fromarray(cropped_image_rgb)
    # Save the cropped image to a temporary file
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    cropped_image_pil.save(temp_file.name)

    return temp_file.name

# new v1.1 checkpoints require no timm anymore
device = "cuda" if torch.cuda.is_available() else "cpu"
langs = ["en","th"] # Replace with your languages - optional but recommended
det_processor, det_model = load_det_processor(), load_det_model()
rec_model, rec_processor = load_rec_model(), load_rec_processor()



COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
feature_extractor = DetrFeatureExtractor()

model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition-v1.1-all")



def compute_boxes(image_path):
    image = Image.open(image_path).convert("RGB")
    width, height = image.size

    encoding = feature_extractor(image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**encoding)

    results = feature_extractor.post_process_object_detection(outputs, threshold=0.7, target_sizes=[(height, width)])[0]
    boxes = results['boxes'].tolist()
    labels = results['labels'].tolist()

    return boxes,labels

def extract_table(image_path):
    image = Image.open(image_path)
    boxes,labels = compute_boxes(image_path)


    cropped_table_visualized = image.copy()
    draw = ImageDraw.Draw(cropped_table_visualized)

    for cell in boxes:
        draw.rectangle(cell, outline="red")
    bbox_table = f"{str(uuid.uuid4())}.png"
    cropped_table_visualized.save(bbox_table)
    cell_locations = []

    for box_row, label_row in zip(boxes, labels):
        if label_row == 2:
            for box_col, label_col in zip(boxes, labels):
                if label_col == 1:
                    cell_box = (box_col[0], box_row[1], box_col[2], box_row[3])
                    cell_locations.append(cell_box)

    cell_locations.sort(key=lambda x: (x[1], x[0]))

    num_columns = 0
    box_old = cell_locations[0]

    for box in cell_locations[1:]:
        x1, y1, x2, y2 = box
        x1_old, y1_old, x2_old, y2_old = box_old
        num_columns += 1
        if y1 > y1_old:
            break

        box_old = box

    headers = []
    for box in cell_locations[:num_columns]:
        x1, y1, x2, y2 = box
        cell_image = resize_image(image.crop((x1, y1, x2, y2)))
        # new_width = cell_image.width *4
        # new_height = cell_image.height *4
        # cell_image = cell_image.resize((new_width, new_height), resample=Image.LANCZOS)
        # cell_text = pytesseract.image_to_string(cell_image, lang='tha+eng')
        # print(cell_text)

        plt.figure()
        plt.imshow(cell_image)
        plt.axis("off")
        plt.title("Cropped Cell Image")
        plt.show()

        predictions = run_ocr([cell_image], [langs], det_model, det_processor, rec_model, rec_processor)
        texts = [line.text for line in predictions[0].text_lines]
        all_text = ' '.join(texts)
        print(all_text)
        if all_text:
          headers.append(all_text)
        else:
          headers.append('')


    df = pd.DataFrame(columns=headers)

    row = []
    for box in cell_locations[num_columns:]:
        x1, y1, x2, y2 = box
        cell_image = resize_image(image.crop((x1, y1, x2, y2)))
        # new_width = cell_image.width * 4
        # new_height = cell_image.height * 4
        # cell_image = cell_image.resize((new_width, new_height), resample=Image.LANCZOS)
        # cell_text = pytesseract.image_to_string(cell_image, lang='tha+eng')
        # print(cell_text)

        plt.figure()
        plt.imshow(cell_image)
        plt.axis("off")
        plt.title("Cropped Cell Image")
        plt.show()
        predictions = run_ocr([cell_image], [langs], det_model, det_processor, rec_model, rec_processor)
        texts = [line.text for line in predictions[0].text_lines]
        all_text = ''.join(texts)
        print(all_text)
        if all_text:
          headers.append(all_text)
        else:
          headers.append('')

        row.append(all_text)

        if len(row) == num_columns:
            df.loc[len(df)] = row
            print(row)
            row = []
    filepath = f"{str(uuid.uuid4())}.csv"
    df.to_csv(filepath, index=False)
    return filepath, bbox_table

# Function to process the uploaded file
def process_file(uploaded_file):
    images_table = convert_pdf_images(uploaded_file)
    croped_table = crop_table(images_table)
    filepath, bbox_table = extract_table(croped_table)
    os.remove(images_table)
    os.remove(croped_table)
    return filepath, bbox_table  # Return the file path for download

# Function to clear the inputs and outputs
def clear_inputs():
    return None, None, None  # Clear both input and output

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Upload a PDF, Process it, and Download the Processed File")

    with gr.Row():
        upload = gr.File(label="Upload PDF", type="filepath", file_types=[".pdf"])
        download = gr.File(label="Download Processed PDF")
    with gr.Row():
        process_button = gr.Button("Process")
        clear_button = gr.Button("Clear")  # Custom clear button
    image_display = gr.Image(label="Processed Image")

    # Trigger the file processing with the button click
    process_button.click(process_file, inputs=upload, outputs=[download, image_display])

    # Trigger clearing inputs and outputs
    clear_button.click(clear_inputs, inputs=None, outputs=[upload, download, image_display])

# Launch the interface
demo.launch()

# print(process_file("/content/give me a example table - give me a example table.pdf"))