70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
from functools import partial
|
|
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,
|
|
dilate_page_components,
|
|
normalize_to_gray_scale,
|
|
threshold_image,
|
|
invert_image,
|
|
fill_rectangles,
|
|
)
|
|
from cv_analysis.utils.conversion import 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) -> List[Rectangle]:
|
|
rectangles = rcompose(
|
|
find_segments,
|
|
remove_included,
|
|
connect_related_rectangles,
|
|
remove_included,
|
|
)(image)
|
|
|
|
return rectangles
|
|
|
|
|
|
def find_segments(image: np.ndarray) -> List[Rectangle]:
|
|
rectangles = rcompose(
|
|
prepare_for_initial_detection,
|
|
__find_segments,
|
|
partial(prepare_for_meta_detection, image.copy()),
|
|
__find_segments,
|
|
)(image)
|
|
|
|
return rectangles
|
|
|
|
|
|
def prepare_for_initial_detection(image: np.ndarray) -> np.ndarray:
|
|
return compose(dilate_page_components, normalize_to_gray_scale)(image)
|
|
|
|
|
|
def __find_segments(image: np.ndarray) -> List[Rectangle]:
|
|
def to_rectangle_if_valid(contour, hierarchy):
|
|
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 prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]) -> np.ndarray:
|
|
|
|
image = fill_rectangles(image, rectangles)
|
|
image = threshold_image(image)
|
|
image = invert_image(image)
|
|
image = normalize_to_gray_scale(image)
|
|
|
|
return image
|
|
|
|
|
|
def is_likely_segment(rectangle: Rectangle, min_area: float = 100) -> bool:
|
|
# FIXME: Parameterize via factory
|
|
return cv2.contourArea(rectangle, False) > min_area
|