File size: 8,325 Bytes
30420b9 c884cd4 30420b9 c884cd4 30420b9 c884cd4 30420b9 c884cd4 30420b9 c884cd4 30420b9 c884cd4 30420b9 c884cd4 30420b9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
import cv2
import fitz
import numpy as np
import os
import pandas as pd
import pytesseract
import warnings
def pdf2png(folderpath):
doc = fitz.open(folderpath + '/opinion.pdf')
zoom = 1
mat = fitz.Matrix(zoom, zoom)
for (i, p) in enumerate(doc):
pix = p.get_pixmap(matrix=mat)
pix.save(folderpath + '/' + str(i) + '.png')
def is_leftmost(image, x, y_top, y_bot):
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (7,7), 0)
thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
left_portion = thresh[y_top:y_bot, :x]
return np.sum(left_portion) == 0
def get_indents(filename, body_bbox):
doc = fitz.open(filename)
indented_lines = []
for (pg_ind, page) in enumerate(doc):
image = cv2.imread(filename.replace('opinion.pdf', str(pg_ind) + '.png'))
body_rect = fitz.Rect(body_bbox)
pg_dict = page.get_text('dict', clip=body_rect)
all_lines = [(int(line['bbox'][0]), int(line['bbox'][1]), int(line['bbox'][2]), int(line['bbox'][3]), line)for block in pg_dict['blocks'] for line in block['lines']]
body_text = page.get_text("text", clip=body_rect).strip()
baseline = min([l[0] for l in all_lines])
indented_inds = [i for (i,l) in enumerate(all_lines) if (l[0]-baseline > 9 and l[0]-baseline < 30 and is_leftmost(image, l[0]-10, l[1], l[3]))]
for i in indented_inds:
indented_lines.append((i, l[0], l[1], l[2], l[3]))
return indented_lines
def get_footnote_bbox(filename):
footnotes_bbox = (None, None, None, None)
x1p, y1p, x2p, y2p = get_page_bbox(filename)
x1h, y1h, x2h, y2h = get_header_bbox(filename)
image = cv2.imread(filename)
im_h, im_w, im_d = image.shape
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)[1]
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 1))
dilate = cv2.dilate(thresh, kernel, iterations=1)
cnts = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
cnts = sorted(cnts, key=lambda x: cv2.boundingRect(x)[1])
for (i, c) in enumerate(cnts):
x, y, w, h = cv2.boundingRect(c)
if h < 7 and w > 50 and y > y1p and x - x1p < 30:
footnotes_bbox = (x, y, x2p, y2p)
return footnotes_bbox
def get_header_bbox(filename):
image = cv2.imread(filename)
im_h, im_w, im_d = image.shape
base_image = image.copy()
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (9,9), 0)
thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (200,10))
dilate = cv2.dilate(thresh, kernel, iterations=1)
cnts = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
cnts = sorted(cnts, key=lambda x: cv2.boundingRect(x)[1])
for (i,c) in enumerate(cnts):
x,y,w,h = cv2.boundingRect(c)
break
header_bbox = (x, y, x+w, y+40)
return header_bbox
def get_page_bbox(filename):
image = cv2.imread(filename)
im_h, im_w, im_d = image.shape
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (7, 7), 0)
thresh = cv2.threshold(blur, 240, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 10))
dilate = cv2.dilate(thresh, kernel, iterations=1)
cnts = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
cnts = sorted(cnts, key=lambda x: cv2.boundingRect(x)[1])
header_bbox = get_header_bbox(filename)
all_x1 = [cv2.boundingRect(c)[0] for c in cnts]
all_y1 = [cv2.boundingRect(c)[1] for c in cnts]
all_x2 = [cv2.boundingRect(c)[0] + cv2.boundingRect(c)[2] for c in cnts]
all_y2 = [cv2.boundingRect(c)[1] + cv2.boundingRect(c)[3] for c in cnts]
return min(all_x1), header_bbox[1], max(all_x2), max(all_y2)
def get_case_separator(filename):
new_case_line = (None, None, None, None)
x1p, y1p, x2p, y2p = get_page_bbox(filename)
x1h, y1h, x2h, y2h = get_header_bbox(filename)
image = cv2.imread(filename)
im_h, im_w, im_d = image.shape
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (7, 7), 0)
thresh = cv2.threshold(blur, 240, 255, cv2.THRESH_BINARY_INV)[1]
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 1))
dilate = cv2.dilate(thresh, kernel, iterations=1)
cnts = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
cnts = sorted(cnts, key=lambda x: cv2.boundingRect(x)[1])
for (i, c) in enumerate(cnts):
x, y, w, h = cv2.boundingRect(c)
x_center = (x1p + x2p) / 2
if h < 8 and w > 70 and ((x - x1p) < x_center and (x - x1p) > 0.3 * x_center) and (y > y1p and y > y1h): #
new_case_line = (x1p, y, x2p, y)
break
return new_case_line
def get_page_elements(filename):
page_bbox = get_page_bbox(filename)
header_bbox = get_header_bbox(filename)
fn_bbox = get_footnote_bbox(filename)
case_separator_bbox = get_case_separator(filename)
if fn_bbox[0] is not None:
body_bbox = (page_bbox[0], header_bbox[3], page_bbox[2], fn_bbox[1])
else:
body_bbox = (page_bbox[0], header_bbox[3], page_bbox[2], page_bbox[3])
indent_lines = get_indents(filename, body_bbox)
image = cv2.imread(filename)
cv2.rectangle(image, (page_bbox[0], page_bbox[1]), (page_bbox[2], page_bbox[3]), (0, 0, 0), 4)
cv2.rectangle(image, (header_bbox[0], header_bbox[1]), (header_bbox[2], header_bbox[3]), (0, 255, 0), 2)
cv2.rectangle(image, (body_bbox[0], body_bbox[1]), (body_bbox[2], body_bbox[3]), (255, 0, 0), 2)
if fn_bbox[0] is not None:
cv2.rectangle(image, (fn_bbox[0], fn_bbox[1]), (fn_bbox[2], fn_bbox[3]), (0, 0, 255), 2)
if case_separator_bbox[0] is not None:
cv2.rectangle(image, (case_separator_bbox[0], case_separator_bbox[1]),
(case_separator_bbox[2], case_separator_bbox[3]), (255, 0, 255), 2)
for (i, il) in enumerate(il):
cv2.circle(image, (il[1]-10, int(0.5*(il[2] + il[4]))), radius=1, color=(240, 32, 160), thickness=2)
return page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, indent_lines, image
def process_file(folderpath):
pdf2png(folderpath)
files = [f for f in os.listdir(folderpath) if '.png' in f.lower() and "processed" not in f.lower()]
data = {'Pg Ind':[],
'Header X1':[], 'Header Y1': [], 'Header X2': [], 'Header Y2':[],
'Body X1':[], 'Body Y1': [], 'Body X2': [], 'Body Y2':[],
'Footer X1':[], 'Footer Y1': [], 'Footer X2': [], 'Footer Y2':[],
'Page X1':[], 'Page Y1': [], 'Page X2': [], 'Page Y2':[],
'Case Separator Y': [],
'Indent Lines': [],
}
data_df = pd.DataFrame(data)
for (i,f) in enumerate(files):
page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, indent_lines, image = get_page_elements(folderpath +'/' + f)
ind = int(f.split('.png')[0])
row = {'Pg Ind':[ind],
'Header X1':[header_bbox[0]], 'Header Y1': [header_bbox[1]], 'Header X2': [header_bbox[2]], 'Header Y2':[header_bbox[3]],
'Body X1':[body_bbox[0]], 'Body Y1': [body_bbox[1]], 'Body X2': [body_bbox[2]], 'Body Y2':[body_bbox[3]],
'Footer X1':[fn_bbox[0]], 'Footer Y1': [fn_bbox[1]], 'Footer X2': [fn_bbox[2]], 'Footer Y2':[fn_bbox[3]],
'Page X1':[page_bbox[0]], 'Page Y1': [page_bbox[1]], 'Page X2': [page_bbox[2]], 'Page Y2':[page_bbox[3]],
'Case Separator Y': [case_separator_bbox[1]],
'Indent Lines': [indent_lines]
}
row_df = pd.DataFrame(row)
data_df = pd.concat([data_df, row_df], ignore_index=True)
cv2.imwrite(folderpath + '/' + str(ind) + '-processed.png', image)
data_df['Pg Ind'] = data_df['Pg Ind'].astype('int')
data_df.to_csv(folderpath +'/data.csv', index=False)
|