limitedonly41/Table-Structure-Recognition-Demo
0
1import matplotlib.pyplot as plt2import matplotlib.patches as patches3from matplotlib.patches import Patch4import io5import cv26from PIL import Image, ImageDraw, ImageFont7import numpy as np8import csv9import pandas as pd10 11from ultralytics import YOLO12import torch13 14from paddleocr import PaddleOCR15import postprocess16 17import gradio as gr18 19 20device = "cuda" if torch.cuda.is_available() else "cpu"21detection_model = YOLO('yolov8/runs/detect/yolov8s-custom-detection/weights/best.pt').to(device)22structure_model = YOLO('yolov8/runs/detect/yolov8s-custom-structure-all/weights/best.pt').to(device)23ocr_model = PaddleOCR(use_angle_cls=True, lang="uk", det_limit_side_len=1920) # TODO use large det_limit_side_len to get better OCR result24 25detection_class_names = ['table', 'table rotated']26structure_class_names = [27 'table', 'table column', 'table row', 'table column header',28 'table projected row header', 'table spanning cell', 'no object'29]30structure_class_map = {k: v for v, k in enumerate(structure_class_names)}31structure_class_thresholds = {32 "table": 0.5,33 "table column": 0.5,34 "table row": 0.5,35 "table column header": 0.5,36 "table projected row header": 0.5,37 "table spanning cell": 0.5,38 "no object": 1039}40 41 42def table_detection(image):43 imgsz = 80044 pred = detection_model.predict(image, imgsz=imgsz)45 pred = pred[0].boxes46 result = pred.cpu().numpy()47 result_list = [list(result.xywhn[i]) + [result.conf[i], result.cls[i]] for i in range(result.shape[0])]48 return result_list49 50 51def table_structure(image):52 imgsz = 102453 pred = structure_model.predict(image, imgsz=imgsz)54 pred = pred[0].boxes55 result = pred.cpu().numpy()56 result_list = [list(result.xywhn[i]) + [result.conf[i], result.cls[i]] for i in range(result.shape[0])]57 return result_list58 59 60def crop_image(image, detection_result):61 # crop_filenames = []62 width = image.shape[1]63 height = image.shape[0]64 # print(width, height)65 crop_image = image66 for i, result in enumerate(detection_result[:1]): # TODO only return first detected table67 class_id = int(result[5])68 score = float(result[4])69 min_x = result[0]70 min_y = result[1]71 w = result[2]72 h = result[3]73 74 # x1 = max(0, int((min_x-w/2-0.02)*width)) # TODO expand 2%75 # y1 = max(0, int((min_y-h/2-0.02)*height)) # TODO expand 2%76 # x2 = min(width, int((min_x+w/2+0.02)*width)) # TODO expand 2%77 # y2 = min(height, int((min_y+h/2+0.02)*height)) # TODO expand 2%78 x1 = max(0, int((min_x-w/2)*width)-10) # TODO expand 10px79 y1 = max(0, int((min_y-h/2)*height)-10) # TODO expand 10px80 x2 = min(width, int((min_x+w/2)*width)+10) # TODO expand 10px81 y2 = min(height, int((min_y+h/2)*height)+10) # TODO expand 10px82 # print(x1, y1, x2, y2)83 crop_image = image[y1:y2, x1:x2, :]84 # crop_filename = filename[:-4]+'_'+str(i)+'_'+detection_class_names[class_id]+filename[-4:]85 # crop_filenames.append(crop_filename)86 # cv2.imwrite(crop_filename, crop_image)87 return crop_image88 89 90def convert_stucture(ocr_result, image, structure_result):91 width = image.shape[1]92 height = image.shape[0]93 # print(width, height)94 95 bboxes = []96 scores = []97 labels = []98 for i, result in enumerate(structure_result):99 class_id = int(result[5])100 score = float(result[4])101 min_x = result[0]102 min_y = result[1]103 w = result[2]104 h = result[3]105 106 x1 = int((min_x-w/2)*width)107 y1 = int((min_y-h/2)*height)108 x2 = int((min_x+w/2)*width)109 y2 = int((min_y+h/2)*height)110 # print(x1, y1, x2, y2)111 112 bboxes.append([x1, y1, x2, y2])113 scores.append(score)114 labels.append(class_id)115 116 table_objects = []117 for bbox, score, label in zip(bboxes, scores, labels):118 table_objects.append({'bbox': bbox, 'score': score, 'label': label})119 # print('table_objects:', table_objects)120 121 table = {'objects': table_objects, 'page_num': 0}122 123 table_class_objects = [obj for obj in table_objects if obj['label'] == structure_class_map['table']]124 if len(table_class_objects) > 1:125 table_class_objects = sorted(table_class_objects, key=lambda x: x['score'], reverse=True)126 try:127 table_bbox = list(table_class_objects[0]['bbox'])128 except:129 table_bbox = (0,0,1000,1000)130 # print('table_class_objects:', table_class_objects)131 # print('table_bbox:', table_bbox)132 133 page_tokens = ocr_result134 tokens_in_table = [token for token in page_tokens if postprocess.iob(token['bbox'], table_bbox) >= 0.5]135 # print('tokens_in_table:', tokens_in_table)136 137 table_structures, cells, confidence_score = postprocess.objects_to_cells(table, table_objects, tokens_in_table, structure_class_names, structure_class_thresholds)138 139 return table_structures, cells, confidence_score140 141 142def visualize_cells(image, table_structures, cells):143 width = image.shape[1]144 height = image.shape[0]145 # print(width, height)146 empty_image = np.zeros((height, width, 3), np.uint8)147 empty_image.fill(255)148 empty_image = Image.fromarray(cv2.cvtColor(empty_image, cv2.COLOR_BGR2RGB))149 draw = ImageDraw.Draw(empty_image)150 fontStyle = ImageFont.truetype("SimSong.ttc", 10, encoding="utf-8")151 152 num_cols = len(table_structures['columns'])153 num_rows = len(table_structures['rows'])154 data_rows = [['' for _ in range(num_cols)] for _ in range(num_rows)]155 for i, cell in enumerate(cells):156 bbox = cell['bbox']157 x1 = int(bbox[0])158 y1 = int(bbox[1])159 x2 = int(bbox[2])160 y2 = int(bbox[3])161 col_num = cell['column_nums'][0]162 row_num = cell['row_nums'][0]163 spans = cell['spans']164 text = ''165 for span in spans:166 if 'text' in span:167 text += span['text'] 168 data_rows[row_num][col_num] = text169 170 # print('text:', text)171 text_len = len(text)172 # print('text_len:', text_len)173 cell_width = x2-x1174 # print('cell_width:', cell_width)175 num_per_line = cell_width//10176 # print('num_per_line:', num_per_line)177 if num_per_line != 0:178 line_num = text_len//num_per_line179 else:180 line_num = 0181 # print('line_num:', line_num)182 new_text = text[:num_per_line]+'\n'183 for j in range(line_num):184 new_text += text[(j+1)*num_per_line:(j+2)*num_per_line]+'\n'185 # print('new_text:', new_text)186 text = new_text187 188 cv2.rectangle(image, (x1, y1), (x2, y2), color=(0,255,0))189 # cv2.putText(image, str(row_num)+'-'+str(col_num), (x1, y1+30), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,255))190 191 # cv2.rectangle(empty_image, (x1, y1), (x2, y2), color=(0,0,255))192 # cv2.putText(empty_image, str(row_num)+'-'+str(col_num), (x1-10, y1), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,255))193 # cv2.putText(empty_image, text, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,255))194 draw.rectangle([(x1, y1), (x2, y2)], (255,255,255), (0,255,0))195 # draw.text((x1-20, y1), str(row_num)+'-'+str(col_num), (255,0,0), font=fontStyle)196 # draw.text((x1, y1), text, (0,0,255), font=fontStyle)197 198 df = pd.DataFrame(data_rows)199 df.columns = df.columns.astype(str)200 return image, df, df.to_json()201 202 203def ocr(image):204 result = ocr_model.ocr(image, cls=True)205 result = result[0]206 new_result = []207 if result is not None:208 bounding_boxes = [line[0] for line in result]209 txts = [line[1][0] for line in result]210 scores = [line[1][1] for line in result]211 # print('txts:', txts)212 # print('scores:', scores)213 # print('bounding_boxes:', bounding_boxes)214 for label, bbox in zip(txts, bounding_boxes):215 new_result.append({'bbox': [bbox[0][0], bbox[0][1], bbox[2][0], bbox[2][1]], 'text': label})216 217 return new_result218 219 220def detect_and_crop_table(image):221 detection_result = table_detection(image)222 # print('detection_result:', detection_result)223 cropped_table = crop_image(image, detection_result)224 225 return cropped_table226 227 228def recognize_table(image, ocr_result):229 structure_result = table_structure(image)230 print('structure_result:', structure_result)231 table_structures, cells, confidence_score = convert_stucture(ocr_result, image, structure_result)232 print('table_structures:', table_structures)233 print('cells:', cells)234 print('confidence_score:', confidence_score)235 image, df, data = visualize_cells(image, table_structures, cells)236 237 return image, df, data238 239 240def process_pdf(image):241 image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)242 243 cropped_table = detect_and_crop_table(image)244 245 ocr_result = ocr(cropped_table)246 # print('ocr_result:', ocr_result)247 248 image, df, data = recognize_table(cropped_table, ocr_result)249 print('df:', df)250 251 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)252 253 return image, df, data254 255 256title = "Demo: table detection & recognition with Table Structure Recognition (Yolov8)."257description = """Demo for table extraction with the Table Structure Recognition (Yolov8)."""258examples = [['image.png'], ['mistral_paper.png']]259 260app = gr.Interface(fn=process_pdf, 261 inputs=gr.Image(type="numpy"), 262 outputs=[gr.Image(type="numpy", label="Detected table"), gr.Dataframe(label="Table as CSV"), gr.JSON(label="Data as JSON")],263 title=title,264 description=description,265 examples=examples)266app.queue()267# app.launch(debug=True, share=True)268app.launch()