Update ImageProcessor.py

ImageProcessor.py

@@ -18,22 +18,23 @@ 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]
+    left_portion = thresh[int((y_top+y_bot)/2), :x]
     return np.sum(left_portion) == 0
 
-def get_indents(filename, body_bbox):
-    doc = fitz.open(filename)
+def get_indents(filename, body_bbox, page):
+    # 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]))
+
+    # for (pg_ind, page) in enumerate(doc):
+    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 l[0]-baseline < 30 and is_leftmost(image, l[0]-10, l[1], l[3]))]
+    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
 
 def get_footnote_bbox(filename):
@@ -118,7 +119,7 @@ def get_case_separator(filename):
             break
     return new_case_line
 
-def get_page_elements(filename):
+def get_page_elements(filename, page):
     page_bbox = get_page_bbox(filename)
     header_bbox = get_header_bbox(filename)
     fn_bbox = get_footnote_bbox(filename)
@@ -127,7 +128,7 @@ def get_page_elements(filename):
         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)
+    indent_lines = get_indents(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)
@@ -138,13 +139,14 @@ def get_page_elements(filename):
     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)
+    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 process_file(folderpath):
     pdf2png(folderpath)
+    doc = fitz.open(folderpath + '/opinion.pdf')
     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':[],
@@ -156,8 +158,9 @@ def process_file(folderpath):
            }
     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])
+        page = doc[ind]
+        page_bbox, header_bbox, fn_bbox, body_bbox, case_separator_bbox, indent_lines, 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]],
@@ -170,5 +173,4 @@ def process_file(folderpath):
         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)
-
+    data_df.to_csv(folderpath +'/data.csv', index=False)