diff --git a/Makefile b/Makefile index 25e0921..b916e8b 100644 --- a/Makefile +++ b/Makefile @@ -83,7 +83,7 @@ smoke-test: tests/resources/2files/kant_aufklaerung_1784_0020.tif eynollah -m $(CURDIR) layout -di $( None: t0 = time.time() - - self.logger.info(f"Processing file: {self.writer.image_filename}") + self.logger.info(img_filename) + + image = self.cache_images(image_filename=img_filename) + writer = EynollahXmlWriter( + dir_out=dir_out, + image_filename=img_filename, + image_width=image['img'].shape[1], + image_height=image['img'].shape[0], + ) + + if os.path.exists(writer.output_filename): + if overwrite: + self.logger.warning("will overwrite existing output file '%s'", writer.output_filename) + else: + self.logger.warning("will skip input for existing output file '%s'", writer.output_filename) + return + + self.logger.info(f"Processing file: {writer.image_filename}") self.logger.info("Step 1/5: Image Enhancement") + + num_col_classifier, _ = self.run_enhancement(image) + writer.scale_x = image['scale_x'] + writer.scale_y = image['scale_y'] - img_res, is_image_enhanced, num_col_classifier, _ = \ - self.run_enhancement() - - self.logger.info(f"Image: {self.image.shape[1]}x{self.image.shape[0]}, " - f"{self.dpi} DPI, {num_col_classifier} columns") - if is_image_enhanced: - self.logger.info("Enhancement applied") - + self.logger.info(f"Image: {image['img_res'].shape[1]}x{image['img_res'].shape[0]}, " + f"scale {image['scale_x']:.1f}x{image['scale_y']:.1f}, " + f"{image['dpi']} DPI, {num_col_classifier} columns") self.logger.info(f"Enhancement complete ({time.time() - t0:.1f}s)") - # Image Extraction Mode self.logger.info("Step 2/5: Image Extraction Mode") + t1 = time.time() + page_cont, image_page, _ = self.extract_page(image) + page = Region(page_cont) - _, _, _, polygons_of_images, \ - image_page, page_coord, cont_page = \ - self.get_regions_light_v_extract_only_images(img_res, num_col_classifier) + _, _, images_cont = self.get_early_layout( + image['img_res'], num_col_classifier) + self.logger.debug("Found %d images", len(images_cont)) - pcgts = self.writer.build_pagexml_no_full_layout( - found_polygons_text_region=[], - page_coord=page_coord, - order_of_texts=[], - all_found_textline_polygons=[], - all_box_coord=[], - found_polygons_text_region_img=polygons_of_images, - found_polygons_marginals_left=[], - found_polygons_marginals_right=[], - all_found_textline_polygons_marginals_left=[], - all_found_textline_polygons_marginals_right=[], - all_box_coord_marginals_left=[], - all_box_coord_marginals_right=[], - slopes=[], - slopes_marginals_left=[], - slopes_marginals_right=[], - cont_page=cont_page, - polygons_seplines=[], - found_polygons_tables=[], - ) + # FIXME: post-hoc cropping (remove when models support it, and replace image['img_res'] with image_page) + page_box = cv2.boundingRect(page.contour) + images_cont = [cont - [page_box[:2]] + for cont in images_cont] if self.plotter: - self.plotter.write_images_into_directory(polygons_of_images, image_page) - + self.plotter.write_images_into_directory(images_cont, image_page, + name=image['name']) self.logger.info("Image extraction complete") - return pcgts + + images = [Region(cont) for cont in images_cont] + # can be empty if above page frame + images = [image for image in images if image.area] + pcgts = writer.build_pagexml( + page=page, + img_bin=self.imread(image, binary=True) if self.input_binary else None, + num_col=num_col_classifier, + images=images, + ) + writer.write_pagexml(pcgts) + self.logger.info("Job done in %.1fs", time.time() - t0) diff --git a/src/eynollah/eynollah.py b/src/eynollah/eynollah.py index becf802..eea3b6e 100644 --- a/src/eynollah/eynollah.py +++ b/src/eynollah/eynollah.py @@ -24,7 +24,8 @@ from difflib import SequenceMatcher as sq import math import os import time -from typing import Optional +from typing import Optional, List, Tuple +from itertools import compress from functools import partial from pathlib import Path import multiprocessing as mp @@ -39,7 +40,7 @@ from scipy.ndimage import gaussian_filter1d try: import matplotlib.pyplot as plt except ImportError: - plt = None + plt = None # type: ignore from .model_zoo import EynollahModelZoo from .utils.contour import ( @@ -53,14 +54,8 @@ from .utils.contour import ( return_contours_of_image, return_contours_of_interested_region, return_parent_contours, - dilate_textregion_contours, - dilate_textline_contours, match_deskewed_contours, estimate_skew_contours, - polygon2contour, - contour2polygon, - join_polygons, - make_intersection, ) from .utils.rotate import rotate_image from .utils.separate_lines import ( @@ -71,12 +66,15 @@ from .utils.marginals import get_marginals from .utils.resize import resize_image from .utils.shm import share_ndarray from .utils import ( + Region, + TextRegion, ensure_array, pairwise, + itemgetter, is_image_filename, isNaN, crop_image_inside_box, - box2rect, + box2slice, find_num_col, otsu_copy_binary, seg_mask_label, @@ -99,8 +97,8 @@ MAX_SLOPE = 999 KERNEL = np.ones((5, 5), np.uint8) -_instance = None -def _set_instance(instance): +_instance: Optional["Eynollah"] = None +def _set_instance(instance: "Eynollah"): global _instance _instance = instance def _run_single(*args, **kwargs): @@ -299,7 +297,7 @@ class Eynollah: img_new = np.copy(img) img_is_resized = False #elif conf_col < 0.8 and img_h_new >= 8000: - elif img_h_new >= 8000: + elif conf_col < 0.9 and img_h_new >= 8000: # don't upsample if too large img_new = np.copy(img) img_is_resized = False @@ -493,7 +491,7 @@ class Eynollah: height_mid = img_height_model - 2 * margin img_h = img.shape[0] img_w = img.shape[1] - prediction = None + prediction: np.ndarray = None # type: ignore nxf = math.ceil((img_w - 2.0 * margin) / width_mid) nyf = math.ceil((img_h - 2.0 * margin) / height_mid) @@ -808,14 +806,13 @@ class Eynollah: return segmentation, confidence def extract_page(self, image): - cropped_page = img = image['img_res'] + page_cropped = img = image['img_res'] h, w = img.shape[:2] - page_coord = [0, h, 0, w] - cont_page = [np.array([[[0, 0]], - [[w, 0]], - [[w, h]], - [[0, h]]])] - mask_page = np.ones((h, w), dtype=np.uint8) + page_cont = np.array([[[0, 0]], + [[w, 0]], + [[w, h]], + [[0, h]]]) + page_mask = np.ones((h, w), dtype=np.uint8) if not self.ignore_page_extraction: self.logger.debug("enter extract_page") #cv2.GaussianBlur(img, (5, 5), 0) @@ -823,26 +820,13 @@ class Eynollah: contours, _ = cv2.findContours(prediction, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if len(contours): areas = np.array(list(map(cv2.contourArea, contours))) - cnt = contours[np.argmax(areas)] - cont_page = [cnt] - x, y, w, h = cv2.boundingRect(cnt) - #if x <= 30: - #w += x - #x = 0 - #if (self.image.shape[1] - (x + w)) <= 30: - #w = w + (self.image.shape[1] - (x + w)) - #if y <= 30: - #h = h + y - #y = 0 - #if (self.image.shape[0] - (y + h)) <= 30: - #h = h + (self.image.shape[0] - (y + h)) - box = [x, y, w, h] - cropped_page, page_coord = crop_image_inside_box(box, img) - mask_page = np.zeros((h, w), dtype=np.uint8) - mask_page = cv2.fillPoly(mask_page, pts=[cnt - [x, y]], color=1) - + page_cont = contours[np.argmax(areas)] + box = (x, y, w, h) = cv2.boundingRect(page_cont) + page_cropped = img[box2slice(box)] + page_mask = np.zeros((h, w), dtype=np.uint8) + page_mask = cv2.fillPoly(page_mask, pts=[page_cont - [x, y]], color=1) self.logger.debug("exit extract_page") - return page_coord, cont_page, cropped_page, mask_page + return page_cont, page_cropped, page_mask def early_page_for_num_of_column_classification(self, img): if not self.ignore_page_extraction: @@ -888,20 +872,27 @@ class Eynollah: self.logger.debug("exit extract_text_regions") return prediction_regions - def get_textlines_of_a_textregion_sorted(self, textlines_textregion, cx_textline, cy_textline, w_h_textline): - N = len(cy_textline) + def get_textlines_of_a_textregion_sorted( + self, + textlines_cont: List[np.ndarray], + textlines_conf: List[float], + textlines_cx: List[float], + textlines_cy: List[float], + textlines_w_h: List[Tuple[int, int]], + ): + N = len(textlines_cont) if N <= 1: - return textlines_textregion + return textlines_cont, textlines_conf - cx_textline = np.array(cx_textline) - cy_textline = np.array(cy_textline) - diff_cy = np.abs(np.diff(np.sort(cy_textline))) - diff_cx = np.abs(np.diff(np.sort(cx_textline))) + textlines_cx = np.array(textlines_cx) + textlines_cy = np.array(textlines_cy) + diff_cx = np.abs(np.diff(np.sort(textlines_cx))) + diff_cy = np.abs(np.diff(np.sort(textlines_cy))) if N > 1: mean_y_diff = np.median(diff_cy) mean_x_diff = np.median(diff_cx) - count_hor = np.count_nonzero(np.diff(w_h_textline, axis=0) > 0) + count_hor = np.count_nonzero(np.diff(textlines_w_h, axis=0) > 0) count_ver = N - count_hor else: mean_y_diff = 0 @@ -909,95 +900,111 @@ class Eynollah: count_hor = 1 count_ver = 0 + sorted_textlines_cont = [] + sorted_textlines_conf = [] if count_hor >= count_ver: row_threshold = mean_y_diff / 1.5 if mean_y_diff > 0 else 10 rows = [] - for prev_idx, curr_idx in pairwise(np.argsort(cy_textline)): + for prev_idx, curr_idx in pairwise(np.argsort(textlines_cy)): if not len(rows): rows.append([prev_idx]) - if abs(cy_textline[curr_idx] - cy_textline[prev_idx]) <= row_threshold: + if abs(textlines_cy[curr_idx] - textlines_cy[prev_idx]) <= row_threshold: rows[-1].append(curr_idx) else: rows.append([curr_idx]) - sorted_textlines = [] for row in rows: - for idx in np.argsort(cx_textline[row]): - sorted_textlines.append(textlines_textregion[row[idx]]) + for idx in np.argsort(textlines_cx[row]): + sorted_textlines_cont.append(textlines_cont[row[idx]]) + sorted_textlines_conf.append(textlines_conf[row[idx]]) else: col_threshold = mean_x_diff / 1.5 if mean_x_diff > 0 else 10 cols = [] - for prev_idx, curr_idx in pairwise(np.argsort(cx_textline)): + for prev_idx, curr_idx in pairwise(np.argsort(textlines_cx)): if not len(cols): cols.append([prev_idx]) - if abs(cx_textline[curr_idx] - cx_textline[prev_idx]) <= col_threshold: + if abs(textlines_cx[curr_idx] - textlines_cx[prev_idx]) <= col_threshold: cols[-1].append(curr_idx) else: cols.append([curr_idx]) - sorted_textlines = [] for col in cols: - for idx in np.argsort(cy_textline[col]): - sorted_textlines.append(textlines_textregion[col[idx]]) + for idx in np.argsort(textlines_cy[col]): + sorted_textlines_cont.append(textlines_cont[col[idx]]) + sorted_textlines_conf.append(textlines_conf[col[idx]]) - return sorted_textlines + return sorted_textlines_cont, sorted_textlines_conf - def get_slopes_and_deskew_new_light2(self, contours_par, textline_mask_tot, slope_deskew): + def get_slopes_and_deskew_new_light2( + self, parents: List[TextRegion], + textline_mask_tot: np.ndarray, + textline_confidence: np.ndarray, + slope_deskew: float + ): + textlines_cont = return_contours_of_interested_region(textline_mask_tot, 1, 0.0001) + textlines_conf = get_region_confidences(textlines_cont, textline_confidence) + textlines_cx, textlines_cy = find_center_of_contours(textlines_cont) + textlines_w_h = [cv2.boundingRect(polygon)[2:] for polygon in textlines_cont] + textlines_args = np.arange(len(textlines_cont)) - polygons_of_textlines = return_contours_of_interested_region(textline_mask_tot, 1, 0.0001) - cx_textlines, cy_textlines = find_center_of_contours(polygons_of_textlines) - w_h_textlines = [cv2.boundingRect(polygon)[2:] for polygon in polygons_of_textlines] - args_textlines = np.arange(len(polygons_of_textlines)) - - all_found_textline_polygons = [] - slopes = [] - for index, contour in enumerate(contours_par): - results = [cv2.pointPolygonTest(contour, - (cx_textlines[ind], - cy_textlines[ind]), + for index, parent in enumerate(parents): + results = [cv2.pointPolygonTest(parent.contour, + (textlines_cx[ind], + textlines_cy[ind]), False) - for ind in args_textlines] + for ind in textlines_args] results = np.array(results) - indexes_in = args_textlines[results == 1] - textlines_in = self.get_textlines_of_a_textregion_sorted( - [polygons_of_textlines[ind] for ind in indexes_in], - [cx_textlines[ind] for ind in indexes_in], - [cy_textlines[ind] for ind in indexes_in], - [w_h_textlines[ind] for ind in indexes_in]) + indexes_in = textlines_args[results == 1] + get_in = itemgetter(indexes_in) + textlines_in_cont, textlines_in_conf = self.get_textlines_of_a_textregion_sorted( + get_in(textlines_cont), get_in(textlines_conf), + get_in(textlines_cx), get_in(textlines_cy), + get_in(textlines_w_h)) - all_found_textline_polygons.append(textlines_in) #[::-1]) + parent.lines = [Region(cont, conf=conf) #[::-1] + for cont, conf in zip(textlines_in_cont, textlines_in_conf)] try: - slopes.append(estimate_skew_contours(textlines_in)) + parent.skew = estimate_skew_contours(textlines_in_cont) except ValueError: - slopes.append(slope_deskew) + parent.skew = slope_deskew # plt.imshow(textline_mask_tot) # for contour in textlines_in: # plt.plot(*contour[:, 0].T, linewidth=3, color='red') # plt.show() - return all_found_textline_polygons, slopes - - def get_slopes_and_deskew_new_curved(self, contours_par, textline_mask_tot, - num_col, slope_deskew, name): - if not len(contours_par): - return [], [] + def get_slopes_and_deskew_new_curved( + self, parents: List[TextRegion], + textline_mask_tot, + textline_confidence, + num_col, slope_deskew, name + ): + if not len(parents): + return self.logger.debug("enter get_slopes_and_deskew_new_curved") - results = map(partial(do_work_of_slopes_new_curved, - textline_mask_tot_ea=textline_mask_tot, - num_col=num_col, - slope_deskew=slope_deskew, - MAX_SLOPE=MAX_SLOPE, - KERNEL=KERNEL, - logger=self.logger, - plotter=self.plotter, - name=name), - contours_par) - results = list(results) # exhaust prior to release - #textline_polygons, slopes = zip(*results) + kwargs = dict(textline_mask_tot_ea=textline_mask_tot, + num_col=num_col, + slope_deskew=slope_deskew, + MAX_SLOPE=MAX_SLOPE, + KERNEL=KERNEL, + logger=self.logger, + plotter=self.plotter, + name=name + ) + all_lines = [] + for parent in parents: + textlines_cont, skew = do_work_of_slopes_new_curved(parent.contour, **kwargs) + all_lines.extend(textlines_cont) + parent.lines = [Region(cont) for cont in textlines_cont] + parent.skew = skew + # more efficient to run all in one: + all_confs = get_region_confidences(all_lines, textline_confidence) + get_confs = iter(all_confs) + for parent in parents: + for line, conf in zip(parent.lines, get_confs): + line.conf = conf self.logger.debug("exit get_slopes_and_deskew_new_curved") - return tuple(zip(*results)) def textline_contours(self, img, use_patches): self.logger.debug('enter textline_contours') @@ -1105,20 +1112,19 @@ class Eynollah: # mask_texts_only = cv2.morphologyEx(mask_texts_only, cv2.MORPH_OPEN, KERNEL, iterations=1) mask_texts_only = cv2.dilate(mask_texts_only, kernel=np.ones((2, 2), np.uint8), iterations=1) - polygons_seplines, hir_seplines = return_contours_of_image(mask_seps_only) - polygons_seplines = filter_contours_area_of_image( - mask_seps_only, polygons_seplines, hir_seplines, max_area=1, min_area=0.00001, dilate=1) + seplines_cont, seplines_hier = return_contours_of_image(mask_seps_only) + seplines_cont = filter_contours_area_of_image( + mask_seps_only, seplines_cont, seplines_hier, max_area=1, min_area=0.00001, dilate=1) - polygons_of_only_texts = return_contours_of_interested_region(mask_texts_only,1,0.00001) - ##polygons_of_only_texts = dilate_textregion_contours(polygons_of_only_texts) - polygons_of_only_seps = return_contours_of_interested_region(mask_seps_only,1,0.00001) - polygons_of_only_tabs = return_contours_of_interested_region(mask_tabs_only,1,0.00001) + texts_only_cont = return_contours_of_interested_region(mask_texts_only,1,0.00001) + seps_only_cont = return_contours_of_interested_region(mask_seps_only,1,0.00001) + tabs_only_cont = return_contours_of_interested_region(mask_tabs_only,1,0.00001) text_regions_p = np.zeros_like(prediction_regions) - text_regions_p = cv2.fillPoly(text_regions_p, pts=polygons_of_only_seps, color=label_seps) + text_regions_p = cv2.fillPoly(text_regions_p, pts=seps_only_cont, color=label_seps) text_regions_p[mask_images_only == 1] = label_imgs - text_regions_p = cv2.fillPoly(text_regions_p, pts=polygons_of_only_texts, color=label_text) - text_regions_p = cv2.fillPoly(text_regions_p, pts=polygons_of_only_tabs, color=label_tabs) + text_regions_p = cv2.fillPoly(text_regions_p, pts=texts_only_cont, color=label_text) + text_regions_p = cv2.fillPoly(text_regions_p, pts=tabs_only_cont, color=label_tabs) textline_mask_tot_ea[text_regions_p != label_text] = 0 confidence_textline[text_regions_p != label_text] = 0 @@ -1132,8 +1138,8 @@ class Eynollah: #print("inside 4 ", time.time()-t_in) self.logger.debug("exit get_early_layout") return (erosion_hurts, - polygons_seplines, - polygons_of_only_texts, + seplines_cont, + texts_only_cont, regions_without_separators, text_regions_p, textline_mask_tot_ea, @@ -1144,14 +1150,14 @@ class Eynollah: self, contours_only_text_parent, contours_only_text_parent_h, - polygons_of_drop_capitals, + contours_drop_capitals, boxes, textline_mask_tot ): self.logger.debug("enter do_order_of_regions") contours_only_text_parent = ensure_array(contours_only_text_parent) contours_only_text_parent_h = ensure_array(contours_only_text_parent_h) - polygons_of_drop_capitals = ensure_array(polygons_of_drop_capitals) + contours_drop_capitals = ensure_array(contours_drop_capitals) boxes = np.array(boxes, dtype=int) # to be on the safe side c_boxes = np.stack((0.5 * boxes[:, 2:4].sum(axis=1), 0.5 * boxes[:, 0:2].sum(axis=1))) @@ -1189,10 +1195,10 @@ class Eynollah: def order_from_boxes(only_centers: bool): arg_text_con_main = match_boxes(contours_only_text_parent, only_centers, "main") arg_text_con_head = match_boxes(contours_only_text_parent_h, only_centers, "head") - arg_text_con_drop = match_boxes(polygons_of_drop_capitals, only_centers, "drop") + arg_text_con_drop = match_boxes(contours_drop_capitals, only_centers, "drop") args_contours_main = np.arange(len(contours_only_text_parent)) args_contours_head = np.arange(len(contours_only_text_parent_h)) - args_contours_drop = np.arange(len(polygons_of_drop_capitals)) + args_contours_drop = np.arange(len(contours_drop_capitals)) order_by_con_main = np.zeros_like(arg_text_con_main) order_by_con_head = np.zeros_like(arg_text_con_head) order_by_con_drop = np.zeros_like(arg_text_con_drop) @@ -1208,7 +1214,7 @@ class Eynollah: textline_mask_tot[ys, xs], contours_only_text_parent[args_contours_box_main], contours_only_text_parent_h[args_contours_box_head], - polygons_of_drop_capitals[args_contours_box_drop], + contours_drop_capitals[args_contours_box_drop], box[2], box[0]) for tidx, kind in zip(index_by_kind_sorted, kind_of_texts_sorted): @@ -1226,7 +1232,7 @@ class Eynollah: # xml writer will create region ids in order of # - contours_only_text_parent (main text), followed by # - contours_only_text_parent_h (headings), and then - # - polygons_of_drop_capitals, + # - contours_drop_capitals, # and then create regionrefs into these ordered by order_text_new order_text_new = np.argsort(np.concatenate((order_by_con_main, order_by_con_head, @@ -1611,7 +1617,7 @@ class Eynollah: num_col_classifier, erosion_hurts, regions_without_separators, - contours_h=None, + contours_h=[], label_seps_fl=6, ): if not erosion_hurts: @@ -1635,7 +1641,7 @@ class Eynollah: contours_only_text_parent, contours_only_text_parent_h, # not trained on drops directly, but it does work: - polygons_of_drop_capitals, + contours_drop_capitals, text_regions_p, n_batch_inference=1, # 3 (causes OOM on 8 GB GPUs) # input labels as in run_boxes_full_layout @@ -1711,7 +1717,7 @@ class Eynollah: indexes_of_located_cont.extend(args_cont_h[:, np.newaxis] + len(contours_only_text_parent)) - args_cont_drop = np.arange(len(polygons_of_drop_capitals)) + args_cont_drop = np.arange(len(contours_drop_capitals)) indexes_of_located_cont.extend(args_cont_drop[:, np.newaxis] + len(contours_only_text_parent) + len(contours_only_text_parent_h)) @@ -1735,12 +1741,12 @@ class Eynollah: img_header_and_sep[contour[:, 0, 1].max(): contour[:, 0, 1].max() + 12, contour[:, 0, 0].min(): contour[:, 0, 0].max()] = 1 co_text_all.extend(contours_only_text_parent_h) - co_text_all.extend(polygons_of_drop_capitals) + co_text_all.extend(contours_drop_capitals) if not len(co_text_all): return [] - # fill polygons in lower resolution to be faster + # fill contours in lower resolution to be faster height, width = text_regions_p.shape labels_con = np.zeros((height // 6, width // 6, len(co_text_all)), dtype=bool) for i in range(len(co_text_all)): @@ -1816,74 +1822,92 @@ class Eynollah: else: return ordered - def filter_contours_inside_a_bigger_one(self, contours, contours_d, shape, - marginal_cnts=None, type_contour="textregion"): - if type_contour == "textregion": - areas = np.array(list(map(cv2.contourArea, contours))) - areas = areas / float(np.prod(shape[:2])) - cx_main, cy_main = find_center_of_contours(contours) + def do_order_of_regions_heuristic( + self, + textregions_cont, + textregions_h_cont, + drop_caps_cont, + text_regions_p, + regions_without_separators, + num_col_classifier, + erosion_hurts, + ): + boxes = self.run_boxes_order(text_regions_p, + num_col_classifier, + erosion_hurts, + regions_without_separators, + contours_h=textregions_h_cont) + order_text = self.do_order_of_regions( + textregions_cont, + textregions_h_cont, + drop_caps_cont, + boxes, + regions_without_separators) #textline_mask_tot_ea) + return order_text - contours = ensure_array(contours) - indices_small = np.flatnonzero(areas < 1e-3) - indices_large = np.flatnonzero(areas >= 1e-3) - - indices_drop = [] - for ind_small in indices_small: - results = [cv2.pointPolygonTest(contours[ind_large], - (cx_main[ind_small], - cy_main[ind_small]), + def filter_small_regions(self, textregions: List[Region], textregions_d: List[Region], area_factor: float, marginals: List[Region]) -> Tuple[List[Region], List[Region]]: + """ + Split list of contours (and optionally deskewed contours) into + small (<0.1% area) and large (>=0.1%) candidates. Then identify + those small contours whose center point is properly contained + by some large contour (or optionally by some marginal contour). + Remove the latter ones from the list of contours (and deskewed + contours). + """ + areas = np.array([textregion.area for textregion in textregions]) * area_factor + indices_small = np.flatnonzero(areas < 1e-3) + indices_large = np.flatnonzero(areas >= 1e-3) + keep = [True] * len(areas) + for ind_small in indices_small: + results = [cv2.pointPolygonTest(textregions[ind_large].contour, + (textregions[ind_small].cx, + textregions[ind_small].cy), + False) + for ind_large in indices_large] + results = np.array(results) + if np.any(results == 1): + keep[ind_small] = False + elif len(marginals): + results = [cv2.pointPolygonTest(marginal.contour, + (textregions[ind_small].cx, + textregions[ind_small].cy), False) - for ind_large in indices_large] + for marginal in marginals] results = np.array(results) if np.any(results == 1): - indices_drop.append(ind_small) - elif marginal_cnts: - results = [cv2.pointPolygonTest(contour, - (cx_main[ind_small], - cy_main[ind_small]), - False) - for contour in marginal_cnts] - results = np.array(results) - if np.any(results == 1): - indices_drop.append(ind_small) + keep[ind_small] = False - contours = np.delete(contours, indices_drop, axis=0) - if len(contours_d): - contours_d = ensure_array(contours_d) - contours_d = np.delete(contours_d, indices_drop, axis=0) + textregions = list(compress(textregions, keep)) + if len(textregions_d): + textregions_d = list(compress(textregions_d, keep)) - return contours, contours_d + return textregions, textregions_d - else: - contours_of_contours = [] - indexes_parent = [] - indexes_child = [] - for ind_region, textlines in enumerate(contours): - contours_of_contours.extend(textlines) - indexes_parent.extend([ind_region] * len(textlines)) - indexes_child.extend(list(range(len(textlines)))) + def filter_small_textlines(self, textregions: List[TextRegion]) -> List[TextRegion]: + textlines = [] + indexes_parent = [] + indexes_child = [] + all_keep = [] + for ind_region, region in enumerate(textregions): + textlines.extend(region.lines) + indexes_parent.extend([ind_region] * len(region.lines)) + indexes_child.extend(list(range(len(region.lines)))) + all_keep.append([True] * len(region.lines)) - areas = np.array(list(map(cv2.contourArea, contours_of_contours))) - cx, cy = find_center_of_contours(contours_of_contours) + areas = np.array([textline.area for textline in textlines]) + for i, textline in enumerate(textlines): + args_other = np.setdiff1d(np.arange(len(textlines)), i) + areas_other = areas[args_other] + for ind in args_other[areas_other > 1.5 * areas[i]]: + if cv2.pointPolygonTest(textlines[ind].contour, + (textline.cx, + textline.cy), + False) == 1: + all_keep[indexes_parent[i]][indexes_child[i]] = False - textline_in_textregion_index_to_del = {} - for i in range(len(contours_of_contours)): - args_other = np.setdiff1d(np.arange(len(contours_of_contours)), i) - areas_other = areas[args_other] - args_other_larger = args_other[areas_other > 1.5 * areas[i]] - - for ind in args_other_larger: - if cv2.pointPolygonTest(contours_of_contours[ind], - (cx[i], cy[i]), False) == 1: - textline_in_textregion_index_to_del.setdefault( - indexes_parent[i], list()).append( - indexes_child[i]) - - for where, which in textline_in_textregion_index_to_del.items(): - contours[where] = [line for idx, line in enumerate(contours[where]) - if idx not in which] - - return contours + for textregion, keep in zip(textregions, all_keep): + textregion.lines = list(compress(textregion.lines, keep)) + return textregions def return_indexes_of_contours_located_inside_another_list_of_contours( self, contours, centersx_loc, centersy_loc, indexes_loc): @@ -1900,50 +1924,21 @@ class Eynollah: return indexes, centersx, centersy - def filter_contours_without_textline_inside( - self, contours_textregions, contours_textregions_d, - contours_textlines, slopes, conf_contours_textregions): + def filter_textregions_without_textlines(self, textregions, textregions_d): + keep = [len(textregion.lines) > 0 for textregion in textregions] + return (list(compress(textregions, keep)), + list(compress(textregions_d, keep))) - assert len(contours_textregions) == len(contours_textlines) - indices = [ind for ind, lines in enumerate(contours_textlines) - if len(lines)] - def filterfun(lis): - if len(lis) == 0: - return [] - return [lis[ind] for ind in indices] + def separate_marginals_and_order(self, marginals, mid_point_of_page_width): + left = [] + right = [] + for marginal in marginals: + (left, right)[marginal.cx < mid_point_of_page_width].append(marginal) - return (filterfun(contours_textregions), - filterfun(contours_textregions_d), - filterfun(contours_textlines), - filterfun(slopes), - filterfun(conf_contours_textregions), - ) + order_left = itemgetter(np.argsort([marginal.cy for marginal in left])) + order_right = itemgetter(np.argsort([marginal.cy for marginal in right])) - def separate_marginals_to_left_and_right_and_order_from_top_to_down( - self, polygons_of_marginals, all_found_textline_polygons_marginals, - slopes_marginals, conf_marginals, mid_point_of_page_width): - cx_marg, cy_marg = find_center_of_contours(polygons_of_marginals) - cx_marg = ensure_array(cx_marg) - cy_marg = ensure_array(cy_marg) - - def split(lis): - left, right = [], [] - for itm, prop in zip(lis, cx_marg < mid_point_of_page_width): - (left if prop else right).append(itm) - return left, right - - cy_marg_left, cy_marg_right = split(cy_marg) - order_left = np.argsort(cy_marg_left) - order_right = np.argsort(cy_marg_right) - - def splitsort(lis): - left, right = split(lis) - return [left[i] for i in order_left], [right[i] for i in order_right] - - return (*splitsort(polygons_of_marginals), - *splitsort(all_found_textline_polygons_marginals), - *splitsort(slopes_marginals), - *splitsort(conf_marginals)) + return order_left(left), order_right(right) def run(self, overwrite: bool = False, @@ -2067,7 +2062,6 @@ class Eynollah: image_filename=img_filename, image_width=image['img'].shape[1], image_height=image['img'].shape[0], - curved_line=self.curved_line, pcgts=pcgts) if os.path.exists(writer.output_filename): @@ -2090,7 +2084,8 @@ class Eynollah: self.logger.info(f"Enhancement complete ({time.time() - t0:.1f}s)") t1 = time.time() - page_coord, cont_page, image_page, mask_page = self.extract_page(image) + page_cont, image_page, mask_page = self.extract_page(image) + page = Region(page_cont) if not self.ignore_page_extraction: self.logger.debug("Cropped page is %dx%d", image_page.shape[1], image_page.shape[0]) self.logger.info("Cropping took %.1fs", time.time() - t1) @@ -2102,46 +2097,35 @@ class Eynollah: self.logger.info("Step 2/5: Basic Processing Mode") self.logger.info("Skipping layout analysis and reading order detection") - _, _, _, _, _, textline_mask_tot_ea, _, _ = \ + _, _, _, _, _, textline_mask_tot_ea, _, textline_confidence = \ self.get_early_layout(image_page, num_col_classifier) textline_mask_tot_ea *= mask_page - textline_cnt, textline_hir = return_contours_of_image(textline_mask_tot_ea) - all_found_textline_polygons = filter_contours_area_of_image( - textline_mask_tot_ea, textline_cnt, textline_hir, max_area=1, min_area=0.00001) + textline_confidence *= mask_page + textlines_cont = return_contours_of_interested_region(textline_mask_tot_ea, 1, 0.00001) + textlines_conf = get_region_confidences(textlines_cont, textline_confidence) - cx_textlines, cy_textlines = find_center_of_contours(all_found_textline_polygons) - w_h_textlines = [cv2.boundingRect(polygon)[2:] - for polygon in all_found_textline_polygons] - all_found_textline_polygons = self.get_textlines_of_a_textregion_sorted( - #all_found_textline_polygons[::-1] - all_found_textline_polygons, cx_textlines, cy_textlines, w_h_textlines) - all_found_textline_polygons = [all_found_textline_polygons] - all_found_textline_polygons = dilate_textline_contours(all_found_textline_polygons) - all_found_textline_polygons = self.filter_contours_inside_a_bigger_one( - all_found_textline_polygons, None, None, type_contour="textline") - - pcgts = writer.build_pagexml_no_full_layout( - num_col=num_col_classifier, - found_polygons_text_region=cont_page, - page_coord=page_coord, - page_slope=0, - order_of_texts=[0], - all_found_textline_polygons=all_found_textline_polygons, - found_polygons_images=[], - found_polygons_tables=[], - found_polygons_marginals_left=[], - found_polygons_marginals_right=[], - all_found_textline_polygons_marginals_left=[], - all_found_textline_polygons_marginals_right=[], - slopes=[0], - slopes_marginals_left=[], - slopes_marginals_right=[], - cont_page=cont_page, - polygons_seplines=[], - conf_textregions=[0], - ) + textlines_cx, textlines_cy = find_center_of_contours(textlines_cont) + textlines_w_h = [cv2.boundingRect(cont)[2:] + for cont in textlines_cont] + textlines_cont, textlines_conf = self.get_textlines_of_a_textregion_sorted( + textlines_cont, textlines_conf, + textlines_cx, textlines_cy, textlines_w_h) + textregions = [ + TextRegion(page.contour, lines=[ + Region(cont, conf=conf) + for cont, conf in zip(textlines_cont, textlines_conf)]) + ] + textregions = self.filter_small_textlines(textregions) self.logger.info("Basic processing complete") + + pcgts = writer.build_pagexml( + page=page, + img_bin=self.imread(image, binary=True) if self.input_binary else None, + num_col=num_col_classifier, + order_of_texts=[0], + textregions=textregions, + ) if writer.pcgts is None: writer.write_pagexml(pcgts) self.logger.info("Job done in %.1fs", time.time() - t0) @@ -2151,8 +2135,8 @@ class Eynollah: self.logger.info("Step 2/5: Layout Analysis") (erosion_hurts, - polygons_seplines, - polygons_text_early, + seplines_cont, + text_early_cont, regions_without_separators, text_regions_p, textline_mask_tot_ea, @@ -2179,19 +2163,24 @@ class Eynollah: if (abs(slope_deskew) > 45 and ((text_regions_p == label_text).sum()) <= 0.3 * image_page.size): slope_deskew = 0 + page.skew = slope_deskew if self.plotter: self.plotter.save_deskewed_image(slope_deskew, image['img'], image['name']) t3 = time.time() self.logger.info("Deskewing took %.1fs", t3 - t2) - page_coord = np.array(page_coord) - page_box = (slice(*page_coord[:2]), - slice(*page_coord[2:])) - polygons_seplines = [contour - page_coord[::2][::-1][np.newaxis, np.newaxis] - for contour in polygons_seplines] + # FIXME: post-hoc cropping (remove when models support it, and replace image['img_res'] with image_page) + page_box = cv2.boundingRect(page.contour) + seplines_conf = get_region_confidences(seplines_cont, regions_confidence) + seplines = [Region(cont - [page_box[:2]], + conf=conf) + for cont, conf in zip(seplines_cont, seplines_conf)] + page_box = box2slice(page_box) regions_without_separators = regions_without_separators[page_box] * mask_page text_regions_p = text_regions_p[page_box] * mask_page textline_mask_tot_ea = textline_mask_tot_ea[page_box] * mask_page + regions_confidence = regions_confidence[page_box] * mask_page + textline_confidence = textline_confidence[page_box] * mask_page num_col, num_col_classifier = \ self.run_columns(text_regions_p, @@ -2200,27 +2189,13 @@ class Eynollah: t4 = time.time() textline_mask_tot_ea_org = np.copy(textline_mask_tot_ea) - if not num_col and len(polygons_text_early) == 0 or not image_page.size: + if not num_col and len(text_early_cont) == 0 or not image_page.size: self.logger.info("No columns detected - generating empty PAGE-XML") - pcgts = writer.build_pagexml_no_full_layout( + pcgts = writer.build_pagexml( + page=page, + img_bin=self.imread(image, binary=True) if self.input_binary else None, num_col=0, - found_polygons_text_region=[], - page_coord=page_coord, - page_slope=slope_deskew, - order_of_texts=[], - all_found_textline_polygons=[], - found_polygons_images=[], - found_polygons_tables=[], - found_polygons_marginals_left=[], - found_polygons_marginals_right=[], - all_found_textline_polygons_marginals_left=[], - all_found_textline_polygons_marginals_right=[], - slopes=[], - slopes_marginals_left=[], - slopes_marginals_right=[], - cont_page=cont_page, - polygons_seplines=[], ) if writer.pcgts is None: writer.write_pagexml(pcgts) @@ -2255,30 +2230,37 @@ class Eynollah: regions_without_separators[text_regions_p == label_drop_fl] = 1 # also cover in reading-order textline_mask_tot_ea_org[text_regions_p == label_drop_fl] = 0 # skip for textlines textline_mask_tot_ea[text_regions_p == label_drop_fl] = 1 # needed for reading order - polygons_of_drop_capitals = return_contours_of_interested_region(text_regions_p, - label_drop_fl, - min_area=0.00003) - conf_drops = get_region_confidences(polygons_of_drop_capitals, regionsfl_confidence) + drop_caps_cont = return_contours_of_interested_region(text_regions_p, label_drop_fl, + min_area=0.00003) + drop_caps_conf = get_region_confidences(drop_caps_cont, regionsfl_confidence) + drop_caps = [Region(cont, conf=conf) + for cont, conf in zip(drop_caps_cont, drop_caps_conf)] t6 = time.time() self.logger.info("Full layout took %.1fs", t6 - t5) else: + drop_caps = [] t6 = time.time() self.logger.info("Step 3/5: Contour extraction") min_area_mar = 0.00001 marginal_mask = (text_regions_p == label_marg_fl).astype(np.uint8) marginal_mask = cv2.dilate(marginal_mask, KERNEL, iterations=2) - polygons_of_marginals = return_contours_of_interested_region(marginal_mask, 1, - min_area_mar) - polygons_of_tables = return_contours_of_interested_region(text_regions_p, label_tabs, - min_area_mar) - polygons_of_images = return_contours_of_interested_region(text_regions_p, label_imgs_fl) - conf_marginals = get_region_confidences(polygons_of_marginals, regions_confidence) - conf_images = get_region_confidences(polygons_of_images, regions_confidence) - conf_tables = get_region_confidences(polygons_of_tables, regions_confidence) + marginals_cont = return_contours_of_interested_region(marginal_mask, 1, min_area_mar) + marginals_conf = get_region_confidences(marginals_cont, regions_confidence) + marginals = [Region(cont, conf=conf) + for cont, conf in zip(marginals_cont, marginals_conf)] + tables_cont = return_contours_of_interested_region(text_regions_p, label_tabs, min_area_mar) + tables_conf = get_region_confidences(tables_cont, regions_confidence) + tables = [Region(cont, conf=conf) + for cont, conf in zip(tables_cont, tables_conf)] + images_cont = return_contours_of_interested_region(text_regions_p, label_imgs_fl) + images_conf = get_region_confidences(images_cont, regions_confidence) + images = [Region(cont, conf=conf) + for cont, conf in zip(images_cont, images_conf)] - polygons_of_textregions = return_contours_of_interested_region(text_regions_p, label_text, - min_area=MIN_AREA_REGION) + textregions_cont = return_contours_of_interested_region(text_regions_p, label_text, + min_area=MIN_AREA_REGION) + textregions = [TextRegion(cont, lines=[]) for cont in textregions_cont] if np.abs(slope_deskew) >= SLOPE_THRESHOLD and not self.reading_order_machine_based: (text_regions_p_d, @@ -2289,129 +2271,87 @@ class Eynollah: textline_mask_tot_ea, regions_without_separators) - polygons_of_textregions_d = return_contours_of_interested_region(text_regions_p_d, label_text, - min_area=MIN_AREA_REGION) - if (len(polygons_of_textregions) and - len(polygons_of_textregions_d)): - polygons_of_textregions_d = \ + textregions_cont_d = return_contours_of_interested_region(text_regions_p_d, label_text, + min_area=MIN_AREA_REGION) + textregions_d = [TextRegion(cont, lines=[]) for cont in textregions_cont_d] + if (len(textregions) and + len(textregions_d)): + textregions_cont_d = \ match_deskewed_contours( slope_deskew, - polygons_of_textregions, - polygons_of_textregions_d, + textregions, + textregions_d, text_regions_p.shape, text_regions_p_d.shape) + textregions_d = [TextRegion(cont, lines=[]) for cont in textregions_cont_d] else: - polygons_of_textregions_d = [] - (polygons_of_textregions, - polygons_of_textregions_d) = self.filter_contours_inside_a_bigger_one( - polygons_of_textregions, - polygons_of_textregions_d, - text_regions_p.shape, - marginal_cnts=polygons_of_marginals) - polygons_of_textregions = dilate_textregion_contours(polygons_of_textregions) - conf_textregions = get_region_confidences(polygons_of_textregions, regions_confidence) + textregions_d = [] + + area_factor = np.reciprocal(np.prod(text_regions_p.shape).astype(float)) + textregions, textregions_d = self.filter_small_regions( + textregions, textregions_d, + area_factor, + marginals) + textregions_conf = get_region_confidences(textregions_cont, regions_confidence) + for textregion, conf in zip(textregions, textregions_conf): + textregion.conf = conf - if not len(polygons_of_textregions): - polygons_of_textregions = polygons_of_marginals - polygons_of_marginals = [] - conf_textregions = conf_marginals - conf_marginals = [] t7 = time.time() self.logger.info("Region contours took %.1fs", t7 - t6) if not self.curved_line: self.logger.info("Mode: Light line detection") - all_found_textline_polygons, slopes = \ - self.get_slopes_and_deskew_new_light2( - polygons_of_textregions, textline_mask_tot_ea_org, + args = (textline_mask_tot_ea_org, + textline_confidence, slope_deskew) - all_found_textline_polygons_marginals, slopes_marginals = \ - self.get_slopes_and_deskew_new_light2( - polygons_of_marginals, textline_mask_tot_ea_org, - slope_deskew) - - all_found_textline_polygons = dilate_textline_contours( - all_found_textline_polygons) - all_found_textline_polygons = self.filter_contours_inside_a_bigger_one( - all_found_textline_polygons, None, None, type_contour="textline") - all_found_textline_polygons_marginals = dilate_textline_contours( - all_found_textline_polygons_marginals) + self.get_slopes_and_deskew_new_light2(textregions, *args) + self.get_slopes_and_deskew_new_light2(marginals, *args) + textregions = self.filter_small_textlines(textregions) else: self.logger.info("Mode: Curved line detection") textline_mask_tot_ea_erode = cv2.erode(textline_mask_tot_ea_org, kernel=KERNEL, iterations=2) - all_found_textline_polygons, slopes = \ - self.get_slopes_and_deskew_new_curved( - polygons_of_textregions, textline_mask_tot_ea_erode, - num_col_classifier, slope_deskew, image['name']) - all_found_textline_polygons = small_textlines_to_parent_adherence2( - all_found_textline_polygons, textline_mask_tot_ea, num_col_classifier) - all_found_textline_polygons_marginals, slopes_marginals = \ - self.get_slopes_and_deskew_new_curved( - polygons_of_marginals, textline_mask_tot_ea_erode, - num_col_classifier, slope_deskew, image['name']) - all_found_textline_polygons_marginals = small_textlines_to_parent_adherence2( - all_found_textline_polygons_marginals, textline_mask_tot_ea, num_col_classifier) - (polygons_of_textregions, - polygons_of_textregions_d, - all_found_textline_polygons, - slopes, - conf_textregions) = \ - self.filter_contours_without_textline_inside( - polygons_of_textregions, - polygons_of_textregions_d, - all_found_textline_polygons, - slopes, - conf_textregions) + args = (textline_mask_tot_ea_erode, + textline_confidence, + num_col_classifier, + slope_deskew, + image['name']) + self.get_slopes_and_deskew_new_curved(textregions, *args) + small_textlines_to_parent_adherence2(textregions, area_factor, num_col_classifier) + self.get_slopes_and_deskew_new_curved(marginals, *args) + small_textlines_to_parent_adherence2(marginals, area_factor, num_col_classifier) + + textregions, textregions_d = self.filter_textregions_without_textlines( + textregions, textregions_d) t8 = time.time() self.logger.info("Line contours took %.1fs", t8 - t7) - (polygons_of_marginals_left, - polygons_of_marginals_right, - all_found_textline_polygons_marginals_left, - all_found_textline_polygons_marginals_right, - slopes_marginals_left, - slopes_marginals_right, - conf_marginals_left, - conf_marginals_right) = \ - self.separate_marginals_to_left_and_right_and_order_from_top_to_down( - polygons_of_marginals, - all_found_textline_polygons_marginals, - slopes_marginals, - conf_marginals, - 0.5 * text_regions_p.shape[1]) - # FIXME: get_region_confidences w/ textline_confidence on all types of textlines... + (marginals_left, + marginals_right) = self.separate_marginals_and_order( + marginals, 0.5 * text_regions_p.shape[1]) if self.full_layout: (text_regions_p, - polygons_of_textregions, - polygons_of_textregions_h, - polygons_of_textregions_d, - polygons_of_textregions_h_d, - all_found_textline_polygons, - all_found_textline_polygons_h, - slopes, - slopes_h, - conf_textregions, - conf_textregions_h) = split_textregion_main_vs_head( + textregions, + textregions_h, + textregions_d, + textregions_h_d) = split_textregion_main_vs_head( text_regions_p, regions_fully, - polygons_of_textregions, - polygons_of_textregions_d, - all_found_textline_polygons, - slopes, - conf_textregions) + textregions, + textregions_d) if self.plotter: self.plotter.save_plot_of_layout(text_regions_p, image_page, image['name']) self.plotter.save_plot_of_layout_all(text_regions_p, image_page, image['name']) else: - polygons_of_drop_capitals = [] - polygons_of_textregions_h = [] - polygons_of_textregions_h_d = [] + textregions_h = [] + textregions_h_d = [] + def contours(regions): + return [region.contour for region in regions] if self.plotter: - self.plotter.write_images_into_directory(polygons_of_images, image_page, + self.plotter.write_images_into_directory(contours(images), image_page, image['scale_x'], image['scale_y'], image['name']) t_order = time.time() @@ -2424,92 +2364,46 @@ class Eynollah: if self.reading_order_machine_based: self.logger.info("Using machine-based detection") order_text = self.do_order_of_regions_with_model( - polygons_of_textregions, - polygons_of_textregions_h, - polygons_of_drop_capitals, + contours(textregions), + contours(textregions_h) if not self.headers_off else [], + contours(drop_caps), text_regions_p) else: if np.abs(slope_deskew) < SLOPE_THRESHOLD: - boxes = self.run_boxes_order(text_regions_p, num_col_classifier, erosion_hurts, - regions_without_separators, - contours_h=(None if self.headers_off or not self.full_layout - else polygons_of_textregions_h)) - order_text = self.do_order_of_regions( - polygons_of_textregions, - polygons_of_textregions_h, - polygons_of_drop_capitals, - boxes, regions_without_separators) #textline_mask_tot_ea) + order_text = self.do_order_of_regions_heuristic( + contours(textregions), + contours(textregions_h) if not self.headers_off else [], + contours(drop_caps), + text_regions_p, + regions_without_separators, + num_col_classifier, + erosion_hurts) else: - boxes_d = self.run_boxes_order(text_regions_p_d, num_col_classifier, erosion_hurts, - regions_without_separators_d, - contours_h=(None if self.headers_off or not self.full_layout - else polygons_of_textregions_h_d)) - - order_text = self.do_order_of_regions( - polygons_of_textregions_d, - polygons_of_textregions_h_d, - polygons_of_drop_capitals, - boxes_d, regions_without_separators_d) #textline_mask_tot_ea_d) + order_text = self.do_order_of_regions_heuristic( + contours(textregions_d), + contours(textregions_h_d) if not self.headers_off else [], + contours(drop_caps), + text_regions_p_d, + regions_without_separators_d, + num_col_classifier, + erosion_hurts) self.logger.info(f"Detection of reading order took {time.time() - t_order:.1f}s") self.logger.info("Step 5/5: Output Generation") - if self.full_layout: - pcgts = writer.build_pagexml_full_layout( - num_col=num_col_classifier, - found_polygons_text_region=polygons_of_textregions, - found_polygons_text_region_h=polygons_of_textregions_h, - page_coord=page_coord, - page_slope=slope_deskew, - order_of_texts=order_text, - all_found_textline_polygons=all_found_textline_polygons, - all_found_textline_polygons_h=all_found_textline_polygons_h, - found_polygons_images=polygons_of_images, - found_polygons_tables=polygons_of_tables, - found_polygons_drop_capitals=polygons_of_drop_capitals, - found_polygons_marginals_left=polygons_of_marginals_left, - found_polygons_marginals_right=polygons_of_marginals_right, - all_found_textline_polygons_marginals_left=all_found_textline_polygons_marginals_left, - all_found_textline_polygons_marginals_right=all_found_textline_polygons_marginals_right, - slopes=slopes, - slopes_h=slopes_h, - slopes_marginals_left=slopes_marginals_left, - slopes_marginals_right=slopes_marginals_right, - cont_page=cont_page, - polygons_seplines=polygons_seplines, - conf_textregions=conf_textregions, - conf_textregions_h=conf_textregions_h, - conf_marginals_left=conf_marginals_left, - conf_marginals_right=conf_marginals_right, - conf_images=conf_images, - conf_tables=conf_tables, - conf_drops=conf_drops, - ) - else: - pcgts = writer.build_pagexml_no_full_layout( - num_col=num_col_classifier, - found_polygons_text_region=polygons_of_textregions, - page_coord=page_coord, - page_slope=slope_deskew, - order_of_texts=order_text, - all_found_textline_polygons=all_found_textline_polygons, - found_polygons_images=polygons_of_images, - found_polygons_tables=polygons_of_tables, - found_polygons_marginals_left=polygons_of_marginals_left, - found_polygons_marginals_right=polygons_of_marginals_right, - all_found_textline_polygons_marginals_left=all_found_textline_polygons_marginals_left, - all_found_textline_polygons_marginals_right=all_found_textline_polygons_marginals_right, - slopes=slopes, - slopes_marginals_left=slopes_marginals_left, - slopes_marginals_right=slopes_marginals_right, - cont_page=cont_page, - polygons_seplines=polygons_seplines, - conf_textregions=conf_textregions, - conf_marginals_left=conf_marginals_left, - conf_marginals_right=conf_marginals_right, - conf_images=conf_images, - conf_tables=conf_tables, - ) - + pcgts = writer.build_pagexml( + page=page, + img_bin=self.imread(image, binary=True) if self.input_binary else None, + num_col=num_col_classifier, + order_of_texts=order_text, + textregions=textregions, + textregions_h=textregions_h, + images=images, + tables=tables, + drop_caps=drop_caps, + marginals_left=marginals_left, + marginals_right=marginals_right, + seplines=seplines, + ) if writer.pcgts is None: writer.write_pagexml(pcgts) self.logger.info("Job done in %.1fs", time.time() - t0) diff --git a/src/eynollah/image_enhancer.py b/src/eynollah/image_enhancer.py index fe1e16d..f515549 100644 --- a/src/eynollah/image_enhancer.py +++ b/src/eynollah/image_enhancer.py @@ -3,9 +3,9 @@ Image enhancer. The output can be written as same scale of input or in new predi """ import logging +import time import os from typing import Optional -from pathlib import Path import cv2 @@ -48,6 +48,8 @@ class Enhancer(Eynollah): dir_out: Optional[str] = None, overwrite: bool = False, ) -> None: + t0 = time.time() + self.logger.info(img_filename) image = self.cache_images(image_filename=img_filename, image_pil=img_pil) output_filename = os.path.join(dir_out or "", image['name'] + '.png') @@ -67,6 +69,7 @@ class Enhancer(Eynollah): cv2.imwrite(output_filename, img_res) self.logger.info("output filename: '%s'", output_filename) + self.logger.info("Job done in %.1fs", time.time() - t0) def run(self, overwrite: bool = False, @@ -78,6 +81,7 @@ class Enhancer(Eynollah): Enlarge and enhance the scanned images """ if dir_in: + t0_tot = time.time() ls_imgs = [os.path.join(dir_in, image_filename) for image_filename in filter(is_image_filename, os.listdir(dir_in))] @@ -87,8 +91,8 @@ class Enhancer(Eynollah): raise ValueError("run requires either a single image filename or a directory") for img_filename in ls_imgs: - self.logger.info(img_filename) - self.run_single(img_filename, dir_out=dir_out, overwrite=overwrite) + if dir_in: + self.logger.info("All jobs done in %.1fs", time.time() - t0_tot) diff --git a/src/eynollah/mb_ro_on_layout.py b/src/eynollah/mb_ro_on_layout.py deleted file mode 100644 index 5bab608..0000000 --- a/src/eynollah/mb_ro_on_layout.py +++ /dev/null @@ -1,806 +0,0 @@ -""" -Machine learning based reading order detection -""" - -# pyright: reportCallIssue=false -# pyright: reportUnboundVariable=false -# pyright: reportArgumentType=false - -import logging -import os -import time -from typing import Optional -from pathlib import Path -import xml.etree.ElementTree as ET - -import cv2 -import numpy as np -import statistics - -from .eynollah import Eynollah -from .model_zoo import EynollahModelZoo -from .utils.resize import resize_image -from .utils.contour import ( - find_new_features_of_contours, - return_contours_of_image, - return_parent_contours, -) -from .utils import is_xml_filename - -DPI_THRESHOLD = 298 -KERNEL = np.ones((5, 5), np.uint8) - - -class Reorder(Eynollah): - def __init__( - self, - *, - model_zoo: EynollahModelZoo, - logger : Optional[logging.Logger] = None, - device: str = '', - ): - self.logger = logger or logging.getLogger('eynollah.mbreorder') - self.model_zoo = model_zoo - - self.setup_models(device=device) - - def setup_models(self, device=''): - loadable = ['reading_order'] - self.model_zoo.load_models(*loadable, device=device) - for model in loadable: - self.logger.debug("model %s has input shape %s", model, - self.model_zoo.get(model).input_shape) - - - def read_xml(self, xml_file): - tree1 = ET.parse(xml_file, parser = ET.XMLParser(encoding='utf-8')) - root1=tree1.getroot() - alltags=[elem.tag for elem in root1.iter()] - link=alltags[0].split('}')[0]+'}' - - index_tot_regions = [] - tot_region_ref = [] - - y_len, x_len = 0, 0 - for jj in root1.iter(link+'Page'): - y_len=int(jj.attrib['imageHeight']) - x_len=int(jj.attrib['imageWidth']) - - for jj in root1.iter(link+'RegionRefIndexed'): - index_tot_regions.append(jj.attrib['index']) - tot_region_ref.append(jj.attrib['regionRef']) - - if (link+'PrintSpace' in alltags) or (link+'Border' in alltags): - co_printspace = [] - if link+'PrintSpace' in alltags: - region_tags_printspace = np.unique([x for x in alltags if x.endswith('PrintSpace')]) - else: - region_tags_printspace = np.unique([x for x in alltags if x.endswith('Border')]) - - for tag in region_tags_printspace: - if link+'PrintSpace' in alltags: - tag_endings_printspace = ['}PrintSpace','}printspace'] - else: - tag_endings_printspace = ['}Border','}border'] - - if tag.endswith(tag_endings_printspace[0]) or tag.endswith(tag_endings_printspace[1]): - for nn in root1.iter(tag): - c_t_in = [] - sumi = 0 - for vv in nn.iter(): - # check the format of coords - if vv.tag == link + 'Coords': - coords = bool(vv.attrib) - if coords: - p_h = vv.attrib['points'].split(' ') - c_t_in.append( - np.array([[int(x.split(',')[0]), int(x.split(',')[1])] for x in p_h])) - break - else: - pass - - if vv.tag == link + 'Point': - c_t_in.append([int(float(vv.attrib['x'])), int(float(vv.attrib['y']))]) - sumi += 1 - elif vv.tag != link + 'Point' and sumi >= 1: - break - co_printspace.append(np.array(c_t_in)) - img_printspace = np.zeros( (y_len,x_len,3) ) - img_printspace=cv2.fillPoly(img_printspace, pts =co_printspace, color=(1,1,1)) - img_printspace = img_printspace.astype(np.uint8) - - imgray = cv2.cvtColor(img_printspace, cv2.COLOR_BGR2GRAY) - _, thresh = cv2.threshold(imgray, 0, 255, 0) - contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - cnt_size = np.array([cv2.contourArea(contours[j]) for j in range(len(contours))]) - cnt = contours[np.argmax(cnt_size)] - x, y, w, h = cv2.boundingRect(cnt) - - bb_coord_printspace = [x, y, w, h] - - else: - bb_coord_printspace = None - - - region_tags=np.unique([x for x in alltags if x.endswith('Region')]) - co_text_paragraph=[] - co_text_drop=[] - co_text_heading=[] - co_text_header=[] - co_text_marginalia=[] - co_text_catch=[] - co_text_page_number=[] - co_text_signature_mark=[] - co_sep=[] - co_img=[] - co_table=[] - co_graphic=[] - co_graphic_text_annotation=[] - co_graphic_decoration=[] - co_noise=[] - - co_text_paragraph_text=[] - co_text_drop_text=[] - co_text_heading_text=[] - co_text_header_text=[] - co_text_marginalia_text=[] - co_text_catch_text=[] - co_text_page_number_text=[] - co_text_signature_mark_text=[] - co_sep_text=[] - co_img_text=[] - co_table_text=[] - co_graphic_text=[] - co_graphic_text_annotation_text=[] - co_graphic_decoration_text=[] - co_noise_text=[] - - id_paragraph = [] - id_header = [] - id_heading = [] - id_marginalia = [] - - for tag in region_tags: - if tag.endswith('}TextRegion') or tag.endswith('}Textregion'): - for nn in root1.iter(tag): - for child2 in nn: - tag2 = child2.tag - if tag2.endswith('}TextEquiv') or tag2.endswith('}TextEquiv'): - for childtext2 in child2: - if childtext2.tag.endswith('}Unicode') or childtext2.tag.endswith('}Unicode'): - if "type" in nn.attrib and nn.attrib['type']=='drop-capital': - co_text_drop_text.append(childtext2.text) - elif "type" in nn.attrib and nn.attrib['type']=='heading': - co_text_heading_text.append(childtext2.text) - elif "type" in nn.attrib and nn.attrib['type']=='signature-mark': - co_text_signature_mark_text.append(childtext2.text) - elif "type" in nn.attrib and nn.attrib['type']=='header': - co_text_header_text.append(childtext2.text) - ###elif "type" in nn.attrib and nn.attrib['type']=='catch-word': - ###co_text_catch_text.append(childtext2.text) - ###elif "type" in nn.attrib and nn.attrib['type']=='page-number': - ###co_text_page_number_text.append(childtext2.text) - elif "type" in nn.attrib and nn.attrib['type']=='marginalia': - co_text_marginalia_text.append(childtext2.text) - else: - co_text_paragraph_text.append(childtext2.text) - c_t_in_drop=[] - c_t_in_paragraph=[] - c_t_in_heading=[] - c_t_in_header=[] - c_t_in_page_number=[] - c_t_in_signature_mark=[] - c_t_in_catch=[] - c_t_in_marginalia=[] - - - sumi=0 - for vv in nn.iter(): - # check the format of coords - if vv.tag==link+'Coords': - - coords=bool(vv.attrib) - if coords: - #print('birda1') - p_h=vv.attrib['points'].split(' ') - - - - if "type" in nn.attrib and nn.attrib['type']=='drop-capital': - - c_t_in_drop.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif "type" in nn.attrib and nn.attrib['type']=='heading': - ##id_heading.append(nn.attrib['id']) - c_t_in_heading.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - - elif "type" in nn.attrib and nn.attrib['type']=='signature-mark': - - c_t_in_signature_mark.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - #print(c_t_in_paragraph) - elif "type" in nn.attrib and nn.attrib['type']=='header': - #id_header.append(nn.attrib['id']) - c_t_in_header.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - - ###elif "type" in nn.attrib and nn.attrib['type']=='catch-word': - ###c_t_in_catch.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - - ###elif "type" in nn.attrib and nn.attrib['type']=='page-number': - - ###c_t_in_page_number.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif "type" in nn.attrib and nn.attrib['type']=='marginalia': - #id_marginalia.append(nn.attrib['id']) - - c_t_in_marginalia.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - else: - #id_paragraph.append(nn.attrib['id']) - - c_t_in_paragraph.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - break - else: - pass - - - if vv.tag==link+'Point': - if "type" in nn.attrib and nn.attrib['type']=='drop-capital': - - c_t_in_drop.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - elif "type" in nn.attrib and nn.attrib['type']=='heading': - #id_heading.append(nn.attrib['id']) - c_t_in_heading.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - - elif "type" in nn.attrib and nn.attrib['type']=='signature-mark': - - c_t_in_signature_mark.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - elif "type" in nn.attrib and nn.attrib['type']=='header': - #id_header.append(nn.attrib['id']) - c_t_in_header.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - - ###elif "type" in nn.attrib and nn.attrib['type']=='catch-word': - ###c_t_in_catch.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - ###sumi+=1 - - ###elif "type" in nn.attrib and nn.attrib['type']=='page-number': - - ###c_t_in_page_number.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - ###sumi+=1 - - elif "type" in nn.attrib and nn.attrib['type']=='marginalia': - #id_marginalia.append(nn.attrib['id']) - - c_t_in_marginalia.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - else: - #id_paragraph.append(nn.attrib['id']) - c_t_in_paragraph.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - elif vv.tag!=link+'Point' and sumi>=1: - break - - if len(c_t_in_drop)>0: - co_text_drop.append(np.array(c_t_in_drop)) - if len(c_t_in_paragraph)>0: - co_text_paragraph.append(np.array(c_t_in_paragraph)) - id_paragraph.append(nn.attrib['id']) - if len(c_t_in_heading)>0: - co_text_heading.append(np.array(c_t_in_heading)) - id_heading.append(nn.attrib['id']) - - if len(c_t_in_header)>0: - co_text_header.append(np.array(c_t_in_header)) - id_header.append(nn.attrib['id']) - if len(c_t_in_page_number)>0: - co_text_page_number.append(np.array(c_t_in_page_number)) - if len(c_t_in_catch)>0: - co_text_catch.append(np.array(c_t_in_catch)) - - if len(c_t_in_signature_mark)>0: - co_text_signature_mark.append(np.array(c_t_in_signature_mark)) - - if len(c_t_in_marginalia)>0: - co_text_marginalia.append(np.array(c_t_in_marginalia)) - id_marginalia.append(nn.attrib['id']) - - - elif tag.endswith('}GraphicRegion') or tag.endswith('}graphicregion'): - for nn in root1.iter(tag): - c_t_in=[] - c_t_in_text_annotation=[] - c_t_in_decoration=[] - sumi=0 - for vv in nn.iter(): - # check the format of coords - if vv.tag==link+'Coords': - coords=bool(vv.attrib) - if coords: - p_h=vv.attrib['points'].split(' ') - - if "type" in nn.attrib and nn.attrib['type']=='handwritten-annotation': - c_t_in_text_annotation.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - elif "type" in nn.attrib and nn.attrib['type']=='decoration': - c_t_in_decoration.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - else: - c_t_in.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - - - break - else: - pass - - - if vv.tag==link+'Point': - if "type" in nn.attrib and nn.attrib['type']=='handwritten-annotation': - c_t_in_text_annotation.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - elif "type" in nn.attrib and nn.attrib['type']=='decoration': - c_t_in_decoration.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - else: - c_t_in.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - if len(c_t_in_text_annotation)>0: - co_graphic_text_annotation.append(np.array(c_t_in_text_annotation)) - if len(c_t_in_decoration)>0: - co_graphic_decoration.append(np.array(c_t_in_decoration)) - if len(c_t_in)>0: - co_graphic.append(np.array(c_t_in)) - - - - elif tag.endswith('}ImageRegion') or tag.endswith('}imageregion'): - for nn in root1.iter(tag): - c_t_in=[] - sumi=0 - for vv in nn.iter(): - # check the format of coords - if vv.tag==link+'Coords': - coords=bool(vv.attrib) - if coords: - p_h=vv.attrib['points'].split(' ') - c_t_in.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - break - else: - pass - - - if vv.tag==link+'Point': - c_t_in.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - elif vv.tag!=link+'Point' and sumi>=1: - break - co_img.append(np.array(c_t_in)) - co_img_text.append(' ') - - - elif tag.endswith('}SeparatorRegion') or tag.endswith('}separatorregion'): - for nn in root1.iter(tag): - c_t_in=[] - sumi=0 - for vv in nn.iter(): - # check the format of coords - if vv.tag==link+'Coords': - coords=bool(vv.attrib) - if coords: - p_h=vv.attrib['points'].split(' ') - c_t_in.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - break - else: - pass - - - if vv.tag==link+'Point': - c_t_in.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - elif vv.tag!=link+'Point' and sumi>=1: - break - co_sep.append(np.array(c_t_in)) - - - - elif tag.endswith('}TableRegion') or tag.endswith('}tableregion'): - for nn in root1.iter(tag): - c_t_in=[] - sumi=0 - for vv in nn.iter(): - # check the format of coords - if vv.tag==link+'Coords': - coords=bool(vv.attrib) - if coords: - p_h=vv.attrib['points'].split(' ') - c_t_in.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - break - else: - pass - - - if vv.tag==link+'Point': - c_t_in.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - elif vv.tag!=link+'Point' and sumi>=1: - break - co_table.append(np.array(c_t_in)) - co_table_text.append(' ') - - elif tag.endswith('}NoiseRegion') or tag.endswith('}noiseregion'): - for nn in root1.iter(tag): - c_t_in=[] - sumi=0 - for vv in nn.iter(): - # check the format of coords - if vv.tag==link+'Coords': - coords=bool(vv.attrib) - if coords: - p_h=vv.attrib['points'].split(' ') - c_t_in.append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - break - else: - pass - - - if vv.tag==link+'Point': - c_t_in.append([ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ]) - sumi+=1 - - elif vv.tag!=link+'Point' and sumi>=1: - break - co_noise.append(np.array(c_t_in)) - co_noise_text.append(' ') - - img = np.zeros( (y_len,x_len,3) ) - img_poly=cv2.fillPoly(img, pts =co_text_paragraph, color=(1,1,1)) - - img_poly=cv2.fillPoly(img, pts =co_text_heading, color=(2,2,2)) - img_poly=cv2.fillPoly(img, pts =co_text_header, color=(2,2,2)) - img_poly=cv2.fillPoly(img, pts =co_text_marginalia, color=(3,3,3)) - img_poly=cv2.fillPoly(img, pts =co_img, color=(4,4,4)) - img_poly=cv2.fillPoly(img, pts =co_sep, color=(5,5,5)) - - return tree1, root1, bb_coord_printspace, id_paragraph, id_header+id_heading, co_text_paragraph, co_text_header+co_text_heading,\ - tot_region_ref,x_len, y_len,index_tot_regions, img_poly - - def return_indexes_of_contours_loctaed_inside_another_list_of_contours(self, contours, contours_loc, cx_main_loc, cy_main_loc, indexes_loc): - indexes_of_located_cont = [] - center_x_coordinates_of_located = [] - center_y_coordinates_of_located = [] - #M_main_tot = [cv2.moments(contours_loc[j]) - #for j in range(len(contours_loc))] - #cx_main_loc = [(M_main_tot[j]["m10"] / (M_main_tot[j]["m00"] + 1e-32)) for j in range(len(M_main_tot))] - #cy_main_loc = [(M_main_tot[j]["m01"] / (M_main_tot[j]["m00"] + 1e-32)) for j in range(len(M_main_tot))] - - for ij in range(len(contours)): - results = [cv2.pointPolygonTest(contours[ij], (cx_main_loc[ind], cy_main_loc[ind]), False) - for ind in range(len(cy_main_loc)) ] - results = np.array(results) - indexes_in = np.where((results == 0) | (results == 1)) - indexes = indexes_loc[indexes_in]# [(results == 0) | (results == 1)]#np.where((results == 0) | (results == 1)) - - indexes_of_located_cont.append(indexes) - center_x_coordinates_of_located.append(np.array(cx_main_loc)[indexes_in] ) - center_y_coordinates_of_located.append(np.array(cy_main_loc)[indexes_in] ) - - return indexes_of_located_cont, center_x_coordinates_of_located, center_y_coordinates_of_located - - def do_order_of_regions_with_model(self, contours_only_text_parent, contours_only_text_parent_h, text_regions_p): - height1 =672#448 - width1 = 448#224 - - height2 =672#448 - width2= 448#224 - - height3 =672#448 - width3 = 448#224 - - inference_bs = 3 - - ver_kernel = np.ones((5, 1), dtype=np.uint8) - hor_kernel = np.ones((1, 5), dtype=np.uint8) - - - min_cont_size_to_be_dilated = 10 - if len(contours_only_text_parent)>min_cont_size_to_be_dilated: - cx_conts, cy_conts, x_min_conts, x_max_conts, y_min_conts, y_max_conts, _ = find_new_features_of_contours(contours_only_text_parent) - args_cont_located = np.array(range(len(contours_only_text_parent))) - - diff_y_conts = np.abs(y_max_conts[:]-y_min_conts) - diff_x_conts = np.abs(x_max_conts[:]-x_min_conts) - - mean_x = statistics.mean(diff_x_conts) - median_x = statistics.median(diff_x_conts) - - - diff_x_ratio= diff_x_conts/mean_x - - args_cont_located_excluded = args_cont_located[diff_x_ratio>=1.3] - args_cont_located_included = args_cont_located[diff_x_ratio<1.3] - - contours_only_text_parent_excluded = [contours_only_text_parent[ind] for ind in range(len(contours_only_text_parent)) if diff_x_ratio[ind]>=1.3]#contours_only_text_parent[diff_x_ratio>=1.3] - contours_only_text_parent_included = [contours_only_text_parent[ind] for ind in range(len(contours_only_text_parent)) if diff_x_ratio[ind]<1.3]#contours_only_text_parent[diff_x_ratio<1.3] - - - cx_conts_excluded = [cx_conts[ind] for ind in range(len(cx_conts)) if diff_x_ratio[ind]>=1.3]#cx_conts[diff_x_ratio>=1.3] - cx_conts_included = [cx_conts[ind] for ind in range(len(cx_conts)) if diff_x_ratio[ind]<1.3]#cx_conts[diff_x_ratio<1.3] - - cy_conts_excluded = [cy_conts[ind] for ind in range(len(cy_conts)) if diff_x_ratio[ind]>=1.3]#cy_conts[diff_x_ratio>=1.3] - cy_conts_included = [cy_conts[ind] for ind in range(len(cy_conts)) if diff_x_ratio[ind]<1.3]#cy_conts[diff_x_ratio<1.3] - - #print(diff_x_ratio, 'ratio') - text_regions_p = text_regions_p.astype('uint8') - - if len(contours_only_text_parent_excluded)>0: - textregion_par = np.zeros((text_regions_p.shape[0], text_regions_p.shape[1])).astype('uint8') - textregion_par = cv2.fillPoly(textregion_par, pts=contours_only_text_parent_included, color=(1,1)) - else: - textregion_par = (text_regions_p[:,:]==1)*1 - textregion_par = textregion_par.astype('uint8') - - text_regions_p_textregions_dilated = cv2.erode(textregion_par , hor_kernel, iterations=2) - text_regions_p_textregions_dilated = cv2.dilate(text_regions_p_textregions_dilated , ver_kernel, iterations=4) - text_regions_p_textregions_dilated = cv2.erode(text_regions_p_textregions_dilated , hor_kernel, iterations=1) - text_regions_p_textregions_dilated = cv2.dilate(text_regions_p_textregions_dilated , ver_kernel, iterations=5) - text_regions_p_textregions_dilated[text_regions_p[:,:]>1] = 0 - - - contours_only_dilated, hir_on_text_dilated = return_contours_of_image(text_regions_p_textregions_dilated) - contours_only_dilated = return_parent_contours(contours_only_dilated, hir_on_text_dilated) - - indexes_of_located_cont, center_x_coordinates_of_located, center_y_coordinates_of_located = self.return_indexes_of_contours_loctaed_inside_another_list_of_contours(contours_only_dilated, contours_only_text_parent_included, cx_conts_included, cy_conts_included, args_cont_located_included) - - - if len(args_cont_located_excluded)>0: - for ind in args_cont_located_excluded: - indexes_of_located_cont.append(np.array([ind])) - contours_only_dilated.append(contours_only_text_parent[ind]) - center_y_coordinates_of_located.append(0) - - array_list = [np.array([elem]) if isinstance(elem, int) else elem for elem in indexes_of_located_cont] - flattened_array = np.concatenate([arr.ravel() for arr in array_list]) - #print(len( np.unique(flattened_array)), 'indexes_of_located_cont uniques') - - missing_textregions = list( set(np.array(range(len(contours_only_text_parent))) ) - set(np.unique(flattened_array)) ) - #print(missing_textregions, 'missing_textregions') - - for ind in missing_textregions: - indexes_of_located_cont.append(np.array([ind])) - contours_only_dilated.append(contours_only_text_parent[ind]) - center_y_coordinates_of_located.append(0) - - - if contours_only_text_parent_h: - for vi in range(len(contours_only_text_parent_h)): - indexes_of_located_cont.append(int(vi+len(contours_only_text_parent))) - - array_list = [np.array([elem]) if isinstance(elem, int) else elem for elem in indexes_of_located_cont] - flattened_array = np.concatenate([arr.ravel() for arr in array_list]) - - y_len = text_regions_p.shape[0] - x_len = text_regions_p.shape[1] - - img_poly = np.zeros((y_len,x_len), dtype='uint8') - ###img_poly[text_regions_p[:,:]==1] = 1 - ###img_poly[text_regions_p[:,:]==2] = 2 - ###img_poly[text_regions_p[:,:]==3] = 4 - ###img_poly[text_regions_p[:,:]==6] = 5 - - ##img_poly[text_regions_p[:,:]==1] = 1 - ##img_poly[text_regions_p[:,:]==2] = 2 - ##img_poly[text_regions_p[:,:]==3] = 3 - ##img_poly[text_regions_p[:,:]==4] = 4 - ##img_poly[text_regions_p[:,:]==5] = 5 - - img_poly = np.copy(text_regions_p) - - img_header_and_sep = np.zeros((y_len,x_len), dtype='uint8') - if contours_only_text_parent_h: - _, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, _ = find_new_features_of_contours( - contours_only_text_parent_h) - for j in range(len(cy_main)): - img_header_and_sep[int(y_max_main[j]):int(y_max_main[j])+12, - int(x_min_main[j]):int(x_max_main[j])] = 1 - co_text_all_org = contours_only_text_parent + contours_only_text_parent_h - if len(contours_only_text_parent)>min_cont_size_to_be_dilated: - co_text_all = contours_only_dilated + contours_only_text_parent_h - else: - co_text_all = contours_only_text_parent + contours_only_text_parent_h - else: - co_text_all_org = contours_only_text_parent - if len(contours_only_text_parent)>min_cont_size_to_be_dilated: - co_text_all = contours_only_dilated - else: - co_text_all = contours_only_text_parent - - if not len(co_text_all): - return [], [] - - labels_con = np.zeros((int(y_len /6.), int(x_len/6.), len(co_text_all)), dtype=bool) - - co_text_all = [(i/6).astype(int) for i in co_text_all] - for i in range(len(co_text_all)): - img = labels_con[:,:,i].astype(np.uint8) - - #img = cv2.resize(img, (int(img.shape[1]/6), int(img.shape[0]/6)), interpolation=cv2.INTER_NEAREST) - - cv2.fillPoly(img, pts=[co_text_all[i]], color=(1,)) - labels_con[:,:,i] = img - - - labels_con = resize_image(labels_con.astype(np.uint8), height1, width1).astype(bool) - img_header_and_sep = resize_image(img_header_and_sep, height1, width1) - img_poly = resize_image(img_poly, height3, width3) - - - - input_1 = np.zeros((inference_bs, height1, width1, 3)) - ordered = [list(range(len(co_text_all)))] - index_update = 0 - #print(labels_con.shape[2],"number of regions for reading order") - while index_update>=0: - ij_list = ordered.pop(index_update) - i = ij_list.pop(0) - - ante_list = [] - post_list = [] - tot_counter = 0 - batch = [] - for j in ij_list: - img1 = labels_con[:,:,i].astype(float) - img2 = labels_con[:,:,j].astype(float) - img1[img_poly==5] = 2 - img2[img_poly==5] = 2 - img1[img_header_and_sep==1] = 3 - img2[img_header_and_sep==1] = 3 - - input_1[len(batch), :, :, 0] = img1 / 3. - input_1[len(batch), :, :, 2] = img2 / 3. - input_1[len(batch), :, :, 1] = img_poly / 5. - - tot_counter += 1 - batch.append(j) - if tot_counter % inference_bs == 0 or tot_counter == len(ij_list): - y_pr = self.model_zoo.get('reading_order').predict(input_1, verbose=0) - for jb, j in enumerate(batch): - if y_pr[jb][0]>=0.5: - post_list.append(j) - else: - ante_list.append(j) - batch = [] - - if len(ante_list): - ordered.insert(index_update, ante_list) - index_update += 1 - ordered.insert(index_update, [i]) - if len(post_list): - ordered.insert(index_update + 1, post_list) - - index_update = -1 - for index_next, ij_list in enumerate(ordered): - if len(ij_list) > 1: - index_update = index_next - break - - ordered = [i[0] for i in ordered] - - ##id_all_text = np.array(id_all_text)[index_sort] - - - if len(contours_only_text_parent)>min_cont_size_to_be_dilated: - org_contours_indexes = [] - for ind in range(len(ordered)): - region_with_curr_order = ordered[ind] - if region_with_curr_order < len(contours_only_dilated): - if np.isscalar(indexes_of_located_cont[region_with_curr_order]): - org_contours_indexes = org_contours_indexes + [indexes_of_located_cont[region_with_curr_order]] - else: - arg_sort_located_cont = np.argsort(center_y_coordinates_of_located[region_with_curr_order]) - org_contours_indexes = org_contours_indexes + list(np.array(indexes_of_located_cont[region_with_curr_order])[arg_sort_located_cont]) ##org_contours_indexes + list ( - else: - org_contours_indexes = org_contours_indexes + [indexes_of_located_cont[region_with_curr_order]] - - region_ids = ['region_%04d' % i for i in range(len(co_text_all_org))] - return org_contours_indexes, region_ids - else: - region_ids = ['region_%04d' % i for i in range(len(co_text_all_org))] - return ordered, region_ids - - - - - def run(self, - overwrite: bool = False, - xml_filename: Optional[str] = None, - dir_in: Optional[str] = None, - dir_out: Optional[str] = None, - ): - """ - Get image and scales, then extract the page of scanned image - """ - self.logger.debug("enter run") - t0_tot = time.time() - - if dir_in: - ls_xmls = [os.path.join(dir_in, xml_filename) - for xml_filename in filter(is_xml_filename, - os.listdir(dir_in))] - elif xml_filename: - ls_xmls = [xml_filename] - else: - raise ValueError("run requires either a single image filename or a directory") - - for xml_filename in ls_xmls: - self.logger.info(xml_filename) - t0 = time.time() - - file_name = Path(xml_filename).stem - (tree_xml, root_xml, bb_coord_printspace, id_paragraph, id_header, - co_text_paragraph, co_text_header, tot_region_ref, - x_len, y_len, index_tot_regions, img_poly) = self.read_xml(xml_filename) - - id_all_text = id_paragraph + id_header - - order_text_new, id_of_texts_tot = self.do_order_of_regions_with_model(co_text_paragraph, co_text_header, img_poly[:,:,0]) - - id_all_text = np.array(id_all_text)[order_text_new] - - alltags=[elem.tag for elem in root_xml.iter()] - - - - link=alltags[0].split('}')[0]+'}' - name_space = alltags[0].split('}')[0] - name_space = name_space.split('{')[1] - - page_element = root_xml.find(link+'Page') - - - old_ro = root_xml.find(".//{*}ReadingOrder") - - if old_ro is not None: - page_element.remove(old_ro) - - #print(old_ro, 'old_ro') - ro_subelement = ET.Element('ReadingOrder') - - ro_subelement2 = ET.SubElement(ro_subelement, 'OrderedGroup') - ro_subelement2.set('id', "ro357564684568544579089") - - for index, id_text in enumerate(id_all_text): - new_element_2 = ET.SubElement(ro_subelement2, 'RegionRefIndexed') - new_element_2.set('regionRef', id_all_text[index]) - new_element_2.set('index', str(index)) - - if (link+'PrintSpace' in alltags) or (link+'Border' in alltags): - page_element.insert(1, ro_subelement) - else: - page_element.insert(0, ro_subelement) - - alltags=[elem.tag for elem in root_xml.iter()] - - ET.register_namespace("",name_space) - assert dir_out - tree_xml.write(os.path.join(dir_out, file_name+'.xml'), - xml_declaration=True, - method='xml', - encoding="utf-8", - default_namespace=None) - - #sys.exit() - diff --git a/src/eynollah/model_zoo/model_zoo.py b/src/eynollah/model_zoo/model_zoo.py index 5291bf8..74321d2 100644 --- a/src/eynollah/model_zoo/model_zoo.py +++ b/src/eynollah/model_zoo/model_zoo.py @@ -25,6 +25,7 @@ MODEL_VRAM_LIMITS = { "table": 1818, "reading_order": 632, "ocr": 2600, # 850 for bs 8 + "extract_images": 954, } class EynollahModelZoo: @@ -48,7 +49,7 @@ class EynollahModelZoo: self._overrides = [] if model_overrides: self.override_models(*model_overrides) - self._loaded: Dict[str, Union[Predictor, AnyModel]] = {} + self._loaded: Dict[str, Predictor] = {} @property def model_overrides(self): @@ -162,7 +163,7 @@ class EynollahModelZoo: model._name = model_category return model - def get(self, model_category: str) -> Union[Predictor, AnyModel]: + def get(self, model_category: str) -> Predictor: if model_category not in self._loaded: raise ValueError(f'Model "{model_category}" not previously loaded with "load_model(..)"') return self._loaded[model_category] @@ -245,7 +246,7 @@ class EynollahModelZoo: self.logger.warning("no GPU device available") return device0 - def _load_keras_model(self, model_category, model_path, device=''): + def _load_keras_model(self, model_category, model_path, device='') -> AnyModel: os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 from ocrd_utils import tf_disable_interactive_logs tf_disable_interactive_logs() @@ -279,7 +280,7 @@ class EynollahModelZoo: return model - def _load_serving_model(self, model_category, model_path, device=''): + def _load_serving_model(self, model_category, model_path, device='') -> AnyModel: from ocrd_utils import tf_disable_interactive_logs tf_disable_interactive_logs() import tensorflow as tf @@ -303,7 +304,7 @@ class EynollahModelZoo: return model - def _load_onnx_model(self, model_category, model_path, device=''): + def _load_onnx_model(self, model_category, model_path, device='') -> AnyModel: import onnxruntime as ort import numpy as np from ocrd_utils import config diff --git a/src/eynollah/plot.py b/src/eynollah/plot.py index 608ca4f..00e15c1 100644 --- a/src/eynollah/plot.py +++ b/src/eynollah/plot.py @@ -1,10 +1,11 @@ +from typing import Optional +import os.path try: import matplotlib.pyplot as plt import matplotlib.patches as mpatches except ImportError: plt = mpatches = None import numpy as np -import os.path import cv2 from scipy.ndimage import gaussian_filter1d @@ -21,11 +22,11 @@ class EynollahPlotter: self, *, dir_out, - dir_of_all, - dir_save_page, - dir_of_deskewed, - dir_of_layout, - dir_of_cropped_images, + dir_of_all: Optional[str] = None, + dir_save_page: Optional[str] = None, + dir_of_deskewed: Optional[str] = None, + dir_of_layout: Optional[str] = None, + dir_of_cropped_images: Optional[str] = None, ): self.dir_out = dir_out self.dir_of_all = dir_of_all @@ -172,6 +173,8 @@ class EynollahPlotter: x, y, w, h = cv2.boundingRect(cont_ind) box = [x, y, w, h] image, _ = crop_image_inside_box(box, image_page) + if not image.size: + continue image = resize_image(image, int(image.shape[0] / scale_y), int(image.shape[1] / scale_x)) diff --git a/src/eynollah/processor.py b/src/eynollah/processor.py index 264d364..0c5ad1d 100644 --- a/src/eynollah/processor.py +++ b/src/eynollah/processor.py @@ -1,11 +1,12 @@ from functools import cached_property from typing import Optional +from PIL import Image from ocrd_models import OcrdPage from ocrd import OcrdPageResultImage, Processor, OcrdPageResult from eynollah.model_zoo.model_zoo import EynollahModelZoo -from .eynollah import Eynollah, EynollahXmlWriter +from .eynollah import Eynollah class EynollahProcessor(Processor): @cached_property @@ -95,4 +96,9 @@ class EynollahProcessor(Processor): img_pil=page_image, pcgts=pcgts, # ocrd.Processor will handle OCRD_EXISTING_OUTPUT more flexibly overwrite=True) + if self.parameter['binarize'] and (img_alt := next( + (img for img in pcgts.Page.AlternativeImage + if img.comments == "binarized"), None)): + result.images.append(OcrdPageResultImage( + Image.fromarray(img_alt.filename), '.IMG-BIN', img_alt)) return result diff --git a/src/eynollah/reorder.py b/src/eynollah/reorder.py new file mode 100644 index 0000000..7affd46 --- /dev/null +++ b/src/eynollah/reorder.py @@ -0,0 +1,359 @@ +""" +Machine learning based reading order detection +""" + +# pyright: reportCallIssue=false +# pyright: reportUnboundVariable=false +# pyright: reportArgumentType=false + +import logging +import os +import time +from typing import Optional +from pathlib import Path +import xml.etree.ElementTree as ET + +import cv2 +import numpy as np +import statistics + +from ocrd_utils import ( + polygon_from_points, + xywh_from_points, +) + +from .eynollah import Eynollah +from .model_zoo import EynollahModelZoo +from .utils.resize import resize_image +from .utils.contour import ( + find_new_features_of_contours, + return_contours_of_image, + return_parent_contours, +) +from .utils import is_xml_filename + +DPI_THRESHOLD = 298 +KERNEL = np.ones((5, 5), np.uint8) + + +class Reorder(Eynollah): + def __init__( + self, + *, + model_zoo: EynollahModelZoo, + logger : Optional[logging.Logger] = None, + device: str = '', + model_based: bool = True, + # also expose these on CLI? + ignore_page_extraction: bool = False, + right2left : bool = False, + input_binary: bool = False, + num_col_upper: int = 0, + num_col_lower: int = 0, + ): + self.logger = logger or logging.getLogger('eynollah.reorder') + self.model_based = model_based + self.ignore_page_extraction = ignore_page_extraction + self.tables = True # for find_num_col + self.right2left = right2left + self.input_binary = input_binary + self.num_col_upper = num_col_upper + self.num_col_lower = num_col_lower + self.model_zoo = model_zoo + self.setup_models(device=device) + + def setup_models(self, device=''): + if self.model_based: + loadable = ['reading_order'] + else: + loadable = ['page', 'col_classifier'] + if self.input_binary: + loadable.append('binarization') + self.model_zoo.load_models(*loadable, device=device) + for model in loadable: + self.logger.debug("model %s has input shape %s", model, + self.model_zoo.get(model).input_shape) + + def read_xml(self, + xml_file: str, + label_text=1, + label_head=2, + label_imgs=5, + label_seps=6, + label_marg=8, + label_drop=4, + label_tabs=10, + ): + tree1 = ET.parse(xml_file, parser=ET.XMLParser(encoding='utf-8')) + root1 = tree1.getroot() + alltags=[elem.tag for elem in root1.iter()] + link=alltags[0].split('}')[0]+'}' + + index_tot_regions = [] + tot_region_ref = [] + + page = root1.find(link+'Page') + height = int(page.get('imageHeight', 0)) + width = int(page.get('imageWidth', 0)) + skew = -float(page.get('orientation', 0)) + img_filename = page.get('imageFilename', '') + + for jj in root1.iter(link+'RegionRefIndexed'): + index_tot_regions.append(jj.attrib['index']) + tot_region_ref.append(jj.attrib['regionRef']) + + bb_coord_printspace = None + if (link+'PrintSpace' in alltags or + link+'Border' in alltags): + tag_printspace = next(x for x in alltags + if x.endswith(('Border', 'PrintSpace'))) + nn = page.find(tag_printspace) + coords = nn.find(link + 'Coords') + if points := coords.attrib.get('points'): + xywh = xywh_from_points(points) + bb_coord_printspace = [xywh['x'], + xywh['y'], + xywh['w'], + xywh['h']] + + seps_cont = [] + imgs_cont = [] + tabs_cont = [] + text_para_cont = [] + text_para_ids = [] + text_drop_cont = [] + text_drop_ids = [] + text_head_cont = [] + text_head_ids = [] + text_marg_cont = [] + text_marg_ids = [] + + for nn in root1.iter(): + if not nn.tag.endswith('Region'): + continue + if (coords := nn.find(link + 'Coords')) is None: + continue + if (points := coords.attrib.get('points')) is None: + continue + if (id_ := nn.get('id')) is None: + continue + poly = polygon_from_points(points) + cont = np.array(poly, dtype=int)[:, np.newaxis] + if nn.tag.endswith('}TextRegion'): + type_ = nn.get('type', '') + if type_ == 'drop-capital': + text_drop_cont.append(cont) + text_drop_ids.append(id_) + elif type_ == 'heading': + text_head_cont.append(cont) + text_head_ids.append(id_) + elif type_ == 'header': + # FIXME: do not keep that mapping + text_head_cont.append(cont) + text_head_ids.append(id_) + elif type_ == 'marginalia': + text_marg_cont.append(cont) + text_marg_ids.append(id_) + else: + text_para_cont.append(cont) + text_para_ids.append(id_) + + elif nn.tag.endswith('}TableRegion'): + tabs_cont.append(cont) + + elif nn.tag.endswith('}GraphicRegion'): + imgs_cont.append(cont) + + elif nn.tag.endswith('}ImageRegion'): + imgs_cont.append(cont) + + elif nn.tag.endswith('}SeparatorRegion'): + seps_cont.append(cont) + + img = np.zeros((height, width), dtype=np.uint8) + img = cv2.fillPoly(img, pts=text_para_cont, color=label_text) + img = cv2.fillPoly(img, pts=text_head_cont, color=label_head) + img = cv2.fillPoly(img, pts=text_marg_cont, color=label_marg) + img = cv2.fillPoly(img, pts=text_drop_cont, color=label_drop) + img = cv2.fillPoly(img, pts=tabs_cont, color=label_tabs) + img = cv2.fillPoly(img, pts=imgs_cont, color=label_imgs) + img = cv2.fillPoly(img, pts=seps_cont, color=label_seps) + + return (tree1, root1, + bb_coord_printspace, + text_para_ids, text_head_ids, text_drop_ids, + text_para_cont, text_head_cont, text_drop_cont, + tot_region_ref, + width, height, skew, img_filename, + index_tot_regions, + img) + + def run(self, + overwrite: bool = False, + xml_filename: Optional[str] = None, + dir_in: Optional[str] = None, + dir_imgs: Optional[str] = None, + dir_out: Optional[str] = None, + ): + """ + Get image and scales, then extract the page of scanned image + """ + self.logger.debug("enter run") + + if dir_in: + t0_tot = time.time() + ls_xmls = [os.path.join(dir_in, xml_filename) + for xml_filename in filter(is_xml_filename, + os.listdir(dir_in))] + elif xml_filename: + ls_xmls = [xml_filename] + else: + raise ValueError("run requires either a single image filename or a directory") + + for xml_filename in ls_xmls: + self.run_single(xml_filename, + dir_out=dir_out, + dir_imgs=dir_imgs, + overwrite=overwrite) + + if dir_in: + self.logger.info("All jobs done in %.1fs", time.time() - t0_tot) + + def run_single(self, + xml_filename: str, + dir_imgs: Optional[str] = None, + dir_out: Optional[str] = None, + overwrite: bool = False, + label_imgs=5, + label_seps=6, + ) -> None: + self.logger.info(xml_filename) + t0 = time.time() + + file_name = Path(xml_filename).stem + (tree_xml, root_xml, + _, # FIXME: crop img_poly and contours (bb_coord_printspace) + para_ids, head_ids, drop_ids, + para_cont, head_cont, drop_cont, + _, # FIXME: do not ignore existing RO (tot_region_ref) + width, height, skew, img_filename, + _, # FIXME: do not ignore existing RO (index_tot_regions) + region_labels) = self.read_xml(xml_filename) + + all_text_ids = np.array(para_ids + head_ids + drop_ids) + + self.logger.debug("ordering %d paragraphs, %d headings and %d drop-capitals", + len(para_ids), len(head_ids), len(drop_ids)) + if self.model_based: + order_text = self.do_order_of_regions_with_model( + para_cont, + head_cont, + drop_cont, + region_labels) + else: + if img_filename and ( + os.path.exists(img_path := img_filename) or + os.path.exists(img_path := os.path.join('..', img_filename)) or + dir_imgs and + os.path.exists(img_path := os.path.join(dir_imgs, img_filename)) or + dir_imgs and + os.path.exists(img_path := os.path.join(dir_imgs, os.path.basename(img_filename)))): + img_filename = img_path + else: + xml_basename = Path(xml_filename).with_suffix('') + def try_suffixes(basename, suffixes): + for suf in suffixes: + if (filename := basename.with_suffix(suf)).exists(): + yield filename + elif dir_imgs and (filename := (dir_imgs / basename).with_suffix(suf)).exists(): + yield filename + img_filename = next(try_suffixes(xml_basename, + ['.tif', '.TIF', + '.jpg', '.JPG', + '.png', '.PNG', + '.jpeg', '.gif'])) + # load and analyse image + image = self.cache_images(image_filename=img_filename) + _, num_col, _ = self.resize_and_enhance_image_with_column_classifier(image) + if image['img_res'].shape[:2] != region_labels.shape: + # image was resized by col-classifier + # so bring label map to same size + # (order rules have similar expectations) + region_labels = resize_image(region_labels, *image['img_res'].shape[:2]) + scale_factor = np.array([[image['scale_x'], image['scale_y']]]) + para_cont = [(cont * scale_factor).astype(int) for cont in para_cont] + head_cont = [(cont * scale_factor).astype(int) for cont in head_cont] + drop_cont = [(cont * scale_factor).astype(int) for cont in drop_cont] + + # in Eynollah: regions_without_separators + nonsep_labels = np.copy(region_labels) + nonsep_labels[label_seps] = 0 + nonsep_labels[label_imgs] = 0 + + # deskew + if np.abs(skew) >= 0.13: # in Eynollah: SLOPE_THRESHOLD + _, region_labels, nonsep_labels = self.get_deskewed_masks( + skew, np.zeros((1, 1)), region_labels, nonsep_labels) + # also deskew contours + # (directly instead of match_deskewed_contours) + # rotate_image() does not enlarge canvas, + # so our calculation must compensate + h_o, w_o = image['img_res'].shape[:2] + M = cv2.getRotationMatrix2D((0.5 * w_o, 0.5 * h_o), -skew, 1.0)[:2, :2] + cos = np.abs(M[0, 0]) + sin = np.abs(M[0, 1]) + off = np.array([[0.5 * (w_o * cos + h_o * sin - w_o), + 0.5 * (w_o * sin + h_o * cos - h_o)]], + dtype=int) + if skew > 0: + off[0, 1] = -off[0, 1] + else: + off[0, 0] = -off[0, 0] + para_cont = [np.dot(cont, M).astype(int) - off for cont in para_cont] + head_cont = [np.dot(cont, M).astype(int) - off for cont in head_cont] + drop_cont = [np.dot(cont, M).astype(int) - off for cont in drop_cont] + + order_text = self.do_order_of_regions_heuristic( + para_cont, + head_cont, + drop_cont, + region_labels, + nonsep_labels, + num_col, + False) # in Eynollah: erosion_hurts + + all_text_ids = all_text_ids[order_text] + + alltags=[elem.tag for elem in root_xml.iter()] + + link=alltags[0].split('}')[0]+'}' + ET.register_namespace("", link[1:-1]) + + page = root_xml.find(link+'Page') + ro_old = page.find(link+'ReadingOrder') + if ro_old is not None: + page.remove(ro_old) + + ro_new = ET.Element('ReadingOrder') + ro_group = ET.SubElement(ro_new, 'OrderedGroup') + ro_group.set('id', "ro357564684568544579089") + + for index, id_text in enumerate(all_text_ids): + ro_ref = ET.SubElement(ro_group, 'RegionRefIndexed') + ro_ref.set('regionRef', id_text) + ro_ref.set('index', str(index)) + + pos = len(page.findall(link+'AlternativeImage') + + page.findall(link+'Border') + + page.findall(link+'PrintSpace')) + page.insert(pos, ro_new) + + output_filename = os.path.join(dir_out or "", file_name + '.xml') + self.logger.info("output filename: '%s'", output_filename) + tree_xml.write(output_filename, + xml_declaration=True, + method='xml', + encoding="utf-8", + default_namespace=None) + self.logger.info("Job done in %.1fs", time.time() - t0) + diff --git a/src/eynollah/training/reload-models-v0.8.mk b/src/eynollah/training/reload-models-v0.8.mk index eab5ffb..54ea7a7 100644 --- a/src/eynollah/training/reload-models-v0.8.mk +++ b/src/eynollah/training/reload-models-v0.8.mk @@ -22,6 +22,7 @@ CURRENT_MODELS += model_eynollah_ocr_cnnrnn_20250930 CURRENT_MODELS += eynollah-binarization_20210425 CURRENT_MODELS += eynollah-column-classifier_20210425 CURRENT_MODELS += eynollah-enhancement_20210425 +CURRENT_MODELS += eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 # tf (SavedModel format) for training # onnx conversion for fast inference diff --git a/src/eynollah/utils/__init__.py b/src/eynollah/utils/__init__.py index 15e8110..cdf0a47 100644 --- a/src/eynollah/utils/__init__.py +++ b/src/eynollah/utils/__init__.py @@ -2,7 +2,8 @@ from typing import Iterable, List, Tuple from logging import getLogger import time import math -from itertools import islice +from dataclasses import dataclass, field +from itertools import islice, compress try: import matplotlib.pyplot as plt @@ -10,16 +11,18 @@ try: except ImportError: plt = None import numpy as np -from shapely import geometry +from shapely import geometry, prepared, ops import cv2 from scipy.signal import find_peaks from scipy.ndimage import gaussian_filter1d from skimage import morphology from .is_nan import isNaN -from .contour import (contours_in_same_horizon, +from .contour import (contour2polygon, + contours_in_same_horizon, find_center_of_contours, find_new_features_of_contours, + polygon2contour, return_contours_of_image, return_parent_contours) @@ -39,6 +42,32 @@ def batched(iterable, n): while batch := tuple(islice(iterator, n)): yield batch +def itemgetter(seq): + # replacement for operator.itemgetter + # (which is ambiguous about its return type: + # single-item if 1 argument, tuple otherwise) + def fun(obj): + return list(obj[idx] for idx in seq) + return fun + +@dataclass +class Region: + contour: np.ndarray + area: int = 0 + cx: float = 0 + cy: float = 0 + skew: float = 0 + conf: float = 0 + def __post_init__(self): + self.area = cv2.contourArea(self.contour) + moments = cv2.moments(self.contour) + self.cx = moments['m10'] / (moments['m00'] or 1e-32) + self.cy = moments['m01'] / (moments['m00'] or 1e-32) + +@dataclass +class TextRegion(Region): + lines: List[Region] = field(default_factory=list) + def return_multicol_separators_x_start_end( regions_without_separators, peak_points, top, bot, x_min_hor_some, x_max_hor_some, cy_hor_some, y_min_hor_some, y_max_hor_some): @@ -800,93 +829,16 @@ def fill_bb_of_drop_capitals( return full_prediction == label_drop_fl_model -def check_any_text_region_in_model_one_is_main_or_header( - regions_model_1, regions_model_full, - contours_only_text_parent, - all_box_coord, all_found_textline_polygons, - slopes, - contours_only_text_parent_d_ordered, conf_contours): - - cx_main, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, y_corr_x_min_from_argmin = \ - find_new_features_of_contours(contours_only_text_parent) - - length_con=x_max_main-x_min_main - height_con=y_max_main-y_min_main - - all_found_textline_polygons_main=[] - all_found_textline_polygons_head=[] - - all_box_coord_main=[] - all_box_coord_head=[] - - slopes_main=[] - slopes_head=[] - - contours_only_text_parent_main=[] - contours_only_text_parent_head=[] - - conf_contours_main=[] - conf_contours_head=[] - - contours_only_text_parent_main_d=[] - contours_only_text_parent_head_d=[] - - for ii, con in enumerate(contours_only_text_parent): - img = np.zeros(regions_model_1.shape[:2]) - img = cv2.fillPoly(img, pts=[con], color=255) - - all_pixels=((img == 255)*1).sum() - pixels_header=( ( (img == 255) & (regions_model_full[:,:,0]==2) )*1 ).sum() - pixels_main=all_pixels-pixels_header - - if (pixels_header>=pixels_main) and ( (length_con[ii]/float(height_con[ii]) )>=1.3 ): - regions_model_1[:,:][(regions_model_1[:,:]==1) & (img == 255) ]=2 - contours_only_text_parent_head.append(con) - if len(contours_only_text_parent_d_ordered): - contours_only_text_parent_head_d.append(contours_only_text_parent_d_ordered[ii]) - all_box_coord_head.append(all_box_coord[ii]) - slopes_head.append(slopes[ii]) - all_found_textline_polygons_head.append(all_found_textline_polygons[ii]) - conf_contours_head.append(None) - else: - regions_model_1[:,:][(regions_model_1[:,:]==1) & (img == 255) ]=1 - contours_only_text_parent_main.append(con) - conf_contours_main.append(conf_contours[ii]) - if len(contours_only_text_parent_d_ordered): - contours_only_text_parent_main_d.append(contours_only_text_parent_d_ordered[ii]) - all_box_coord_main.append(all_box_coord[ii]) - slopes_main.append(slopes[ii]) - all_found_textline_polygons_main.append(all_found_textline_polygons[ii]) - - #print(all_pixels,pixels_main,pixels_header) - - return (regions_model_1, - contours_only_text_parent_main, - contours_only_text_parent_head, - all_box_coord_main, - all_box_coord_head, - all_found_textline_polygons_main, - all_found_textline_polygons_head, - slopes_main, - slopes_head, - contours_only_text_parent_main_d, - contours_only_text_parent_head_d, - conf_contours_main, - conf_contours_head) - def split_textregion_main_vs_head( - regions_model_1, - regions_model_full, - polygons_of_textregions, - polygons_of_textregions_d, - all_found_textline_polygons, - slopes, - conf_textregions, + regions_model_1: np.ndarray, + regions_model_full: np.ndarray, + textregions: List[Region], + textregions_d: List[Region], label_text=1, label_head_full=2, label_head_final=2, label_main_final=1, -): +) -> Tuple[np.ndarray, List[Region], List[Region], List[Region], List[Region]]: ### to make it faster h_o = regions_model_1.shape[0] @@ -900,19 +852,12 @@ def split_textregion_main_vs_head( (regions_model_full.shape[1] // zoom, regions_model_full.shape[0] // zoom), interpolation=cv2.INTER_NEAREST) - contours_z = [contour // zoom - for contour in polygons_of_textregions] + contours_z = [textregion.contour // zoom + for textregion in textregions] - ### - _, _, x_min_main, x_max_main, y_min_main, y_max_main, _ = \ - find_new_features_of_contours(contours_z) - - length_con=x_max_main-x_min_main - height_con=y_max_main-y_min_main - - main = [] - head = [] + main = np.ones(len(contours_z), dtype=bool) for ii, con in enumerate(contours_z): + width, height = cv2.boundingRect(con)[2:] parent = np.zeros_like(regions_model_1) parent = cv2.fillPoly(parent, pts=[con], color=1) @@ -920,16 +865,14 @@ def split_textregion_main_vs_head( pixels_main = parent.sum() - pixels_head if (( pixels_head >= 0.6 * pixels_main and - length_con[ii] >= 1.3 * height_con[ii] and - length_con[ii] <= 3 * height_con[ii] ) or + width >= 1.3 * height and + width <= 3 * height ) or ( pixels_head >= 0.3 * pixels_main and - length_con[ii] >= 3 * height_con[ii] )): + width >= 3 * height )): - head.append(ii) + main[ii] = False label = label_head_final - else: - main.append(ii) label = label_main_final regions_model_1[(regions_model_1 == label_text) & (parent > 0)] = label @@ -941,35 +884,26 @@ def split_textregion_main_vs_head( # interpolation=cv2.INTER_NEAREST) ### - def select(lis, indexes): - if not len(lis): - return [] - return [lis[ind] for ind in indexes] - return (regions_model_1, - select(polygons_of_textregions, main), - select(polygons_of_textregions, head), - select(polygons_of_textregions_d, main), - select(polygons_of_textregions_d, head), - select(all_found_textline_polygons, main), - select(all_found_textline_polygons, head), - select(slopes, main), - select(slopes, head), - select(conf_textregions, main), - select(conf_textregions, head), + list(compress(textregions, main)), + list(compress(textregions, ~main)), + list(compress(textregions_d, main)), + list(compress(textregions_d, ~main)), ) -def small_textlines_to_parent_adherence2(textlines_con, textline_mask, num_col): +def small_textlines_to_parent_adherence2( + textregions: List[TextRegion], + area_factor: float, + num_col: int, +) -> None: """ - for each region, split up textlines into small and large areas; + for each region, split up textlines into small and large; keep only the ones with large area, but expanded (by merging contours) by all intersecting lines with small area """ textlines_con_new = [] - for region in textlines_con: - areas_cnt_text = np.array(list(map(cv2.contourArea, region))) - areas_cnt_text = areas_cnt_text / float(textline_mask.size) - indexes_textlines = np.arange(len(region)) + for textregion in textregions: + areas = np.array([line.area for line in textregion.lines]) * area_factor if num_col == 0: min_area = 0.0004 @@ -977,52 +911,43 @@ def small_textlines_to_parent_adherence2(textlines_con, textline_mask, num_col): min_area = 0.0003 else: min_area = 0.0001 - indexes_textlines_small = indexes_textlines[areas_cnt_text < min_area] - indexes_textlines_large = indexes_textlines[areas_cnt_text >= min_area] + textlines_small = list(compress(textregion.lines, areas < min_area)) + textlines_large = list(compress(textregion.lines, areas >= min_area)) + textlines_small_poly = [contour2polygon(line.contour) for line in textlines_small] + textlines_large_poly = [contour2polygon(line.contour) for line in textlines_large] + textregion.lines = textlines_large - textlines_small = [region[i] for i in indexes_textlines_small] - textlines_large = [region[i] for i in indexes_textlines_large] - - img_small = np.zeros_like(textline_mask) - img_small = cv2.fillPoly(img_small, pts=textlines_small, color=1) - img_large = np.zeros_like(textline_mask) - img_large = cv2.fillPoly(img_large, pts=textlines_large, color=1) - img_inter = img_small + img_large == 2 - if np.any(img_inter): - indexes_textlines_inter = [] - for contour_small in textlines_small: - intersections = [] - for contour_large in textlines_large: - img0_small = np.zeros_like(textline_mask) - img0_small = cv2.fillPoly(img0_small, pts=[contour_small], color=1) - img0_large = np.zeros_like(textline_mask) - img0_large = cv2.fillPoly(img0_large, pts=[contour_large], color=1) - img0_inter = img0_small + img0_large == 2 - intersections.append(np.count_nonzero(img0_inter)) + if geometry.MultiPolygon(textlines_small_poly).intersects( + geometry.MultiPolygon(textlines_large_poly)): + # FIXME: also consider confidence (less certain lines replaced by better ones)... + textlines_large_prep = [prepared.prep(poly) for poly in textlines_large_poly] + textlines_small_indexes_interlarge = [] + for small_poly in textlines_small_poly: + intersections = [small_poly.intersection(prep.context).area + if prep.intersects(small_poly) else 0 + for prep in textlines_large_prep] idx_large = np.argmax(intersections) - if intersections[idx_large] <= 0: + if intersections[idx_large] == 0: idx_large = -1 - indexes_textlines_inter.append(idx_large) - - indexes_textlines_inter = np.array(indexes_textlines_inter) - for idx_large in set(indexes_textlines_inter): + textlines_small_indexes_interlarge.append(idx_large) + textlines_small_indexes_interlarge = np.array(textlines_small_indexes_interlarge) + for idx_large in set(textlines_small_indexes_interlarge): if idx_large < 0: continue - img0_union = np.zeros_like(textline_mask) - img0_union = cv2.fillPoly(img0_union, pts=[textlines_large[idx_large]], color=255) - indexes_inter_small = np.flatnonzero(indexes_textlines_inter == idx_large) - for idx_small in indexes_inter_small: - img0_union = cv2.fillPoly(img0_union, pts=[textlines_small[idx_small]], color=255) - - _, thresh = cv2.threshold(img0_union, 0, 255, 0) - contours_union, _ = cv2.findContours(thresh.astype(np.uint8), - cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - areas_union = np.array(list(map(cv2.contourArea, contours_union))) - contour_union = contours_union[np.argmax(areas_union)] #contours_union[0] - textlines_large[idx_large] = contour_union - - textlines_con_new.append(textlines_large) - return textlines_con_new + large_poly = textlines_large_poly[idx_large] + indexes_small = np.flatnonzero(textlines_small_indexes_interlarge == idx_large) + for idx_small in indexes_small: + large_poly = large_poly.union(textlines_small_poly[idx_small]) + if large_poly.geom_type == 'GeometryCollection': + # hetergeneous: filter lines and points + large_poly = ops.unary_union([geom for geom in large_poly.geoms + if geom.area > 0]) + if large_poly.geom_type == 'MultiPolygon': + # disjoint: pick largest part + idx_large = np.argmax([geom.area for geom in large_poly.geoms]) + large_poly = large_poly.geoms[idx_large] + # replace original + textregion.lines[idx_large].contour = polygon2contour(large_poly) def order_of_regions(textline_mask, contours_main, contours_head, contours_drop, y_ref, x_ref): """ @@ -1051,7 +976,7 @@ def order_of_regions(textline_mask, contours_main, contours_head, contours_drop, total = len(contours_main) + len(contours_head) + len(contours_drop) assert total == 0 or np.any(textline_mask) - # ax1 = plt.subplot(2, 1, 1, title="order_of_regions textline_mask") + # ax1 = plt.subplot(1, 2, 1, title="order_of_regions textline_mask") # plt.imshow(textline_mask, aspect='auto') y = textline_mask.sum(axis=1) # horizontal projection profile y_padded = np.zeros(len(y) + 40) @@ -1061,15 +986,21 @@ def order_of_regions(textline_mask, contours_main, contours_head, contours_drop, #z = gaussian_filter1d(y_padded, sigma_gaus) #peaks, _ = find_peaks(z, height=0) #peaks = peaks - 20 - # ax2 = plt.subplot(2, 1, 2, title="smoothed horizontal projection", sharex=ax1) - # plt.plot(y) + # ax2 = plt.subplot(1, 2, 2, title="smoothed horizontal projection", sharey=ax1) + # ax2plot = plt.plot(y)[0] + # xdata = ax2plot.get_xdata() + # ydata = ax2plot.get_ydata() + # ax2plot.set_xdata(ydata) + # ax2plot.set_ydata(xdata) + # ax2.set_xlim(*ax1.get_xlim()) + # ax2.set_ylim(*ax1.get_ylim()) zneg_rev = np.max(y_padded) - y_padded zneg = np.zeros(len(zneg_rev) + 40) zneg[20 : len(zneg_rev) + 20] = zneg_rev zneg = gaussian_filter1d(zneg, sigma_gaus) peaks_neg, _ = find_peaks(zneg, height=0) - # plt.vlines(peaks_neg - 40, 0, None, label="peaks") + # plt.hlines(peaks_neg - 40, 0, textline_mask.shape[1], label="peaks") # plt.show() peaks_neg = peaks_neg - 20 - 20 @@ -1288,7 +1219,7 @@ def find_number_of_columns_in_document( separator_mask: np.ndarray, num_col_classifier: int, tables: bool, - contours_h: List[np.ndarray] = None, + contours_h: List[np.ndarray] = [], logger=None ) -> Tuple[int, List[int], np.ndarray, List[int], np.ndarray]: """ @@ -1299,7 +1230,7 @@ def find_number_of_columns_in_document( * separator_mask: mask of (separator-only) region labels * num_col_classifier: predicted (expected) number of columns of the page * tables: whether tables may be present - * contours_h: polygons of potential headings (serving as additional horizontal separators) + * contours_h: contours of potential headings (serving as additional horizontal separators) * logger Returns: a tuple of @@ -1422,7 +1353,7 @@ def find_number_of_columns_in_document( matrix_of_seps_ch[len(cy_seps_hor):,8]=dist_y_ver matrix_of_seps_ch[len(cy_seps_hor):,9]=1 - if contours_h is not None: + if len(contours_h): _, dist_x_head, x_min_head, x_max_head, cy_head, _, y_min_head, y_max_head, _ = \ find_features_of_lines(contours_h) matrix_l_n = np.zeros((len(cy_head), matrix_of_seps_ch.shape[1]), dtype=int) @@ -1451,7 +1382,7 @@ def find_number_of_columns_in_document( (x_max_seps_hor >= .84 * width)] cy_seps_splitters = np.append(cy_seps_splitters, special_separators) - if contours_h is not None: + if len(contours_h): y_min_splitters_head = y_min_head[(x_min_head <= .16 * width) & (x_max_head >= .84 * width)] y_max_splitters_head = y_max_head[(x_min_head <= .16 * width) & diff --git a/src/eynollah/utils/contour.py b/src/eynollah/utils/contour.py index 83243ca..1868567 100644 --- a/src/eynollah/utils/contour.py +++ b/src/eynollah/utils/contour.py @@ -1,4 +1,4 @@ -from typing import Sequence, Union +from typing import List, Sequence, Union, Tuple from numbers import Number from functools import partial import itertools @@ -167,26 +167,34 @@ def dilate_textregion_contours(all_found_textregion_polygons): [polygon2contour(contour2polygon(contour, dilate=6)) for contour in all_found_textregion_polygons]) -def match_deskewed_contours(slope_deskew, contours_o, contours_d, shape_o, shape_d): +def match_deskewed_contours( + slope_deskew: float, + regions_o, #: List[Region], # (cyclic import) + regions_d, #: List[Region], # (cyclic import) + shape_o: Tuple[int, int], + shape_d: Tuple[int, int], +) -> np.ndarray: from . import ensure_array - cntareas_o = np.array([cv2.contourArea(contour) for contour in contours_o]) - cntareas_d = np.array([cv2.contourArea(contour) for contour in contours_d]) + centers_o = np.array([[region.cx, region.cy] for region in regions_o]) # [N, 2] + centers_d = np.array([[region.cx, region.cy] for region in regions_d]) # [N, 2] + cntareas_o = np.array([region.area for region in regions_o]) + cntareas_d = np.array([region.area for region in regions_d]) cntareas_o = cntareas_o / float(np.prod(shape_o[:2])) cntareas_d = cntareas_d / float(np.prod(shape_d[:2])) - contours_o = ensure_array(contours_o) - contours_d = ensure_array(contours_d) + contours_o = ensure_array([region.contour for region in regions_o]) + contours_d = ensure_array([region.contour for region in regions_d]) sort_o = np.argsort(cntareas_o) sort_d = np.argsort(cntareas_d) + centers_o = centers_o[sort_o].T # [2, N] + centers_d = centers_d[sort_d].T # [2, N] contours_o = contours_o[sort_o] contours_d = contours_d[sort_d] cntareas_o = cntareas_o[sort_o] cntareas_d = cntareas_d[sort_d] - centers_o = np.stack(find_center_of_contours(contours_o)) # [2, N] - centers_d = np.stack(find_center_of_contours(contours_d)) # [2, N] center0_o = centers_o[:, -1:] # [2, 1] center0_d = centers_d[:, -1:] # [2, 1] diff --git a/src/eynollah/utils/separate_lines.py b/src/eynollah/utils/separate_lines.py index 689a564..a4c528c 100644 --- a/src/eynollah/utils/separate_lines.py +++ b/src/eynollah/utils/separate_lines.py @@ -1585,6 +1585,7 @@ def do_work_of_slopes_new_curved( mask_line = np.zeros_like(mask_parent) mask_line = cv2.fillPoly(mask_line, pts=[contour], color=1) mask_line = cv2.dilate(mask_line, KERNEL, iterations=5 if num_col == 0 else 4) + mask_line *= mask_parent # plt.subplot(1, 2, 1, title="parent mask") # plt.imshow(mask_parent) # plt.subplot(1, 2, 2, title="single textline") diff --git a/src/eynollah/writer.py b/src/eynollah/writer.py index 47fc32f..4d68ae8 100644 --- a/src/eynollah/writer.py +++ b/src/eynollah/writer.py @@ -3,36 +3,47 @@ from pathlib import Path import os.path import logging -from typing import Optional +from typing import Optional, List, Tuple + import numpy as np +import cv2 from shapely import affinity, clip_by_rect from ocrd_utils import points_from_polygon from ocrd_models.ocrd_page import ( - BorderType, - CoordsType, - TextLineType, - TextEquivType, - TextRegionType, - ImageRegionType, - TableRegionType, - SeparatorRegionType, - to_xml - ) + AlternativeImageType, + BorderType, + CoordsType, + TextLineType, + TextEquivType, + TextRegionType, + ImageRegionType, + TableRegionType, + SeparatorRegionType, + PcGtsType, + to_xml +) +from .utils import Region, TextRegion from .utils.xml import create_page_xml, xml_reading_order from .utils.counter import EynollahIdCounter from .utils.contour import contour2polygon, make_valid class EynollahXmlWriter: - def __init__(self, *, dir_out, image_filename, image_width, image_height, curved_line, pcgts=None): + def __init__( + self, *, + dir_out: Optional[str], + image_filename: str, + image_width: int, + image_height: int, + pcgts: Optional[PcGtsType] = None, + ): self.logger = logging.getLogger('eynollah.writer') self.counter = EynollahIdCounter() self.dir_out = dir_out self.image_filename = image_filename self.output_filename = os.path.join(self.dir_out or "", self.image_filename_stem) + ".xml" - self.curved_line = curved_line self.pcgts = pcgts self.image_height = image_height self.image_width = image_width @@ -40,254 +51,159 @@ class EynollahXmlWriter: self.scale_y = 1.0 @property - def image_filename_stem(self): + def image_filename_stem(self) -> str: return Path(Path(self.image_filename).name).stem - def calculate_points(self, contour, offset=None): - poly = contour2polygon(contour) + def calculate_points( + self, contour: np.ndarray, + offset: Optional[List[int]] = None, + dilate: int = 0, + ) -> str: + poly = contour2polygon(contour, dilate=dilate) if offset is not None: poly = affinity.translate(poly, *offset) poly = affinity.scale(poly, xfact=1 / self.scale_x, yfact=1 / self.scale_y, origin=(0, 0)) poly = make_valid(clip_by_rect(poly, 0, 0, self.image_width, self.image_height)) return points_from_polygon(poly.exterior.coords[:-1]) - def serialize_lines_in_region(self, text_region, all_found_textline_polygons, region_idx, page_coord, slopes, counter, ocr_all_textlines_textregion): - for j, polygon_textline in enumerate(all_found_textline_polygons[region_idx]): - coords = CoordsType() - textline = TextLineType(id=counter.next_line_id, Coords=coords) - if ocr_all_textlines_textregion: - # FIXME: add OCR confidence - textline.set_TextEquiv([TextEquivType(Unicode=ocr_all_textlines_textregion[j])]) + def serialize_lines_in_region( + self, text_region: TextRegionType, + offset: List[int], + counter: EynollahIdCounter, + lines: List[Region], + ) -> None: + for line in lines: + textline = TextLineType( + id=counter.next_line_id, + Coords=CoordsType(points=self.calculate_points(line.contour, offset, 5), + conf=line.conf) + ) text_region.add_TextLine(textline) - text_region.set_orientation(-slopes[region_idx]) - offset = [page_coord[2], page_coord[0]] - coords.set_points(self.calculate_points(polygon_textline, offset)) def write_pagexml(self, pcgts): self.logger.info("output filename: '%s'", self.output_filename) + if img_alt := next( + (img for img in pcgts.Page.AlternativeImage + if img.comments == "binarized" + and isinstance(img.filename, np.ndarray)), None): + img_alt_filename = self.output_filename[:-4] + '.bin.png' + cv2.imwrite(img_alt_filename, img_alt.filename) + img_alt.filename = os.path.basename(img_alt_filename) with open(self.output_filename, 'w') as f: f.write(to_xml(pcgts)) - def build_pagexml_no_full_layout( + def build_pagexml( self, *, - num_col, - found_polygons_text_region, - page_coord, - page_slope, - order_of_texts, - all_found_textline_polygons, - found_polygons_images, - found_polygons_tables, - found_polygons_marginals_left, - found_polygons_marginals_right, - all_found_textline_polygons_marginals_left, - all_found_textline_polygons_marginals_right, - slopes, - slopes_marginals_left, - slopes_marginals_right, - cont_page, - polygons_seplines, - ocr_all_textlines=None, - ocr_all_textlines_marginals_left=None, - ocr_all_textlines_marginals_right=None, - ocr_all_textlines_drop=None, - conf_textregions=None, - conf_marginals_left=None, - conf_marginals_right=None, - conf_images=None, - conf_tables=None, - ): - return self.build_pagexml_full_layout( - num_col=num_col, - found_polygons_text_region=found_polygons_text_region, - found_polygons_text_region_h=[], - page_coord=page_coord, - page_slope=page_slope, - order_of_texts=order_of_texts, - all_found_textline_polygons=all_found_textline_polygons, - all_found_textline_polygons_h=[], - found_polygons_images=found_polygons_images, - found_polygons_tables=found_polygons_tables, - found_polygons_drop_capitals=[], - found_polygons_marginals_left=found_polygons_marginals_left, - found_polygons_marginals_right=found_polygons_marginals_right, - all_found_textline_polygons_marginals_left=all_found_textline_polygons_marginals_left, - all_found_textline_polygons_marginals_right=all_found_textline_polygons_marginals_right, - slopes=slopes, - slopes_h=[], - slopes_marginals_left=slopes_marginals_left, - slopes_marginals_right=slopes_marginals_right, - cont_page=cont_page, - polygons_seplines=polygons_seplines, - ocr_all_textlines=ocr_all_textlines, - ocr_all_textlines_marginals_left=ocr_all_textlines_marginals_left, - ocr_all_textlines_marginals_right=ocr_all_textlines_marginals_right, - conf_textregions=conf_textregions, - conf_marginals_left=conf_marginals_left, - conf_marginals_right=conf_marginals_right, - conf_images=conf_images, - conf_tables=conf_tables, - ) - - def build_pagexml_full_layout( - self, - *, - num_col, - found_polygons_text_region, - found_polygons_text_region_h, - page_coord, - page_slope, - order_of_texts, - all_found_textline_polygons, - all_found_textline_polygons_h, - found_polygons_images, - found_polygons_tables, - found_polygons_drop_capitals, - found_polygons_marginals_left, - found_polygons_marginals_right, - all_found_textline_polygons_marginals_left, - all_found_textline_polygons_marginals_right, - slopes, - slopes_h, - slopes_marginals_left, - slopes_marginals_right, - cont_page, - polygons_seplines, - ocr_all_textlines=None, - ocr_all_textlines_h=None, - ocr_all_textlines_marginals_left=None, - ocr_all_textlines_marginals_right=None, - ocr_all_textlines_drop=None, - conf_textregions=None, - conf_textregions_h=None, - conf_marginals_left=None, - conf_marginals_right=None, - conf_images=None, - conf_tables=None, - conf_drops=None, + page: Region, + img_bin: Optional[np.ndarray] = None, + num_col: int = 1, + order_of_texts: List[int] = [], + textregions: List[TextRegion] = [], + textregions_h: List[TextRegion] = [], + images: List[Region] = [], + tables: List[Region] = [], + drop_caps: List[Region] = [], + marginals_left: List[TextRegion] = [], + marginals_right: List[TextRegion] = [], + seplines: List[Region] = [], ): self.logger.debug('enter build_pagexml') # create the file structure pcgts = self.pcgts if self.pcgts else create_page_xml( self.image_filename, self.image_height, self.image_width) - page = pcgts.get_Page() pcgts.Metadata.Comments = "num_col %d" % num_col - page.set_custom('layout {num_col:%d;} ' % num_col) - page.set_orientation(-page_slope) - if len(cont_page): - page.set_Border(BorderType(Coords=CoordsType(points=self.calculate_points(cont_page[0])))) - - offset = [page_coord[2], page_coord[0]] + if img_bin is not None: + img_alt = AlternativeImageType(filename=img_bin, # will be replaced later + comments="binarized") + pcgts.Page.add_AlternativeImage(img_alt) + pcgts.Page.set_custom('layout {num_col:%d;} ' % num_col) + pcgts.Page.set_orientation(-page.skew) + pcgts.Page.set_Border(BorderType(Coords=CoordsType( + points=self.calculate_points(page.contour)))) + x, y, w, h = cv2.boundingRect(page.contour) + offset = [x, y] counter = EynollahIdCounter() if len(order_of_texts): _counter_marginals = EynollahIdCounter(region_idx=len(order_of_texts)) id_of_marginalia_left = [_counter_marginals.next_region_id - for _ in found_polygons_marginals_left] + for _ in marginals_left] id_of_marginalia_right = [_counter_marginals.next_region_id - for _ in found_polygons_marginals_right] - xml_reading_order(page, order_of_texts, id_of_marginalia_left, id_of_marginalia_right) + for _ in marginals_right] + xml_reading_order(pcgts.Page, order_of_texts, id_of_marginalia_left, id_of_marginalia_right) - for mm, region_contour in enumerate(found_polygons_text_region): + for region in textregions: textregion = TextRegionType( id=counter.next_region_id, type_='paragraph', - Coords=CoordsType(points=self.calculate_points(region_contour, offset)) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 6), + conf=region.conf), + orientation=-region.skew ) - assert textregion.Coords - if conf_textregions: - textregion.Coords.set_conf(conf_textregions[mm]) - page.add_TextRegion(textregion) - if ocr_all_textlines: - ocr_textlines = ocr_all_textlines[mm] - else: - ocr_textlines = None - self.serialize_lines_in_region(textregion, all_found_textline_polygons, mm, page_coord, - slopes, counter, ocr_textlines) + self.serialize_lines_in_region(textregion, offset, counter, region.lines) + pcgts.Page.add_TextRegion(textregion) - self.logger.debug('len(found_polygons_text_region_h) %s', len(found_polygons_text_region_h)) - for mm, region_contour in enumerate(found_polygons_text_region_h): + self.logger.debug('len(textregions_h) %s', len(textregions_h)) + for region in textregions_h: textregion = TextRegionType( id=counter.next_region_id, type_='heading', - Coords=CoordsType(points=self.calculate_points(region_contour, offset)) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 6), + conf=region.conf), + orientation=-region.skew ) - assert textregion.Coords - if conf_textregions_h: - textregion.Coords.set_conf(conf_textregions_h[mm]) - page.add_TextRegion(textregion) - if ocr_all_textlines_h: - ocr_textlines = ocr_all_textlines_h[mm] - else: - ocr_textlines = None - self.serialize_lines_in_region(textregion, all_found_textline_polygons_h, mm, page_coord, - slopes_h, counter, ocr_textlines) + self.serialize_lines_in_region(textregion, offset, counter, region.lines) + pcgts.Page.add_TextRegion(textregion) - for mm, region_contour in enumerate(found_polygons_drop_capitals): - dropcapital = TextRegionType( + for region in drop_caps: + textregion = TextRegionType( id=counter.next_region_id, type_='drop-capital', - Coords=CoordsType(points=self.calculate_points(region_contour, offset)) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 6), + conf=region.conf), + orientation=-region.skew ) - if conf_drops: - dropcapital.Coords.set_conf(conf_drops[mm]) - page.add_TextRegion(dropcapital) - slopes_drop = [0] - if ocr_all_textlines_drop: - ocr_textlines = ocr_all_textlines_drop[mm] - else: - ocr_textlines = None - self.serialize_lines_in_region(dropcapital, [[found_polygons_drop_capitals[mm]]], 0, page_coord, - slopes_drop, counter, ocr_textlines) + self.serialize_lines_in_region(textregion, offset, counter, [region]) + pcgts.Page.add_TextRegion(textregion) - for mm, region_contour in enumerate(found_polygons_marginals_left): - marginal = TextRegionType( + for region in marginals_left: + textregion = TextRegionType( id=counter.next_region_id, type_='marginalia', - Coords=CoordsType(points=self.calculate_points(region_contour, offset)) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 6), + conf=region.conf), + orientation=-region.skew ) - if conf_marginals_left: - marginal.Coords.set_conf(conf_marginals_left[mm]) - page.add_TextRegion(marginal) - if ocr_all_textlines_marginals_left: - ocr_textlines = ocr_all_textlines_marginals_left[mm] - else: - ocr_textlines = None - self.serialize_lines_in_region(marginal, all_found_textline_polygons_marginals_left, mm, page_coord, - slopes_marginals_left, counter, ocr_textlines) + self.serialize_lines_in_region(textregion, offset, counter, region.lines) + pcgts.Page.add_TextRegion(textregion) - for mm, region_contour in enumerate(found_polygons_marginals_right): - marginal = TextRegionType( + for region in marginals_right: + textregion = TextRegionType( id=counter.next_region_id, type_='marginalia', - Coords=CoordsType(points=self.calculate_points(region_contour, offset)) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 6), + conf=region.conf), + orientation=-region.skew ) - if conf_marginals_right: - marginal.Coords.set_conf(conf_marginals_right[mm]) - page.add_TextRegion(marginal) - if ocr_all_textlines_marginals_right: - ocr_textlines = ocr_all_textlines_marginals_right[mm] - else: - ocr_textlines = None - self.serialize_lines_in_region(marginal, all_found_textline_polygons_marginals_right, mm, page_coord, - slopes_marginals_right, counter, ocr_textlines) + self.serialize_lines_in_region(textregion, offset, counter, region.lines) + pcgts.Page.add_TextRegion(textregion) - for mm, region_contour in enumerate(found_polygons_images): + for region in images: image = ImageRegionType( id=counter.next_region_id, - Coords=CoordsType(points=self.calculate_points(region_contour, offset))) - if conf_images: - image.Coords.set_conf(conf_images[mm]) - page.add_ImageRegion(image) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 2), + conf=region.conf)) + pcgts.Page.add_ImageRegion(image) - for region_contour in polygons_seplines: - page.add_SeparatorRegion( - SeparatorRegionType(id=counter.next_region_id, - Coords=CoordsType(points=self.calculate_points(region_contour, offset)))) + for region in seplines: + pcgts.Page.add_SeparatorRegion( + SeparatorRegionType( + id=counter.next_region_id, + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 2), + conf=region.conf))) - for mm, region_contour in enumerate(found_polygons_tables): + for region in tables: table = TableRegionType( id=counter.next_region_id, - Coords=CoordsType(points=self.calculate_points(region_contour, offset))) - if conf_tables: - table.Coords.set_conf(conf_tables[mm]) - page.add_TableRegion(table) + Coords=CoordsType(points=self.calculate_points(region.contour, offset, 6), + conf=region.conf)) + pcgts.Page.add_TableRegion(table) return pcgts diff --git a/tests/cli_tests/test_mbreorder.py b/tests/cli_tests/test_mbreorder.py index e429e98..e452857 100644 --- a/tests/cli_tests/test_mbreorder.py +++ b/tests/cli_tests/test_mbreorder.py @@ -1,21 +1,30 @@ +import pytest from ocrd_modelfactory import page_from_file from ocrd_models.constants import NAMESPACES as NS +@pytest.mark.parametrize( + "options", + [ + [], # defaults + ["--model_based"], + ], ids=str) def test_run_eynollah_mbreorder_filename( tmp_path, resources_dir, run_eynollah_ok_and_check_logs, + options, ): infile = resources_dir / '2files/kant_aufklaerung_1784_0020.xml' outfile = tmp_path /'kant_aufklaerung_1784_0020.xml' run_eynollah_ok_and_check_logs( - 'machine-based-reading-order', + 'reorder', [ '-i', str(infile), + '-dim', str(resources_dir / '2files'), '-o', str(outfile.parent), - ], + ] + options, [ - # FIXME: mbreorder has no logging! + str(infile) ] ) assert outfile.exists() @@ -34,13 +43,15 @@ def test_run_eynollah_mbreorder_directory( ): outdir = tmp_path run_eynollah_ok_and_check_logs( - 'machine-based-reading-order', + 'reorder', [ '-di', str(resources_dir / '2files'), + '-dim', str(resources_dir / '2files'), '-o', str(outdir), ], [ - # FIXME: mbreorder has no logging! + 'Job done in', + 'All jobs done in', ] ) assert len(list(outdir.iterdir())) == 2 diff --git a/tests/conftest.py b/tests/conftest.py index 69f3d28..8ef14ba 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,7 +31,7 @@ def eynollah_subcommands(): 'layout', 'ocr', 'enhancement', - 'machine-based-reading-order', + 'reorder', 'models', ]