binery/Bank_Statement_Parser
2
1from paddleocr import PaddleOCR2import numpy as np3import pandas as pd4 5ocr = PaddleOCR(use_angle_cls=True, lang="ch")6 7 8def filter_columns(columns: np.ndarray):9 for idx, col in enumerate(columns):10 if idx >= len(columns) - 1:11 break12 nxt = columns[idx + 1]13 threshold = ((col[2] - col[0]) + (nxt[2] - nxt[0])) / 214 if (col[2] - columns[idx + 1][0]) > threshold * 0.5:15 col[1], col[2], col[3] = min(col[1], nxt[1]), nxt[2], max(col[3], nxt[3])16 columns = np.delete(columns, idx + 1, 0)17 idx -= 118 return columns19 20 21def process_text(row):22 # * concatenate the text of the cell and return the coordinates and the text of the cell23 coor = np.array([None, None])24 text = ""25 for txt in row:26 coor[0], coor[1] = (27 txt[0][0][1] if coor[0] is None or txt[0][0][1] < coor[0] else coor[0],28 txt[0][2][1] if coor[1] is None or txt[0][2][1] > coor[1] else coor[1],29 )30 text += f"{txt[1][0]} "31 text = text.strip()32 row = [coor, text]33 return row34 35 36def extract_text_of_col(col_img):37 """'38 * extract text from the column image and calculate the average length of the row in the column39 * the average is calculated by summing the length of each row then divide the total by the number of rows inside the column40 * return the text and the average length41 """42 result = ocr.ocr(col_img, cls=False)43 ocr_res = []44 45 for ps, (text, score) in result[0]:46 x1 = min(p[0] for p in ps)47 y1 = min(p[1] for p in ps)48 x2 = max(p[0] for p in ps)49 y2 = max(p[1] for p in ps)50 word_info = {51 'bbox': [x1, y1, x2, y2],52 'text': text53 }54 ocr_res.append(word_info)55 threshold = 056 print(result)57 for idx in range(len(result)):58 summ = 059 length = len(result[idx])60 for line in result[idx]:61 summ += line[0][2][1] - line[0][0][1]62 if length > 0:63 threshold += summ / len(result[idx])64 return result, threshold / len(result),ocr_res65 66 67def prepare_cols(result, threshold):68 """69 ** columns are seperated **70 * add each element from the extracted text to its row according to the coordinate intersection with respect to the average length of the row71 * the intersection is True if the intersected part is bigger than the threshold number (ex: half of the average length of the row)72 * return the column of the arranged rows73 """74 col = []75 for idx in range(len(result)):76 row = []77 for i, line in enumerate(result[idx]):78 if i == 0:79 row.append(line)80 if i == len(result[idx]) - 1:81 col.append(process_text(row))82 continue83 if (84 line[0][0][1] >= row[-1][0][0][1] and line[0][2][1] >= row[-1][0][2][1]85 ) and (86 line[0][2][1] > row[-1][0][0][1]87 and line[0][0][1] < row[-1][0][2][1]88 and (abs(line[0][0][1] - row[-1][0][2][1]) > threshold)89 ):90 row.append(line)91 elif (92 line[0][0][1] <= row[-1][0][0][1] and line[0][2][1] <= row[-1][0][2][1]93 ) and (94 line[0][2][1] > row[-1][0][0][1]95 and line[0][0][1] < row[-1][0][2][1]96 and (abs(line[0][2][1] - row[-1][0][0][1]) > threshold)97 ):98 row.append(line)99 elif (100 line[0][0][1] <= row[-1][0][0][1] and line[0][2][1] >= row[-1][0][2][1]101 ) and (102 line[0][2][1] > row[-1][0][0][1]103 and line[0][0][1] < row[-1][0][2][1]104 and (abs(row[-1][0][2][1] - row[-1][0][0][1]) > threshold)105 ):106 row.append(line)107 elif (108 line[0][0][1] >= row[-1][0][0][1] and line[0][2][1] <= row[-1][0][2][1]109 ) and (110 line[0][2][1] > row[-1][0][0][1]111 and line[0][0][1] < row[-1][0][2][1]112 and (abs(line[0][0][1] - line[0][2][1]) > threshold)113 ):114 row.append(line)115 elif (116 line[0][0][1] == row[-1][0][0][1] and line[0][2][1] == row[-1][0][2][1]117 ) and (118 line[0][2][1] > row[-1][0][0][1]119 and line[0][0][1] < row[-1][0][2][1]120 and (abs(line[0][2][1] - row[-1][0][0][1]) > threshold)121 ):122 row.append(line)123 else:124 col.append(process_text(row))125 row = [line]126 if i == len(result[idx]) - 1:127 col.append(process_text(row))128 return col129 130 131def prepare_coordinates(cols):132 """133 * find the column with the maximum number of rows134 * create a key value pair in which the key is the coordinates of each row in the column with the highest number of rows135 and the value is an empty numpy array which has length of number of detected columns136 """137 max_col = max(cols, key=len)138 array = np.empty(len(cols), dtype=object)139 array.fill(np.nan)140 coor_dict = {tuple(k[0]): array for k in max_col}141 return coor_dict142 143 144def process_cols(cols, threshold):145 coor_dict = prepare_coordinates(cols)146 """147 * loop over each element inside each column and find the right place for it inside the dataframe by using the coordinates intersection with respect to the average length of the row148 * the intersection is True if the intersected part is bigger than the threshold number (ex: half of the average length of the row)149 """150 for idx, col in enumerate(cols):151 for element in col:152 for coor, row in coor_dict.items():153 if (coor[0] >= element[0][0] and coor[1] >= element[0][1]) and (154 (coor[1] > element[0][0])155 and (coor[0] < element[0][1])156 and (abs(coor[0] - element[0][1]) > threshold)157 ):158 new = row.copy()159 new[idx] = element[1]160 coor_dict[coor] = new161 elif (coor[0] <= element[0][0] and coor[1] <= element[0][1]) and (162 (coor[1] > element[0][0])163 and (coor[0] < element[0][1])164 and (abs(coor[1] - element[0][0]) > threshold)165 ):166 new = row.copy()167 new[idx] = element[1]168 coor_dict[coor] = new169 elif (coor[0] >= element[0][0] and coor[1] <= element[0][1]) and (170 (coor[1] > element[0][0])171 and (coor[0] < element[0][1])172 and (abs(coor[1] - coor[0]) > threshold)173 ):174 new = row.copy()175 new[idx] = element[1]176 coor_dict[coor] = new177 elif (coor[0] <= element[0][0] and coor[1] >= element[0][1]) and (178 (coor[1] > element[0][0])179 and (coor[0] < element[0][1])180 and (abs(element[0][1] - element[0][0]) > threshold)181 ):182 new = row.copy()183 new[idx] = element[1]184 coor_dict[coor] = new185 data = [row for row in coor_dict.values()]186 return data187 188 189def valid_row(row):190 return (191 (row[0] is not np.nan)192 or (row[-1] is not np.nan)193 or (row[-2] is not np.nan)194 or (row[-3] is not np.nan)195 )196 197 198def finalize_data(data: list, page_enumeration: int):199 idx = 0200 while idx <= len(data) - 1:201 row = data[idx]202 if valid_row(row) and row[0] is np.nan:203 # * add the date to the valid row if it's empty204 try:205 row[0] = data[idx - 1][0]206 data[idx] = row207 except:208 data.pop(idx)209 idx = (idx - 1) if idx > 0 else idx210 continue211 if not valid_row(row):212 if idx == 0:213 data.pop(idx)214 continue215 for i, col in enumerate(row):216 # * merge description to the previous row if the current row is not valid217 if (col is not None) and (col is not np.nan):218 data[idx - 1][i] = str(data[idx - 1][i]) + f" {col}"219 data.pop(idx)220 idx -= 1221 continue222 idx += 1223 page_idx = ["page"] + [page_enumeration for i in range(len(data) - 1)]224 data: pd.DataFrame = pd.DataFrame(data)225 data.insert(0, "page", page_idx)226 return data227 