File size: 6,908 Bytes
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 | 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 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)
# header_bbox = (145, 45, 465, 155) # For digitized variants
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])
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)
return page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, 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': [],
}
data_df = pd.DataFrame(data)
for (i,f) in enumerate(files):
page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, 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]]
}
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)
|