From 622e2f4fd83180a17f727e768617ea5607e4d244 Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Thu, 20 Jan 2022 15:36:46 +0100
Subject: [PATCH] added first table parsing function: cell detection by
 applying horizontal and vertical kernels

---
 requirements.txt              |  5 +++-
 scripts/annotate_table.py     | 18 +++++++++++++
 table_parsing/__init__.py     |  0
 table_parsing/table_parsig.py | 51 +++++++++++++++++++++++++++++++++++
 4 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 scripts/annotate_table.py
 create mode 100644 table_parsing/__init__.py
 create mode 100644 table_parsing/table_parsig.py

diff --git a/requirements.txt b/requirements.txt
index 391ca2f..7d49ccd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,4 @@
-scikit-image
+opencv-python~=4.5.5.62
+numpy~=1.22.1
+pdf2image~=1.16.0
+matplotlib~=3.5.1
\ No newline at end of file
diff --git a/scripts/annotate_table.py b/scripts/annotate_table.py
new file mode 100644
index 0000000..44e0985
--- /dev/null
+++ b/scripts/annotate_table.py
@@ -0,0 +1,18 @@
+import argparse
+
+from table_parsing.table_parsig import annotate_tables_in_pdf
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("pdf_path")
+    parser.add_argument("page_index", type=int)
+
+    args = parser.parse_args()
+
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    annotate_tables_in_pdf(args.pdf_path, page_index=args.page_index)
diff --git a/table_parsing/__init__.py b/table_parsing/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/table_parsing/table_parsig.py b/table_parsing/table_parsig.py
new file mode 100644
index 0000000..c54aacd
--- /dev/null
+++ b/table_parsing/table_parsig.py
@@ -0,0 +1,51 @@
+from itertools import count
+
+import cv2
+import numpy as np
+import pdf2image
+from matplotlib import pyplot as plt
+
+
+def parse(image: np.array):
+
+    gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    th1, img_bin = cv2.threshold(gray_scale, 150, 225, cv2.THRESH_BINARY)
+    img_bin = ~img_bin
+
+    line_min_width = 4
+    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_final = img_bin_h | img_bin_v
+
+    _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S)
+
+    return labels, stats
+
+
+def parse_tables_in_pdf(pages):
+    return zip(map(parse, pages), count())
+
+
+def annotate_image(image, stats):
+    for x, y, w, h, area in stats[2:]:
+        cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 255), 2)
+
+    return image
+
+
+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)
+
+    _, stats = parse(page)
+    page = annotate_image(page, stats)
+
+    fig, ax = plt.subplots(1, 1)
+    fig.set_size_inches(20, 20)
+    ax.imshow(page)
+    plt.show()