merge hugo into main; deskew already merged

.gitignore

@@ -1,4 +1,5 @@
 __pycache__/
 *.egg-info/
 deskew_model/
-
+/pdfs/
+/results/

vidocp/layout_detection.py

@@ -23,6 +23,7 @@ def find_layout_boxes(image: np.array):
 
     contours = cv2.findContours(img_bin_final, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     contours = imutils.grab_contours(contours)
+
     for c in contours:
         peri = cv2.arcLength(c, True)
         approx = cv2.approxPolyDP(c, 0.04 * peri, True)

vidocp/table_parsing.py

@@ -1,15 +1,22 @@
+from functools import partial
+from itertools import chain, starmap
+from operator import attrgetter
+
 import cv2
 import numpy as np
 from pdf2image import pdf2image
 
 from vidocp.utils.display import show_mpl
-from vidocp.utils.draw import draw_stats
+from vidocp.utils.draw import draw_rectangles
+from vidocp.utils.post_processing import xywh_to_vecs, xywh_to_vec_rect, adjacent1d, remove_isolated
 from vidocp.utils.deskew import deskew_histbased
+from vidocp.layout_parsing import parse_layout
 
 
 def add_external_contours(image, img):
+    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
 
-    contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+    # contours = filter(partial(is_large_enough, min_area=5000000), contours)
 
     for cnt in contours:
         x, y, w, h = cv2.boundingRect(cnt)
@@ -18,42 +25,131 @@ def add_external_contours(image, img):
     return image
 
 
-def isolate_vertical_and_horizontal_components(img_bin):
-
-    line_min_width = 30
+def isolate_vertical_and_horizontal_components(img_bin, bounding_rects):
+    line_min_width = 48
     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)
+    show_mpl(img_bin_h | img_bin_v)
+
+    kernel_h = np.ones((1, 30), np.uint8)
+    kernel_v = np.ones((30, 1), np.uint8)
+    img_bin_h = cv2.dilate(img_bin_h, kernel_h, iterations=2)
+    img_bin_v = cv2.dilate(img_bin_v, kernel_v, iterations=2)
+    show_mpl(img_bin_h | img_bin_v)
+
+    #reduced filtersize from 100 to 80 to minimize splitting narrow cells
+    img_bin_h = apply_motion_blur(img_bin_h, 80, 0)
+    img_bin_v = apply_motion_blur(img_bin_v, 80, 90)
 
     img_bin_final = img_bin_h | img_bin_v
+    show_mpl(img_bin_final)
+    #changed threshold from 110 to 120 to minimize cell splitting
+    th1, img_bin_final = cv2.threshold(img_bin_final, 120, 255, cv2.THRESH_BINARY)
+    img_bin_final = cv2.dilate(img_bin_final, np.ones((1, 1), np.uint8), iterations=1)
+    show_mpl(img_bin_final)
+    # problem if layout parser detects too big of a layout box as in VV-748542.pdf p.22
+    img_bin_final = disconnect_non_existing_cells(img_bin_final, bounding_rects)
+    show_mpl(img_bin_final)
 
     return img_bin_final
 
 
+def disconnect_non_existing_cells(img_bin, bounding_rects):
+    for rect in bounding_rects:
+        x, y, w, h = rect
+        img_bin = cv2.rectangle(img_bin, (x, y), (x + w, y + h), (0, 0, 0), 5)
+    return img_bin
+
+
+# FIXME: does not work yet
+def has_table_shape(rects):
+    assert isinstance(rects, list)
+
+    points = list(chain(*map(xywh_to_vecs, rects)))
+    brect = xywh_to_vec_rect(cv2.boundingRect(np.vstack(points)))
+
+    rects = list(map(xywh_to_vec_rect, rects))
+
+    # print(rects)
+    # print(brect)
+
+    def matches_bounding_rect_corner(rect, x, y):
+        corresp_coords = list(zip(*map(attrgetter(x, y), [brect, rect])))
+        ret = all(starmap(partial(adjacent1d, tolerance=30), corresp_coords))
+        # print()
+        # print(x, y)
+        # print(brect)
+        # print(rect)
+        # print(corresp_coords)
+        # print(ret)
+
+        return ret
+
+    return all(
+        (
+            any(matches_bounding_rect_corner(r, "xmin", "ymin") for r in rects),
+            any(matches_bounding_rect_corner(r, "xmin", "ymax") for r in rects),
+            any(matches_bounding_rect_corner(r, "xmax", "ymax") for r in rects),
+            any(matches_bounding_rect_corner(r, "xmax", "ymin") for r in rects),
+        )
+    )
+
+
+def apply_motion_blur(image, size, angle):
+    k = np.zeros((size, size), dtype=np.float32)
+    k[(size - 1) // 2, :] = np.ones(size, dtype=np.float32)
+    k = cv2.warpAffine(k, cv2.getRotationMatrix2D((size / 2 - 0.5, size / 2 - 0.5), angle, 1.0), (size, size))
+    k = k * (1.0 / np.sum(k))
+    return cv2.filter2D(image, -1, k)
+
+
+def find_table_layout_boxes(image: np.array):
+    layout_boxes = parse_layout(image)
+    table_boxes = []
+    for box in layout_boxes:
+        (x, y, w, h) = box
+        if w * h >= 100000:
+            table_boxes.append(box)
+    return table_boxes
+
+
 def parse_table(image: np.array):
+    def is_large_enough(stat):
+        x1, y1, w, h, area = stat
+        return area > 2000 and w > 35 and h > 25
 
     gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape)>2 else image
-    th1, img_bin = cv2.threshold(gray_scale, 150, 255, cv2.THRESH_BINARY)
+    th1, img_bin = cv2.threshold(gray_scale, 195, 255, cv2.THRESH_BINARY)
     img_bin = ~img_bin
+    show_mpl(img_bin)
 
-    img_bin = isolate_vertical_and_horizontal_components(img_bin)
+    table_layout_boxes = find_table_layout_boxes(image)
+    img_bin = isolate_vertical_and_horizontal_components(img_bin, table_layout_boxes)
     img_bin_final = add_external_contours(img_bin, img_bin)
 
-    _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S)
+    _, _, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S)
 
-    return stats
+    stats = np.vstack(list(filter(is_large_enough, stats)))
+    rects = stats[:, :-1][2:]
+
+    # FIXME: produces false negatives for `data0/043d551b4c4c768b899eaece4466c836.pdf 1 --type table`
+    rects = list(remove_isolated(rects, input_sorted=True))
+
+    return rects
 
 
 def annotate_tables_in_pdf(pdf_path, page_index=0, deskew=False):
-
     page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0]
     page = np.array(page)
     if deskew:
         page = deskew_histbased(page)
 
     stats = parse_table(page)
-    page = draw_stats(page, stats)
+    page = draw_rectangles(page, stats, annotate=True)
+    # if stats:
+    #     page = draw_rectangles(page, stats, annotate=True)
 
     show_mpl(page)

vidocp/utils/draw.py

@@ -13,7 +13,9 @@ def draw_contours(image, contours):
     return image
 
 
-def draw_rectangles(image, rectangles, color=None):
+def draw_rectangles(image, rectangles, color=None, annotate=False):
+    def annotate_rect(x, y, w, h):
+        cv2.putText(image, "+", (x + (w // 2) - 12, y + (h // 2) + 9), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
 
     image = copy_and_normalize_channels(image)
 
@@ -24,33 +26,7 @@ def draw_rectangles(image, rectangles, color=None):
         x, y, w, h = rect
         cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
 
-    return image
-
-
-def draw_stats(image, stats, annotate=False):
-
-    image = copy_and_normalize_channels(image)
-
-    keys = ["x", "y", "w", "h"]
-
-    def annotate_stat(x, y, w, h):
-
-        for i, (s, v) in enumerate(zip(keys, [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, (0, 255, 0), 2)
-
-    def draw_stat(stat):
-
-        x, y, w, h, area = stat
-
-        cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
-
         if annotate:
-            annotate_stat(x, y, w, h)
-
-    for stat in stats[2:]:
-        draw_stat(stat)
+            annotate_rect(x, y, w, h)
 
     return image