current status of trying to connect layout and table parsing

table_parsing/table_parsig.py

@@ -9,19 +9,22 @@ from matplotlib import pyplot as plt
 
 def parse(image: np.array):
     gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    blurred = cv2.GaussianBlur(gray_scale, (5, 5), 1)
-    thresh = cv2.threshold(blurred, 253, 255, cv2.THRESH_BINARY)[1]
+    #plt.imshow(gray_scale)
+    blurred = cv2.GaussianBlur(gray_scale, (7, 7), 2)  #5 5 1
+    thresh = cv2.threshold(blurred, 251, 255, cv2.THRESH_BINARY)[1]
+    #plt.imshow(thresh)
     img_bin = ~thresh
 
-    line_min_width = 10
+    line_min_width = 7
     kernel_h = np.ones((10, line_min_width), np.uint8)
     kernel_v = np.ones((line_min_width, 10), np.uint8)
 
     img_bin_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h)
     img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v)
-
+    #plt.imshow(img_bin_h)
+    #plt.imshow(img_bin_v)
     img_bin_final = img_bin_h | img_bin_v
-
+    plt.imshow(img_bin_final)
     contours = cv2.findContours(img_bin_final, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     contours = imutils.grab_contours(contours)
     for c in contours:
@@ -29,25 +32,48 @@ def parse(image: np.array):
         approx = cv2.approxPolyDP(c, 0.04 * peri, True)
         yield cv2.boundingRect(approx)
 
-def parse_tables(image: np.array, rectangle):
-    gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    th1, img_bin = cv2.threshold(gray_scale, 200, 255, cv2.THRESH_BINARY)
-    img_bin = ~img_bin
+def parse_tables(image: np.array, rects: list):
+    parsed_tables = []
+    for rect in rects:
+        (x,y,w,h) = rect
+        region_of_interest = image[x:x+w, y:y+h]
+        gray = cv2.cvtColor(region_of_interest, cv2.COLOR_BGR2GRAY)
+        thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)[1]
+        img_bin = ~thresh
 
-    line_min_width = 5
-    kernel_h = np.ones((1, line_min_width), np.uint8)
-    kernel_v = np.ones((line_min_width, 1), np.uint8)
+        line_min_width = 5
+        kernel_h = np.ones((1, line_min_width), np.uint8)
+        kernel_v = np.ones((line_min_width, 1), np.uint8)
 
-    img_bin_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h)
-    img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v)
+        img_bin_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h)
+        img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v)
     # find_and_close_internal_gaps(img_bin_v)
-    img_bin_final = img_bin_h | img_bin_v
-    plt.imshow(img_bin_final)
+        img_bin_final = img_bin_h | img_bin_v
+    #plt.imshow(img_bin_final)
     # find_and_close_internal_gaps(img_bin_final)
     # find_and_close_edges(img_bin_final)
 
-    _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S)
-    return labels, stats
+        _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S)
+        parsed_tables.append([(x,y,w,h), stats])
+    return parsed_tables
+        #yield (x,y,w,h), stats, region_of_interest
+    # return stats
+
+def annotate_table(image, parsed_tables):
+    for table in parsed_tables:
+        original_coordinates, stats = table
+        stats = filter_unconnected_cells(stats)
+        for stat in stats:
+            x, y, w, h, area = stat
+            cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 255), 2)
+            for i, (s, v) in enumerate(zip(["x", "y", "w", "h"], [x, y, w, h])):
+                anno = f"{s} = {v}"
+                xann = int(x + 5)
+                yann = int(y + h - (20 * (i + 1)))
+                cv2.putText(image, anno, (xann, yann), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2)
+
+    return image
+
 
 
 def filter_unconnected_cells(stats):
@@ -87,18 +113,7 @@ def find_and_close_edges(img_bin_final):
 
     return img_bin_final
 
-def annotate_image(image, stats):
-    stats = filter_unconnected_cells(stats)
-    for stat in stats:
-        x, y, w, h, area = stat
-        cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 255), 2)
-        for i, (s, v) in enumerate(zip(["x", "y", "w", "h"], [x, y, w, h])):
-            anno = f"{s} = {v}"
-            xann = int(x + 5)
-            yann = int(y + h - (20 * (i + 1)))
-            cv2.putText(image, anno, (xann, yann), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2)
 
-    return image
 
 def parse_tables_in_pdf(pages):
     return zip(map(parse, pages), count())
@@ -118,19 +133,36 @@ def parse_tables_in_pdf(pages):
 
 
 def annotate_boxes(image, rects):
+    print(type(rects))
     for rect in rects:
         (x, y, w, h) = rect
         cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
 
     return image
 
+def filter_tables_or_images(rects):
+    filtered = []
+    for rect in rects:
+        (x,y,w,h) = rect
+        print(w*h)
+        if w * h > 10**6:
+            filtered.append(rect)
+    print(filtered)
+    return filtered
+
+
+
 
 def annotate_tables_in_pdf(pdf_path, page_index=1):
     page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0]
     page = np.array(page)
 
-    asd = parse(page)
-    page = annotate_boxes(page, asd)
+    layout_boxes = parse(page)
+    page = annotate_boxes(page, layout_boxes)
+    parsed_tables = parse_tables(page, filter_tables_or_images(layout_boxes))
+    page = annotate_table(page, parsed_tables)
+
+
 
     fig, ax = plt.subplots(1, 1)
     fig.set_size_inches(20, 20)