Spaces:
Running
Running
| import matplotlib.pyplot as plt | |
| import matplotlib.patches as patches | |
| from matplotlib.patches import Patch | |
| import io | |
| from PIL import Image, ImageDraw | |
| import numpy as np | |
| import csv | |
| import pandas as pd | |
| from torchvision import transforms | |
| from transformers import AutoModelForObjectDetection | |
| import torch | |
| import easyocr | |
| import gradio as gr | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| class MaxResize(object): | |
| def __init__(self, max_size=800): | |
| self.max_size = max_size | |
| def __call__(self, image): | |
| width, height = image.size | |
| current_max_size = max(width, height) | |
| scale = self.max_size / current_max_size | |
| resized_image = image.resize((int(round(scale*width)), int(round(scale*height)))) | |
| return resized_image | |
| detection_transform = transforms.Compose([ | |
| MaxResize(800), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| structure_transform = transforms.Compose([ | |
| MaxResize(1000), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| # load table detection model | |
| # processor = TableTransformerImageProcessor(max_size=800) | |
| model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm").to(device) | |
| # load table structure recognition model | |
| # structure_processor = TableTransformerImageProcessor(max_size=1000) | |
| structure_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition-v1.1-all").to(device) | |
| # load EasyOCR reader | |
| reader = easyocr.Reader(['en']) | |
| # for output bounding box post-processing | |
| def box_cxcywh_to_xyxy(x): | |
| x_c, y_c, w, h = x.unbind(-1) | |
| b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] | |
| return torch.stack(b, dim=1) | |
| def rescale_bboxes(out_bbox, size): | |
| width, height = size | |
| boxes = box_cxcywh_to_xyxy(out_bbox) | |
| boxes = boxes * torch.tensor([width, height, width, height], dtype=torch.float32) | |
| return boxes | |
| def outputs_to_objects(outputs, img_size, id2label): | |
| m = outputs.logits.softmax(-1).max(-1) | |
| pred_labels = list(m.indices.detach().cpu().numpy())[0] | |
| pred_scores = list(m.values.detach().cpu().numpy())[0] | |
| pred_bboxes = outputs['pred_boxes'].detach().cpu()[0] | |
| pred_bboxes = [elem.tolist() for elem in rescale_bboxes(pred_bboxes, img_size)] | |
| objects = [] | |
| for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes): | |
| class_label = id2label[int(label)] | |
| if not class_label == 'no object': | |
| objects.append({'label': class_label, 'score': float(score), | |
| 'bbox': [float(elem) for elem in bbox]}) | |
| return objects | |
| def fig2img(fig): | |
| """Convert a Matplotlib figure to a PIL Image and return it""" | |
| buf = io.BytesIO() | |
| fig.savefig(buf) | |
| buf.seek(0) | |
| image = Image.open(buf) | |
| return image | |
| def visualize_detected_tables(img, det_tables): | |
| plt.imshow(img, interpolation="lanczos") | |
| fig = plt.gcf() | |
| fig.set_size_inches(20, 20) | |
| ax = plt.gca() | |
| for det_table in det_tables: | |
| bbox = det_table['bbox'] | |
| if det_table['label'] == 'table': | |
| facecolor = (1, 0, 0.45) | |
| edgecolor = (1, 0, 0.45) | |
| alpha = 0.3 | |
| linewidth = 2 | |
| hatch='//////' | |
| elif det_table['label'] == 'table rotated': | |
| facecolor = (0.95, 0.6, 0.1) | |
| edgecolor = (0.95, 0.6, 0.1) | |
| alpha = 0.3 | |
| linewidth = 2 | |
| hatch='//////' | |
| else: | |
| continue | |
| rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=linewidth, | |
| edgecolor='none',facecolor=facecolor, alpha=0.1) | |
| ax.add_patch(rect) | |
| rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=linewidth, | |
| edgecolor=edgecolor,facecolor='none',linestyle='-', alpha=alpha) | |
| ax.add_patch(rect) | |
| rect = patches.Rectangle(bbox[:2], bbox[2]-bbox[0], bbox[3]-bbox[1], linewidth=0, | |
| edgecolor=edgecolor,facecolor='none',linestyle='-', hatch=hatch, alpha=0.2) | |
| ax.add_patch(rect) | |
| plt.xticks([], []) | |
| plt.yticks([], []) | |
| legend_elements = [Patch(facecolor=(1, 0, 0.45), edgecolor=(1, 0, 0.45), | |
| label='Table', hatch='//////', alpha=0.3), | |
| Patch(facecolor=(0.95, 0.6, 0.1), edgecolor=(0.95, 0.6, 0.1), | |
| label='Table (rotated)', hatch='//////', alpha=0.3)] | |
| plt.legend(handles=legend_elements, bbox_to_anchor=(0.5, -0.02), loc='upper center', borderaxespad=0, | |
| fontsize=10, ncol=2) | |
| plt.gcf().set_size_inches(10, 10) | |
| plt.axis('off') | |
| return fig | |
| def detect_and_crop_table(image): | |
| # prepare image for the model | |
| # pixel_values = processor(image, return_tensors="pt").pixel_values | |
| pixel_values = detection_transform(image).unsqueeze(0).to(device) | |
| # forward pass | |
| with torch.no_grad(): | |
| outputs = model(pixel_values) | |
| # postprocess to get detected tables | |
| id2label = model.config.id2label | |
| id2label[len(model.config.id2label)] = "no object" | |
| detected_tables = outputs_to_objects(outputs, image.size, id2label) | |
| # visualize | |
| # fig = visualize_detected_tables(image, detected_tables) | |
| # image = fig2img(fig) | |
| # crop first detected table out of image | |
| cropped_table = image.crop(detected_tables[0]["bbox"]) | |
| return cropped_table | |
| def recognize_table(image): | |
| # prepare image for the model | |
| # pixel_values = structure_processor(images=image, return_tensors="pt").pixel_values | |
| pixel_values = structure_transform(image).unsqueeze(0).to(device) | |
| # forward pass | |
| with torch.no_grad(): | |
| outputs = structure_model(pixel_values) | |
| # postprocess to get individual elements | |
| id2label = structure_model.config.id2label | |
| id2label[len(structure_model.config.id2label)] = "no object" | |
| cells = outputs_to_objects(outputs, image.size, id2label) | |
| # visualize cells on cropped table | |
| draw = ImageDraw.Draw(image) | |
| for cell in cells: | |
| draw.rectangle(cell["bbox"], outline="red") | |
| return image, cells | |
| def get_cell_coordinates_by_row(table_data): | |
| # Extract rows and columns | |
| rows = [entry for entry in table_data if entry['label'] == 'table row'] | |
| columns = [entry for entry in table_data if entry['label'] == 'table column'] | |
| # Sort rows and columns by their Y and X coordinates, respectively | |
| rows.sort(key=lambda x: x['bbox'][1]) | |
| columns.sort(key=lambda x: x['bbox'][0]) | |
| # Function to find cell coordinates | |
| def find_cell_coordinates(row, column): | |
| cell_bbox = [column['bbox'][0], row['bbox'][1], column['bbox'][2], row['bbox'][3]] | |
| return cell_bbox | |
| # Generate cell coordinates and count cells in each row | |
| cell_coordinates = [] | |
| for row in rows: | |
| row_cells = [] | |
| for column in columns: | |
| cell_bbox = find_cell_coordinates(row, column) | |
| row_cells.append({'column': column['bbox'], 'cell': cell_bbox}) | |
| # Sort cells in the row by X coordinate | |
| row_cells.sort(key=lambda x: x['column'][0]) | |
| # Append row information to cell_coordinates | |
| cell_coordinates.append({'row': row['bbox'], 'cells': row_cells, 'cell_count': len(row_cells)}) | |
| # Sort rows from top to bottom | |
| cell_coordinates.sort(key=lambda x: x['row'][1]) | |
| return cell_coordinates | |
| def apply_ocr(cell_coordinates, cropped_table): | |
| # let's OCR row by row | |
| data = dict() | |
| max_num_columns = 0 | |
| for idx, row in enumerate(cell_coordinates): | |
| row_text = [] | |
| for cell in row["cells"]: | |
| # crop cell out of image | |
| cell_image = np.array(cropped_table.crop(cell["cell"])) | |
| # apply OCR | |
| result = reader.readtext(np.array(cell_image)) | |
| if len(result) > 0: | |
| text = " ".join([x[1] for x in result]) | |
| row_text.append(text) | |
| if len(row_text) > max_num_columns: | |
| max_num_columns = len(row_text) | |
| data[str(idx)] = row_text | |
| # pad rows which don't have max_num_columns elements | |
| # to make sure all rows have the same number of columns | |
| for idx, row_data in data.copy().items(): | |
| if len(row_data) != max_num_columns: | |
| row_data = row_data + ["" for _ in range(max_num_columns - len(row_data))] | |
| data[str(idx)] = row_data | |
| # write to csv | |
| with open('output.csv','w') as result_file: | |
| wr = csv.writer(result_file, dialect='excel') | |
| for row, row_text in data.items(): | |
| wr.writerow(row_text) | |
| # return as Pandas dataframe | |
| df = pd.read_csv('output.csv') | |
| return df, data | |
| def process_pdf(image): | |
| cropped_table = detect_and_crop_table(image) | |
| image, cells = recognize_table(cropped_table) | |
| cell_coordinates = get_cell_coordinates_by_row(cells) | |
| df, data = apply_ocr(cell_coordinates, image) | |
| return image, df, data | |
| title = "Demo: table detection & recognition with Table Transformer (TATR)." | |
| description = """Demo for table extraction with the Table Transformer. First, table detection is performed on the input image using https://huggingface.co/microsoft/table-transformer-detection, | |
| after which the detected table is extracted and https://huggingface.co/microsoft/table-transformer-structure-recognition-v1.1-all is leveraged to recognize the individual rows, columns and cells. OCR is then performed per cell, row by row.""" | |
| examples = [['image.png'], ['mistral_paper.png']] | |
| app = gr.Interface(fn=process_pdf, | |
| inputs=gr.Image(type="pil"), | |
| outputs=[gr.Image(type="pil", label="Detected table"), gr.Dataframe(label="Table as CSV"), gr.JSON(label="Data as JSON")], | |
| title=title, | |
| description=description, | |
| examples=examples) | |
| app.queue() | |
| app.launch(debug=True) |