diff --git a/cv_analysis/layout_parsing.py b/cv_analysis/layout_parsing.py index f5de783..46ae963 100644 --- a/cv_analysis/layout_parsing.py +++ b/cv_analysis/layout_parsing.py @@ -7,7 +7,10 @@ import cv2 import numpy as np -from cv_analysis.utils.connect_rects import connect_related_rects2 +from cv_analysis.utils.connect_rects import connect_related_rects +from cv_analysis.utils.display import show_image_mpl +from cv_analysis.utils.draw import draw_rectangles +from cv_analysis.utils.label_rects import label_rects from cv_analysis.utils.structures import Rectangle from cv_analysis.utils.postprocessing import ( remove_overlapping, @@ -53,7 +56,7 @@ def fill_in_component_area(image, rect): def parse_layout(image: np.array): image = image.copy() image_ = image.copy() - + #show_image_mpl(image) if len(image_.shape) > 2: image_ = cv2.cvtColor(image_, cv2.COLOR_BGR2GRAY) @@ -80,8 +83,7 @@ def parse_layout(image: np.array): rects = remove_included(rects) rects = map(lambda r: r.xywh(), rects) - rects = connect_related_rects2(rects) - rects = list(map(Rectangle.from_xywh, rects)) - rects = remove_included(rects) - + rects = connect_related_rects(rects) + # rects = list(map(Rectangle.from_xywh, rects)) + # rects = remove_included(rects) return rects diff --git a/cv_analysis/utils/connect_rects.py b/cv_analysis/utils/connect_rects.py index 09d48bb..6304e33 100644 --- a/cv_analysis/utils/connect_rects.py +++ b/cv_analysis/utils/connect_rects.py @@ -75,28 +75,8 @@ def find_related_rects(rects): return rect_pairs, unrel_rects -def connect_related_rects(rects): - rects_to_connect, rects_new = find_related_rects(rects) - while len(rects_to_connect) > 0: - rects_fused = list(starmap(fuse_rects, rects_to_connect)) - rects_fused = list(dict.fromkeys(rects_fused)) - - if len(rects_fused) == 1: - rects_new += rects_fused - rects_fused = [] - - rects_to_connect, connected_rects = find_related_rects(rects_fused) - rects_new += connected_rects - - if len(rects_to_connect) > 1 and len(set(rects_to_connect)) == 1: - rects_new.append(rects_fused[0]) - rects_to_connect = [] - - return rects_new - - -def connect_related_rects2(rects: Iterable[tuple]): +def connect_related_rects(rects: Iterable[tuple]): rects = list(rects) current_idx = 0 diff --git a/cv_analysis/utils/label_rects.py b/cv_analysis/utils/label_rects.py new file mode 100644 index 0000000..1b121dc --- /dev/null +++ b/cv_analysis/utils/label_rects.py @@ -0,0 +1,115 @@ +from itertools import starmap +from typing import Iterable + +import cv2 +import numpy as np + +from cv_analysis.figure_detection.text import remove_primary_text_regions, apply_threshold_to_image +from cv_analysis.table_parsing import preprocess, isolate_vertical_and_horizontal_components, \ + turn_connected_components_into_rects +from cv_analysis.utils.display import show_image_mpl + + +def area_is_bigger_than(rect: tuple, maxarea=100000): + x, y, w, h = rect + return w * h >= maxarea + + +def define_rect(rect_img, original_position): + # print(original_position) + # show_image_mpl(rect_img) + xo, yo, wo, ho = original_position + rect_img_inv = preprocess(rect_img) + # print("pixel density inverted img", pixel_density(rect_img_inv)) + grid_inv = isolate_vertical_and_horizontal_components(rect_img_inv) + cnts, _ = cv2.findContours(image=grid_inv, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) + if cnts: + rects = turn_connected_components_into_rects(grid_inv) + rects = map(lambda r: r.xywh(), rects) + bbox = list((cv2.boundingRect(c) for c in cnts)) + if len(list(rects)) > 1 and len(bbox) == 1: + x, y, w, h = bbox[0] + w_img, h_img = rect_img.shape + if w * h / (w_img * h_img) >= 0.75: + # print("is table") + return (255, 255, 0) + else: + # show_image_mpl(rect_img) + # print(" table detected but to small for layout rect, so cant be table, maybe figure?") + return (0, 255, 255) + + else: + if is_header(yo + ho): + # print("is header component") + return (255, 0, 0) + elif is_footer(yo): + # print("is footer component") + return (0, 255, 0) + else: + # print("img-inv",pixel_density(rect_img_inv)) + # show_image_mpl(rect_img) + # print("grid_in", pixel_density(grid_inv)) + # show_image_mpl(grid_inv) + # print("single cell or no connected components, maybe figure?") + return (0, 255, 255) + + else: + if is_header(yo + ho): + # print("is header text") + return (255, 0, 0) + elif is_footer(yo): + # print("is footer text") + return (0, 255, 0) + else: + # print("is text") + return (0, 0, 255) + + +def is_header(y): + return y < 200 + + +def is_footer(y): + return y > 2100 + + +def is_text(img): + show_image_mpl(img) + cleaned = remove_primary_text_regions(img) + show_image_mpl(cleaned) + return pixel_density(cleaned) < 0.05 + + +def pixel_density(img): + pixels = np.count_nonzero(img) + density = pixels / img.size + return density + + +def annotate_rect(image, rects, rect_labels): + 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, + ) + + for rect, label in zip(rects, rect_labels): + x, y, w, h = rect + cv2.rectangle(image, (x, y), (x + w, y + h), label, 2) + + return image + + +def label_rects(image: np.array, rects: Iterable[tuple]): + def crop_image_rects(rect): + x, y, w, h = rect + return image[y:y + h, x:x + w] + + rect_images = map(crop_image_rects, rects) + rect_labels = starmap(define_rect, zip(rect_images, rects)) + return rect_labels diff --git a/cv_analysis/utils/sort_rects.py b/cv_analysis/utils/sort_rects.py new file mode 100644 index 0000000..a665d79 --- /dev/null +++ b/cv_analysis/utils/sort_rects.py @@ -0,0 +1,89 @@ +from typing import Iterable + +import cv2 +import numpy as np + +from cv_analysis.figure_detection.text import remove_primary_text_regions, apply_threshold_to_image +from cv_analysis.table_parsing import preprocess, isolate_vertical_and_horizontal_components, \ + turn_connected_components_into_rects +from cv_analysis.utils.display import show_image_mpl + + +def area_is_bigger_than(rect: tuple, maxarea=100000): + x, y, w, h = rect + return w * h >= maxarea + + +def define_rect(rect_img, original_position): + show_image_mpl(rect_img) + xo,yo,wo,ho = original_position + if is_header(yo+ho): + print(original_position, " is header") + return "header" + elif is_footer(yo): + print(original_position, " is footer") + return "footer" + elif is_table(rect_img): + print(original_position, " is table") + return "table" + elif is_text(rect_img): + print(original_position, " is text") + return "text" + + +def is_table(rect_img): + rect_img_inv = preprocess(rect_img) + grid_inv = isolate_vertical_and_horizontal_components(rect_img_inv) + cnts, _ = cv2.findContours(image=grid_inv, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) + if cnts: + rects = turn_connected_components_into_rects(grid_inv) + rects = map(lambda r: r.xywh(), rects) + bbox = list((cv2.boundingRect(c) for c in cnts)) + if len(list(rects)) > 1 and len(bbox) == 1: + x, y, w, h = bbox[0] + w_img, h_img = rect_img.shape + if w * h / (w_img * h_img) >= 0.75: + print("is table") + return True + else: + print(" table detected but to small for layout rect, so cant be table, maybe figure?") + return False + + else: + print("single cell or no connected components, maybe figure?") + return False + + else: + print("not a table, but text?") + return False + +def is_header(y): + return y < 200 + + +def is_footer(y): + return y > 2150 + + +def is_text(img): + show_image_mpl(img) + cleaned = remove_primary_text_regions(img) + show_image_mpl(cleaned) + return pixel_density(cleaned) < 0.05 + + +def pixel_density(img): + pixels = np.count_nonzero(img) + density = pixels / img.size + return density + + +def annotate_rect(rect, rect_img): + pass + + +def label_rects(rects: Iterable[tuple], image: np.array): + labeled_rects = {} + for rect in rects: + x, y, w, h = rect + labeled_rects[rect] = define_rect(image[y:y + h, x:x + w], rect)