Removed para-identifier

ImageProcessor.py

@@ -5,6 +5,7 @@ import os
 import pandas as pd
 import pytesseract
 import warnings
+import re
 
 def pdf2png(folderpath):
     doc = fitz.open(folderpath + '/opinion.pdf')
@@ -16,23 +17,27 @@ def pdf2png(folderpath):
 
 def is_leftmost(image, x, y_top, y_bot):
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    blur = cv2.GaussianBlur(gray, (7,7), 0)
+    blur = cv2.GaussianBlur(gray, (9,9), 0)
     thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
     left_portion = thresh[int((y_top+y_bot)/2), :x]
     return np.sum(left_portion) == 0
 
-def get_indents(filename, body_bbox, page):
-    indented_lines = []
+
+def get_line_data(filename, body_bbox, page):
     image = cv2.imread(filename)
     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 is_leftmost(image, l[0]-12, l[1], l[3]))]# and l[0]-baseline < 30
-    for i in indented_inds:
-        indented_lines.append((i, all_lines[i][0], all_lines[i][1], all_lines[i][2], all_lines[i][3]))
-    return indented_lines
+    line_data = []
+    for (i,l) in enumerate(all_lines):
+        if not is_leftmost(image, l[0]-11, l[1], l[3]) and i > 0: # Add it
+            line_data[-1] = list(line_data[-1])
+            line_data[-1][-1] += " " + get_line_text(l[-1])
+            line_data[-1] = tuple(line_data[-1])
+        else:
+            line_data.append((l[0], l[1], l[2], l[3], get_line_text(l[-1])))
+    return line_data
+
 
 def get_footnote_bbox(filename):
     footnotes_bbox = (None, None, None, None)
@@ -125,7 +130,7 @@ def get_page_elements(filename, page):
         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, page)
+    line_data = get_line_data(filename, body_bbox, page)
 
     image = cv2.imread(filename)
     cv2.rectangle(image, (page_bbox[0], page_bbox[1]), (page_bbox[2], page_bbox[3]), (0, 0, 0), 4)
@@ -134,76 +139,17 @@ def get_page_elements(filename, page):
     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(indent_lines):
-        cv2.circle(image, (il[1]-15, int(0.5*(il[2] + il[4]))), radius=1, color=(240, 32, 160), thickness=2)
+        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(indent_lines):
+    #     cv2.circle(image, (il[1]-15, 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 paragraphs(folderpath):
-    doc = fitz.open(folderpath + '/opinion.pdf')
-    df = pd.read_csv(folderpath + '/data.csv')
-    df = df.replace({np.nan: None})
-    nl_inds = df['Indent Lines'].tolist()
-    nl_inds = [eval(nli) for nli in nl_inds]
-    nl_indents = [nli[1] for page_nli in nl_inds for nli in page_nli]
-    nl_inds = [(i, nli[0]) for (i, page_nli) in zip(df['Pg Ind'].tolist(), nl_inds) for nli in page_nli]
-    paras = [([], 0, 0, 0)]  # Text, indent amount, start pg ind, end pg ind
-    para_lines = []
-    for (i, page) in enumerate(doc):
-        ind = df.index[df['Pg Ind'] == i].tolist()[0]
-        body_bbox = [df.iloc[ind]['Body X1'], df.iloc[ind]['Body Y1'], df.iloc[ind]['Body X2'], df.iloc[ind]['Body Y2']]
-        case_separator = df.iloc[ind]['Case Separator Y']
-        if case_separator is not None:
-            body_bbox[-1] = case_separator
-        body_rect = fitz.Rect(body_bbox)
-        pg_dict = page.get_text('dict', clip=body_rect)
-        all_lines = [get_line_text(line) for block in pg_dict['blocks'] for line in block['lines']]
-        for (j, line) in enumerate(all_lines):
-            if line == "":
-                continue
-            if (i, j) in nl_inds:
-                indent_amt = nl_indents[nl_inds.index((i, j))]  # This is for the starting one
-                paras.append(([], indent_amt, i, i))
-            paras[-1] = list(paras[-1])
-            paras[-1][0].append(line.strip())
-            paras[-1][-1] = i  # Update the page ind
-            paras[-1] = tuple(paras[-1])
-    paras = block_quotes(paras)
-    paras_df = pd.DataFrame(data=paras, index=None, columns=['Text', 'Indent Amount', 'Start Pg Ind', 'End Pg Ind'])
-    return paras_df
+    return page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, line_data, image
 
 def get_line_text(line):
     words = []
-    for s in line['spans']:
-        text = s['text'].strip()
-        if text != "":
-            words.append(text)
-    words = " ".join(words)
+    words = "".join(s['text'] for s in line['spans'] if s['text'].strip() != "")
     return words
 
-def block_quotes(paras):
-    modified_paras = []
-    start_para, end_para, end_quote_passed = None, None, None
-    for (i, (para, ind_amt, start_pg_ind, end_pg_ind)) in enumerate(paras):
-        if i == len(paras) - 1:
-            modified_paras.append((para, ind_amt, start_pg_ind, end_pg_ind))
-            break
-        if len(para) == 1 and "“" == para[0][0] and start_para is None:
-            start_para = i
-        if len(para) == 1 and "”" == para[0][-1] and start_para is not None:
-            end_quote_passed = True
-        if len(para) == 1 and (
-                paras[i + 1][1] - ind_amt) < -5 and end_para is None and start_para is not None and end_quote_passed:
-            end_para = i
-        if start_para is not None and end_para is not None:
-            para = [p[0][0] for p in paras[start_para:end_para + 1]]
-            start_para, end_para, end_quote_passed = None, None, False
-        if start_para is None and end_para is None:
-            modified_paras.append((para, ind_amt, start_pg_ind, end_pg_ind))
-    return modified_paras
-
 def process_file(folderpath):
     pdf2png(folderpath)
     doc = fitz.open(folderpath + '/opinion.pdf')
@@ -214,25 +160,23 @@ def process_file(folderpath):
             'Footer X1':[], 'Footer Y1': [], 'Footer X2': [], 'Footer Y2':[],
             'Page X1':[], 'Page Y1': [], 'Page X2': [], 'Page Y2':[],
             'Case Separator Y': [],
-            'Indent Lines': [],
+            'Lines': [],
            }
     data_df = pd.DataFrame(data)
     for (i,f) in enumerate(files):
         ind = int(f.split('.png')[0])
         page = doc[ind]
-        page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, indent_lines, image = get_page_elements(folderpath +'/' + f, page)
+        page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, line_data, image = get_page_elements(folderpath +'/' + f, page)
         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]
+                'Lines': [line_data]
                }
         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)
-    paras_df = paragraphs(folderpath)
-    paras_df.to_csv(folderpath + '/paragraphs.csv', index=False)
+    data_df.to_csv(folderpath +'/data.csv', index=False)