diff --git a/cv_analysis/figure_detection/text.py b/cv_analysis/figure_detection/text.py index 11fd832..a179e27 100644 --- a/cv_analysis/figure_detection/text.py +++ b/cv_analysis/figure_detection/text.py @@ -1,5 +1,7 @@ import cv2 +from cv_analysis.layout_parsing import normalize_to_gray_scale + def remove_primary_text_regions(image): """Removes regions of primary text, meaning no figure descriptions for example, but main text body paragraphs. @@ -35,6 +37,7 @@ def remove_primary_text_regions(image): def apply_threshold_to_image(image): """Converts an image to black and white.""" + image = normalize_to_gray_scale(image) image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) > 2 else image return cv2.threshold(image, 253, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1] diff --git a/cv_analysis/layout_parsing.py b/cv_analysis/layout_parsing.py index cd332a8..d8e6ac6 100644 --- a/cv_analysis/layout_parsing.py +++ b/cv_analysis/layout_parsing.py @@ -1,28 +1,29 @@ from functools import reduce, partial -from typing import Iterable +from typing import Iterable, List import cv2 import numpy as np from funcy import compose, rcompose, lkeep from cv_analysis.utils.common import find_contours -from cv_analysis.utils.conversion import box_to_rectangle +from cv_analysis.utils.conversion import box_to_rectangle, contour_to_rectangle from cv_analysis.utils.merging import connect_related_rectangles from cv_analysis.utils.postprocessing import remove_included, has_no_parent from cv_analysis.utils.rectangle import Rectangle -def parse_layout(image: np.array): - - rectangles = find_segments(image) - rectangles = remove_included(rectangles) - rectangles = connect_related_rectangles(rectangles) - rectangles = remove_included(rectangles) +def parse_layout(image: np.array) -> List[Rectangle]: + rectangles = rcompose( + find_segments, + remove_included, + connect_related_rectangles, + remove_included, + )(image) return rectangles -def find_segments(image): +def find_segments(image: np.ndarray) -> List[Rectangle]: rectangles = rcompose( prepare_for_initial_detection, __find_segments, @@ -33,29 +34,25 @@ def find_segments(image): return rectangles -def prepare_for_initial_detection(image: np.ndarray): +def prepare_for_initial_detection(image: np.ndarray) -> np.ndarray: return compose(dilate_page_components, normalize_to_gray_scale)(image) -def __find_segments(image): +def __find_segments(image: np.ndarray) -> List[Rectangle]: def to_rectangle_if_valid(contour, hierarchy): - return ( - box_to_rectangle(cv2.boundingRect(contour)) - if is_likely_segment(contour) and has_no_parent(hierarchy) - else None - ) + return contour_to_rectangle(contour) if is_likely_segment(contour) and has_no_parent(hierarchy) else None rectangles = lkeep(map(to_rectangle_if_valid, *find_contours(image))) return rectangles -def is_likely_segment(rect, min_area=100): +def is_likely_segment(rectangle: Rectangle, min_area: float = 100) -> bool: # FIXME: Parameterize via factory - return cv2.contourArea(rect, False) > min_area + return cv2.contourArea(rectangle, False) > min_area -def dilate_page_components(image): +def dilate_page_components(image: np.ndarray) -> np.ndarray: # FIXME: Parameterize via factory image = cv2.GaussianBlur(image, (7, 7), 0) # FIXME: Parameterize via factory @@ -63,10 +60,11 @@ def dilate_page_components(image): # FIXME: Parameterize via factory kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) # FIXME: Parameterize via factory - return cv2.dilate(thresh, kernel, iterations=4) + dilate = cv2.dilate(thresh, kernel, iterations=4) + return dilate -def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]): +def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]) -> np.ndarray: image = fill_rectangles(image, rectangles) image = threshold_image(image) @@ -76,28 +74,27 @@ def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle return image -def normalize_to_gray_scale(image): - if len(image.shape) > 2: - image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) +def normalize_to_gray_scale(image: np.ndarray) -> np.ndarray: + image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) > 2 else image return image -def threshold_image(image): +def threshold_image(image: np.ndarray) -> np.ndarray: # FIXME: Parameterize via factory _, image = cv2.threshold(image, 254, 255, cv2.THRESH_BINARY) return image -def invert_image(image): +def invert_image(image: np.ndarray): return ~image -def fill_rectangles(image, rectangles): +def fill_rectangles(image: np.ndarray, rectangles: Iterable[Rectangle]) -> np.ndarray: image = reduce(fill_in_component_area, rectangles, image) return image -def fill_in_component_area(image, rect): +def fill_in_component_area(image: np.ndarray, rect: Rectangle) -> np.ndarray: 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) diff --git a/cv_analysis/utils/conversion.py b/cv_analysis/utils/conversion.py index 3b0e100..8a25abc 100644 --- a/cv_analysis/utils/conversion.py +++ b/cv_analysis/utils/conversion.py @@ -1,4 +1,5 @@ import json +from typing import Sequence import cv2 @@ -9,12 +10,12 @@ def contour_to_rectangle(contour): return box_to_rectangle(cv2.boundingRect(contour)) -def box_to_rectangle(box): +def box_to_rectangle(box: Sequence[int]) -> Rectangle: x, y, w, h = box return Rectangle(x, y, x + w, y + h) -def rectangle_to_box(rectangle): +def rectangle_to_box(rectangle: Rectangle) -> Sequence[int]: return [rectangle.x1, rectangle.y1, rectangle.width, rectangle.height] diff --git a/cv_analysis/utils/postprocessing.py b/cv_analysis/utils/postprocessing.py index 1ce776f..4069fd4 100644 --- a/cv_analysis/utils/postprocessing.py +++ b/cv_analysis/utils/postprocessing.py @@ -1,6 +1,6 @@ from functools import partial from itertools import starmap, compress -from typing import Iterable, List +from typing import Iterable, List, Sequence from cv_analysis.utils.rectangle import Rectangle @@ -17,8 +17,8 @@ def remove_overlapping(rectangles: Iterable[Rectangle]) -> List[Rectangle]: def remove_included(rectangles: Iterable[Rectangle]) -> List[Rectangle]: - keep = [rect for rect in rectangles if not rect.is_included(rectangles)] - return keep + rectangles_to_keep = [rect for rect in rectangles if not rect.is_included(rectangles)] + return rectangles_to_keep def __remove_isolated_unsorted(rectangles: Iterable[Rectangle]) -> List[Rectangle]: @@ -45,9 +45,9 @@ def __remove_isolated_sorted(rectangles: Iterable[Rectangle]) -> List[Rectangle] return rectangles -def remove_isolated(rectangles: Iterable[Rectangle], input_unsorted=True) -> List[Rectangle]: +def remove_isolated(rectangles: Iterable[Rectangle], input_unsorted: bool = True) -> List[Rectangle]: return (__remove_isolated_unsorted if input_unsorted else __remove_isolated_sorted)(rectangles) -def has_no_parent(hierarchy): +def has_no_parent(hierarchy: Sequence[int]) -> bool: return hierarchy[-1] <= 0