Spaces:
Running
Running
| 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")) |