diff --git a/.gitignore b/.gitignore index f3b659b..5360d78 100644 --- a/.gitignore +++ b/.gitignore @@ -13,7 +13,6 @@ build_venv/ /.idea/table_parsing.iml /.idea/vcs.xml /results/ -/data /table_parsing.egg-info /target/ /tests/ @@ -22,3 +21,5 @@ build_venv/ /cv_analysis.egg-info/SOURCES.txt /cv_analysis.egg-info/top_level.txt /.vscode/ +/cv_analysis/test/test_data/example_pages.json +/data/metadata_testing_files.csv diff --git a/config.yaml b/config.yaml index 42bd2e7..fc6bb42 100644 --- a/config.yaml +++ b/config.yaml @@ -23,5 +23,5 @@ deskew: test_dummy: test_dummy visual_logging: - level: $LOGGING_LEVEL_ROOT|INFO + level: $LOGGING_LEVEL_ROOT|DEBUG output_folder: /tmp/debug/ \ No newline at end of file diff --git a/cv_analysis/fig_detection_with_layout.py b/cv_analysis/fig_detection_with_layout.py new file mode 100644 index 0000000..ce1d71b --- /dev/null +++ b/cv_analysis/fig_detection_with_layout.py @@ -0,0 +1,58 @@ +from cv_analysis.layout_parsing import annotate_layout_in_pdf +from cv_analysis.figure_detection import detect_figures +from cv_analysis.table_parsing import tables_in_image, parse_table +from cv_analysis.utils.draw import draw_rectangles +from cv_analysis.utils.display import show_mpl +from cv_analysis.utils.visual_logging import vizlogger +#from PIL import Image + + + +def cut_out_content_structures(layout_rects, page): + large_rects = [] + small_rects = [] + for x, y, w, h in layout_rects: + rect = (x, y, w, h) + if w * h >= 75000: + cropped_page = page[y:(y + h), x:(x + w)] + large_rects.append([rect, cropped_page]) + else: + cropped_page = page[y:(y + h), x:(x + w)] + small_rects.append([rect, cropped_page]) + return large_rects, small_rects + + +def parse_content_structures(page, large_rects, small_rects): + for coordinates, cropped_image in large_rects: + figure_rects = detect_figures(cropped_image) + if len(figure_rects) == 0: # text + page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True) + elif len(parse_table(cropped_image)) > 0: + #elif tables_in_image(cropped_image)[0]: # table + stats = parse_table(cropped_image) + cropped_image = draw_rectangles(cropped_image, stats, color=(255, 0, 0), annotate=True) + x,y,w,h = coordinates + page[y:y+h, x:x+w] = cropped_image + else: # figure + page = draw_rectangles(page, [coordinates], color=(0, 0, 255), annotate=True) + + # for coordinates, cropped_image in small_rects: + # figure_rects = detect_figures(cropped_image) + # if len(figure_rects) == 0 and len(list(find_primary_text_regions(cropped_image))) > 0: + # page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True) + # else: + # page = draw_rectangles(page, [coordinates], color=(0, 255, 255), annotate=True) + return page + + +def detect_figures_with_layout_parsing(pdf_path, page_index=1, show=False): + layout_rects, page = annotate_layout_in_pdf(pdf_path, page_index, return_rects=True) + big_structures, small_structures = cut_out_content_structures(layout_rects, page) + page = parse_content_structures(page, big_structures, small_structures) + vizlogger.debug(page, "figures03_final.png") + if show: + show_mpl(page) + else: + return page + + diff --git a/cv_analysis/figure_detection.py b/cv_analysis/figure_detection.py index 38f3f48..32db8d5 100644 --- a/cv_analysis/figure_detection.py +++ b/cv_analysis/figure_detection.py @@ -1,7 +1,11 @@ import cv2 import numpy as np from pdf2image import pdf2image - +import pandas as pd +from PIL import Image +import timeit +from os import path +from cv_analysis.locations import METADATA_TESTFILES, PNG_FOR_TESTING, PNG_FIGURES_DETECTED from cv_analysis.utils.detection import detect_large_coherent_structures from cv_analysis.utils.display import show_mpl from cv_analysis.utils.draw import draw_rectangles @@ -33,7 +37,7 @@ def detect_figures(image: np.array): def detect_figures_in_pdf(pdf_path, page_index=1, show=False): - page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0] + page = pdf2image.convert_from_path(pdf_path, dpi=300, first_page=page_index + 1, last_page=page_index + 1)[0] page = np.array(page) redaction_contours = detect_figures(page) @@ -41,3 +45,56 @@ def detect_figures_in_pdf(pdf_path, page_index=1, show=False): vizlogger.debug(page, "figures03_final.png") if show: show_mpl(page) + return page + + +def detect_figures_in_test_files(): + def save_as_pdf(pages): + p1, p = pages[0], pages[1:] + out_pdf_path = "/home/lillian/ocr_docs/output_files/fig_detection_pdf.pdf" + p1.save( + out_pdf_path, "PDF", resolution=150.0, save_all=True, append_images=p + ) + path = "/home/lillian/ocr_docs/" + ex_pages = pd.read_csv(path+"/metadata/metadata2.csv") + pages_detected = [] + + t0 = timeit.default_timer() + for name, page_nr in zip(ex_pages.pdf_name, ex_pages.page): + page = pdf2image.convert_from_path(path + "/original/" + name, dpi=300, first_page=page_nr, last_page=page_nr)[0] + page = np.array(page) + redaction_contours = detect_figures(page) + page = draw_rectangles(page, redaction_contours) + pages_detected.append(Image.fromarray(page)) + print(timeit.default_timer()-t0) + + save_as_pdf(pages_detected) + + +def detect_figures_in_png(pdf_path, show=False): + + page = Image.open(pdf_path) + page = np.array(page) + + redaction_contours = detect_figures(page) + page = draw_rectangles(page, redaction_contours) + vizlogger.debug(page, "figures03_final.png") + if show: + show_mpl(page) + return page + + +def detect_figures_in_test_files_png(): + file_name = pd.read_csv(METADATA_TESTFILES) + pages = [] + t0 = timeit.default_timer() + for name in file_name.image_name: + page = detect_figures_in_png(path.join(PNG_FOR_TESTING, name+".png")) + pages.append(Image.fromarray(page)) + t1 = timeit.default_timer() + print(t1-t0) + p1, p = pages[0], pages[1:] + out_pdf_path = path.join(PNG_FIGURES_DETECTED, "fig_detectes.pdf") + p1.save( + out_pdf_path, "PDF", resolution=300.0, save_all=True, append_images=p + ) \ No newline at end of file diff --git a/cv_analysis/layout_parsing.py b/cv_analysis/layout_parsing.py index 290a3d0..b6ae567 100644 --- a/cv_analysis/layout_parsing.py +++ b/cv_analysis/layout_parsing.py @@ -36,17 +36,17 @@ def parse_layout(image: np.array): if len(image_.shape) > 2: image_ = cv2.cvtColor(image_, cv2.COLOR_BGR2GRAY) - vizlogger.debug(image_, "layout01_start.png") + #vizlogger.debug(image_, "layout01_start.png") image_ = cv2.GaussianBlur(image_, (7, 7), 0) - vizlogger.debug(image_, "layout02_blur.png") + #vizlogger.debug(image_, "layout02_blur.png") thresh = cv2.threshold(image_, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1] vizlogger.debug(image_, "layout03_theshold.png") kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) - vizlogger.debug(kernel, "layout04_kernel.png") + #vizlogger.debug(kernel, "layout04_kernel.png") dilate = cv2.dilate(thresh, kernel, iterations=4) - vizlogger.debug(dilate, "layout05_dilate.png") + #vizlogger.debug(dilate, "layout05_dilate.png") rects = list(find_segments(dilate)) @@ -55,16 +55,16 @@ def parse_layout(image: np.array): x, y, w, h = rect cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 0), -1) cv2.rectangle(image, (x, y), (x + w, y + h), (255, 255, 255), 7) - vizlogger.debug(image, "layout06_rectangles.png") + #vizlogger.debug(image, "layout06_rectangles.png") _, image = cv2.threshold(image, 254, 255, cv2.THRESH_BINARY) - vizlogger.debug(image, "layout07_threshold.png") + #vizlogger.debug(image, "layout07_threshold.png") image = ~image - vizlogger.debug(image, "layout08_inverse.png") + #vizlogger.debug(image, "layout08_inverse.png") if len(image.shape) > 2: image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - vizlogger.debug(image, "layout09_convertcolor.png") + #vizlogger.debug(image, "layout09_convertcolor.png") rects = find_segments(image) # <- End of meta detection @@ -75,17 +75,22 @@ def parse_layout(image: np.array): return list(rects) -def annotate_layout_in_pdf(pdf_path, page_index=1, show=False): +def annotate_layout_in_pdf(pdf_path, page_index=1, return_rects=False, show=False): page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0] page = np.array(page) rects = parse_layout(page) - page = draw_rectangles(page, rects) - vizlogger.debug(page, "layout10_output.png") - if show: + if return_rects: + return rects, page + elif show: + page = draw_rectangles(page, rects) + vizlogger.debug(page, "layout10_output.png") show_mpl(page) + else: + page = draw_rectangles(page, rects) + return page """ diff --git a/cv_analysis/locations.py b/cv_analysis/locations.py index 6e56ec1..07e15f6 100644 --- a/cv_analysis/locations.py +++ b/cv_analysis/locations.py @@ -10,6 +10,13 @@ CONFIG_FILE = path.join(PACKAGE_ROOT_DIR, "config.yaml") LOG_FILE = "/tmp/log.log" DVC_DATA_DIR = path.join(PACKAGE_ROOT_DIR, "data") +PDF_FOR_TESTING = path.join(DVC_DATA_DIR, "pdfs_for_testing") +PNG_FOR_TESTING = path.join(DVC_DATA_DIR, "pngs_for_testing") +PNG_FIGURES_DETECTED = path.join(PNG_FOR_TESTING, "figures_detected") +PNG_TABLES_DETECTED = path.join(PNG_FOR_TESTING, "tables_detected_by_tp") +HASHED_PDFS_FOR_TESTING = path.join(PDF_FOR_TESTING, "hashed") +METADATA_TESTFILES = path.join(DVC_DATA_DIR, "metadata_testing_files.csv") + TEST_DIR = path.join(MODULE_DIR, "test") TEST_DATA_DIR = path.join(MODULE_DIR, "test", "test_data") diff --git a/cv_analysis/table_parsing.py b/cv_analysis/table_parsing.py index 5d4b522..2b6344b 100644 --- a/cv_analysis/table_parsing.py +++ b/cv_analysis/table_parsing.py @@ -1,25 +1,26 @@ + from functools import partial from itertools import chain, starmap from operator import attrgetter -from os.path import join + import cv2 import numpy as np + from pdf2image import pdf2image + from cv_analysis.utils.display import show_mpl from cv_analysis.utils.draw import draw_rectangles from cv_analysis.utils.post_processing import xywh_to_vecs, xywh_to_vec_rect, adjacent1d -from cv_analysis.utils.deskew import deskew_histbased +from cv_analysis.utils.deskew import deskew_histbased, deskew from cv_analysis.utils.filters import is_large_enough from cv_analysis.utils.visual_logging import vizlogger from cv_analysis.layout_parsing import parse_layout -def add_external_contours(image, contour_source_image): - contours, _ = cv2.findContours(contour_source_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) - contours = filter(partial(is_large_enough, min_area=5000), contours) - +def add_external_contours(image, img): + contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) for cnt in contours: x, y, w, h = cv2.boundingRect(cnt) cv2.rectangle(image, (x, y), (x + w, y + h), 255, 1) @@ -28,12 +29,12 @@ def add_external_contours(image, contour_source_image): def extend_lines(): - #TODO + # TODO pass def make_table_block_mask(): - #TODO + # TODO pass @@ -80,7 +81,7 @@ def isolate_vertical_and_horizontal_components(img_bin): img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v) img_lines_raw = img_bin_v | img_bin_h vizlogger.debug(img_lines_raw, "tables02_isolate02_img_bin_v.png") - + 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) @@ -100,7 +101,7 @@ def isolate_vertical_and_horizontal_components(img_bin): vizlogger.debug(img_bin_final, "tables10_isolate12_threshold.png") img_bin_final = cv2.dilate(img_bin_final, np.ones((1, 1), np.uint8), iterations=1) vizlogger.debug(img_bin_final, "tables11_isolate13_dilate.png") - + # add contours before lines are extended by blurring img_bin_final = add_external_contours(img_bin_final, img_lines_raw) vizlogger.debug(img_bin_final, "tables11_isolate14_contours_added.png") @@ -166,8 +167,8 @@ def parse_table(image: np.array, show=False): table_layout_boxes = find_table_layout_boxes(image) image = isolate_vertical_and_horizontal_components(image) - #image = add_external_contours(image, image) - #vizlogger.debug(image, "external_contours_added.png") + # image = add_external_contours(image, image) + # vizlogger.debug(image, "external_contours_added.png") _, _, stats, _ = cv2.connectedComponentsWithStats(~image, connectivity=8, ltype=cv2.CV_32S) @@ -200,3 +201,5 @@ def tables_in_image(cropped_image): return True, table_rects else: return False, None + + diff --git a/cv_analysis/test/scripts/export_example_pages.py b/cv_analysis/test/scripts/export_example_pages.py new file mode 100644 index 0000000..79dcd04 --- /dev/null +++ b/cv_analysis/test/scripts/export_example_pages.py @@ -0,0 +1,108 @@ +import hashlib +import os +from os import path +import pandas as pd +from pdf2image import convert_from_path +from itertools import chain +import json +from cv_analysis.locations import PDF_FOR_TESTING, TEST_DATA_DIR, PNG_FOR_TESTING, DVC_DATA_DIR, HASHED_PDFS_FOR_TESTING + + +def read_json(path): + with open(path, encoding='utf-8') as file: + data = json.load(file) + return data + + +def collect_metadata(example_pages, save=False): + metadata = [] + make_metadata_entry = make_metadata_entry_maker() + for name, document_sections in example_pages.items(): + metadata.append(f(name, document_sections, make_metadata_entry)) + metadata = list(chain.from_iterable(metadata)) + if save: + df = pd.DataFrame(data=metadata, columns=["image_name", "pdf_name", "page"]) + df.to_csv(path.join(DVC_DATA_DIR, "metadata_testing_files.csv")) + else: + return pd.DataFrame(data=metadata, columns=["image_name", "pdf_name", "page"]) + + +def f(name, document_sections, make_metadata_entry): + for pages in document_sections: + span = list(range(pages[0], pages[1] + 1)) + for page_nr in span: + yield make_metadata_entry(name, page_nr) + + +def make_metadata_entry_maker(): + i = -1 + + def make_metadata_entry(name, page_nr): + nonlocal i + i += 1 + return [f"fig_table{i:0>3}", name, page_nr] + + return make_metadata_entry + + +def split_pdf(example_pages): + dir_path = PDF_FOR_TESTING + i = 0 + for name, document_sections in example_pages.items(): + for pages in document_sections: + images = convert_from_path(pdf_path=path.join(dir_path, name), dpi=300, first_page=pages[0], + last_page=pages[1]) + for image in images: + fp = path.join(PNG_FOR_TESTING, f"fig_table{i:0>3}.png") + image.save(fp=fp, dpi=(300, 300)) + i += 1 + + +def find_hash(file_path): + BLOCK_SIZE = 65536 + + file_hash = hashlib.sha256() + with open(file_path, 'rb') as f: + fb = f.read(BLOCK_SIZE) + while len(fb) > 0: + file_hash.update(fb) + fb = f.read(BLOCK_SIZE) + + return file_hash.hexdigest() + + +def rename_files_with_hash(example_pages): + files_to_rename = list(example_pages.keys()) + folder = HASHED_PDFS_FOR_TESTING + + # Iterate through the folder + for file in os.listdir(folder): + # Checking if the file is present in the list + if file in files_to_rename: + # construct current name using file name and path + old_name = path.join(folder, file) + # get file name without extension + only_name = path.splitext(file)[0] + + # Adding the new name with extension + hash = find_hash(old_name) + # construct full file path + new_name = path.join(folder, hash + ".pdf") + + # Renaming the file + os.rename(old_name, new_name) + + # verify the result + res = os.listdir(folder) + print(res) + + +def main(): + examples_pages = read_json(path.join(TEST_DATA_DIR, "example_pages.json")) + rename_files_with_hash(examples_pages) + #collect_metadata(examples_pages, save=True) + #split_pdf(examples_pages) + + +if __name__ == "__main__": + main() diff --git a/cv_analysis/utils/display.py b/cv_analysis/utils/display.py index 4d346a5..999c9a2 100644 --- a/cv_analysis/utils/display.py +++ b/cv_analysis/utils/display.py @@ -18,6 +18,7 @@ def save_mpl(image, path): ax.imshow(image, cmap="gray") # plt.close() plt.savefig(path) + plt.close() def show_cv2(image): diff --git a/cv_analysis/utils/post_processing.py b/cv_analysis/utils/post_processing.py index 753e091..1749f2d 100644 --- a/cv_analysis/utils/post_processing.py +++ b/cv_analysis/utils/post_processing.py @@ -25,7 +25,7 @@ def remove_included(rectangles): return b.xmin + tol >= a.xmin and b.ymin + tol >= a.ymin and b.xmax - tol <= a.xmax and b.ymax - tol <= a.ymax def is_not_included(rect, rectangles): - return not any(included(r2, rect) for r2 in rectangles if not rect == r2) + return not any(includes(r2, rect) for r2 in rectangles if not rect == r2) rectangles = list(map(xywh_to_vec_rect, rectangles)) rectangles = filter(partial(is_not_included, rectangles=rectangles), rectangles) diff --git a/cv_analysis/utils/text.py b/cv_analysis/utils/text.py index acfaa48..5d9ccaa 100644 --- a/cv_analysis/utils/text.py +++ b/cv_analysis/utils/text.py @@ -1,5 +1,7 @@ import cv2 +import numpy as np +from cv_analysis.utils.display import show_mpl def remove_primary_text_regions(image): """Removes regions of primary text, meaning no figure descriptions for example, but main text body paragraphs. @@ -14,11 +16,10 @@ def remove_primary_text_regions(image): image = image.copy() cnts = find_primary_text_regions(image) - for cnt in cnts: x, y, w, h = cv2.boundingRect(cnt) cv2.rectangle(image, (x, y), (x + w, y + h), (255, 255, 255), -1) - + #show_mpl(image) return image @@ -36,7 +37,8 @@ def find_primary_text_regions(image): """ def is_likely_primary_text_segments(cnt): - return 800 < cv2.contourArea(cnt) < 15000 + x,y,w,h = cv2.boundingRect(cnt) + return 800 < cv2.contourArea(cnt) < 16000 or w/h > 3 image = image.copy() @@ -45,14 +47,17 @@ def find_primary_text_regions(image): image = cv2.threshold(image, 253, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1] - close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 3)) + close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (17, 7)) #20,3 close = cv2.morphologyEx(image, cv2.MORPH_CLOSE, close_kernel, iterations=1) - dilate_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 3)) + #show_mpl(close) + + dilate_kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(7, 3)) #5,3 dilate = cv2.dilate(close, dilate_kernel, iterations=1) - cnts, _ = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + #show_mpl(dilate) + cnts, _ = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) cnts = filter(is_likely_primary_text_segments, cnts) return cnts diff --git a/data/.gitignore b/data/.gitignore index 09d8485..7b38b1e 100644 --- a/data/.gitignore +++ b/data/.gitignore @@ -1 +1,7 @@ /test_pdf.pdf +/pdfs_for_testing +/figure_detection.png +/layout_parsing.png +/redaction_detection.png +/table_parsing.png +/pngs_for_testing diff --git a/data/pdfs_for_testing.dvc b/data/pdfs_for_testing.dvc new file mode 100644 index 0000000..e85e518 --- /dev/null +++ b/data/pdfs_for_testing.dvc @@ -0,0 +1,5 @@ +outs: +- md5: bb0ce084f7ca54583972da71cb87e22c.dir + size: 367181628 + nfiles: 28 + path: pdfs_for_testing diff --git a/data/pngs_for_testing.dvc b/data/pngs_for_testing.dvc new file mode 100644 index 0000000..630eab7 --- /dev/null +++ b/data/pngs_for_testing.dvc @@ -0,0 +1,5 @@ +outs: +- md5: 4fed91116111b47edf1c6f6a67eb84d3.dir + size: 58125058 + nfiles: 230 + path: pngs_for_testing diff --git a/scripts/annotate.py b/scripts/annotate.py index c92ecf1..002e245 100644 --- a/scripts/annotate.py +++ b/scripts/annotate.py @@ -10,7 +10,7 @@ def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("pdf_path") parser.add_argument("page_index", type=int) - parser.add_argument("--type", choices=["table", "redaction", "layout", "figure"]) + parser.add_argument("--type", choices=["table", "redaction", "layout", "figure", "figures"]) parser.add_argument("--show", action="store_true", default=False) args = parser.parse_args() @@ -20,7 +20,6 @@ def parse_args(): if __name__ == "__main__": args = parse_args() - # print(args.show) if args.type == "table": annotate_tables_in_pdf(args.pdf_path, page_index=args.page_index, show=args.show) elif args.type == "redaction": @@ -28,4 +27,4 @@ if __name__ == "__main__": elif args.type == "layout": annotate_layout_in_pdf(args.pdf_path, page_index=args.page_index, show=args.show) elif args.type == "figure": - detect_figures_in_pdf(args.pdf_path, page_index=args.page_index, show=args.show) + detect_figures_in_pdf(args.pdf_path, page_index=args.page_index, show=True)