Update TextProcessor.py

TextProcessor.py

@@ -4,27 +4,30 @@ import fitz
 import re
 import cv2
 
+
 def paragraphs(folderpath):
     doc = fitz.open(folderpath + '/opinion.pdf')
     df = pd.read_csv(folderpath + '/data.csv').replace({np.nan: None})
     indices = list(df.index)
     pg_indices = df['Pg Ind'].tolist()
 
-    x1s, y1s, x2s, y2s, line_texts, line_inds, pg_inds, baselines = [], [], [], [], [], [], [], {}
+    x1s, y1s, x2s, y2s, line_texts, line_inds, pg_inds, baselines, rights = [], [], [], [], [], [], [], {}, {}
     paras = []
     for (i, pg_ind) in enumerate(pg_indices):
         lines = eval(df[df['Pg Ind'] == i]['Lines'].tolist()[0])
-        pg_x1s = []
-        for (j,n) in enumerate(lines):
+        pg_x1s, pg_x2s = [], []
+        for (j, n) in enumerate(lines):
             x1s.append(n[0])
             y1s.append(n[1])
             x2s.append(n[2])
             y2s.append(n[3])
             line_texts.append(n[4])
             pg_x1s.append(n[0])
+            pg_x2s.append(n[2])
             pg_inds.append(i)
             line_inds.append(j)
         baselines[i] = min(pg_x1s)
+        rights[i] = max(pg_x2s)
 
     for (j, line_text) in enumerate(line_texts):
         if j == 0:
@@ -32,29 +35,28 @@ def paragraphs(folderpath):
             continue
 
         if len(line_texts[j]) > 0:
-            prior_endswith_period = re.search('\.([^A-z]{0,2})$', line_texts[j-1].strip()) is not None
-
-            prior_is_section_header = re.search('^([ABCDEIVX]+)$', line_texts[j-1].strip()) is not None
+            prior_endswith_period = re.search('\.([^A-z]{0,2})$', line_texts[j - 1].strip()) is not None
+            prior_is_section_header = re.search('^([ABCDEIVX]+)$', line_texts[j - 1].strip()) is not None
             current_is_section_header = re.search('^([ABCDEIVX]+)$', line_texts[j].strip()) is not None
-            prior_is_asterisk = re.search('^([\s\*]+)$', line_texts[j-1].strip()) is not None
+            prior_is_asterisk = re.search('^([\s\*]+)$', line_texts[j - 1].strip()) is not None
             current_is_asterisk = re.search('^([\s\*]+)$', line_texts[j].strip()) is not None
-            prior_is_date = re.search('(\[[A-z\s0-9]*,\s[0-9]*]+)$', line_texts[j-1].strip()) is not None
+            prior_is_date = re.search('(\[[A-z\s0-9]*,\s[0-9]*]+)$', line_texts[j - 1].strip()) is not None
 
-            current_upper = line_text[0].isupper()
-            current_tabbed = x1s[j]-baselines[pg_inds[j]] > 7
-            prior_tabbed = x1s[j-1]-baselines[pg_inds[j-1]] > 7
-            prior_supertabbed = x1s[j-1]-baselines[pg_inds[j-1]] > 18
-            current_supertabbed = x1s[j]-baselines[pg_inds[j]] > 18
-            prior_more_left = x1s[j]-x1s[j-1] > 7
+            current_tabbed = x1s[j] - baselines[pg_inds[j]] > 7
+            prior_tabbed = x1s[j - 1] - baselines[pg_inds[j - 1]] > 7
+            prior_supertabbed = x1s[j - 1] - baselines[pg_inds[j - 1]] > 18
+            current_supertabbed = x1s[j] - baselines[pg_inds[j]] > 18
+            prior_more_left = x1s[j] - x1s[j - 1] > 7
+            prior_right_margin = abs(x2s[j - 1] - rights[pg_inds[j - 1]]) < 10
 
             is_section_header = (prior_is_section_header or current_is_section_header or prior_is_asterisk or current_is_asterisk or prior_is_date)
-            prior_period_current_tabbed = (prior_endswith_period and current_tabbed and (prior_more_left or (not prior_supertabbed and not current_supertabbed)))
+            prior_period_current_tabbed = (prior_endswith_period and current_tabbed and (prior_more_left or prior_right_margin or (not prior_supertabbed and not current_supertabbed)))
 
             if is_section_header or prior_period_current_tabbed:
                 paras.append(para)
                 para = []
             #     print('\n')
-            # print(str(j) + ':\t' + str(prior_endswith_period) + '\t' + str(current_tabbed)  + '\t' + str(prior_more_left)+ '\t' + str(prior_supertabbed) + '\t' + str(current_supertabbed) +'\t' + line_text)
+            # print(str(j) + ':\t' + str(prior_endswith_period) + '\t' + str(current_tabbed)  + '\t' + str(prior_more_left)+ '\t' + str(prior_supertabbed) + '\t' + str(prior_right_margin)+ '\t' + str(current_supertabbed) +'\t' + line_text)
             para.append((pg_inds[j], line_inds[j], line_text))
     paras.append(para)
     paras_df = pd.DataFrame({'Lines': paras})
@@ -78,4 +80,4 @@ def process_file(folderpath):
         x1, y1, x2, y2, para_first_line, pg_ind = indent
         image = cv2.imread(folderpath + '/' + str(pg_ind) + '-processed.png')
         cv2.circle(image, (x1 - 15, int(0.5 * (y1 + y2))), radius=1, color=(240, 32, 160), thickness=2)
-        cv2.imwrite(folderpath + '/' + str(pg_ind) + '-processed.png', image)
+        cv2.imwrite(folderpath + '/' + str(pg_ind) + '-processed.png', image)