From 12e1c7aea83bedb6fb1d96d965ce6e19e60f3234 Mon Sep 17 00:00:00 2001 From: Robert Sachunsky Date: Sun, 19 Jul 2026 14:48:35 +0200 Subject: [PATCH] =?UTF-8?q?extract=5Fimages:=20make=20it=20work=20(again),?= =?UTF-8?q?=20adapt=20to=20dataclasses=E2=80=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix model loading for predictor - add `extract_images` model to ModelZoo and Makefile, too - rename `get_regions_light_v_extract_only_images` → `get_early_layout` (with similar refactoring as done in superclass previously) - separate `run_single()` from `run()` (as in superclass) - adapt to new `Region` dataclasses, too --- src/eynollah/cli/cli_extract_images.py | 2 +- src/eynollah/extract_images.py | 223 +++++++++----------- src/eynollah/model_zoo/model_zoo.py | 1 + src/eynollah/training/reload-models-v0.8.mk | 1 + 4 files changed, 107 insertions(+), 120 deletions(-) diff --git a/src/eynollah/cli/cli_extract_images.py b/src/eynollah/cli/cli_extract_images.py index acd31f1..4dbdff1 100644 --- a/src/eynollah/cli/cli_extract_images.py +++ b/src/eynollah/cli/cli_extract_images.py @@ -82,7 +82,7 @@ def extract_images_cli( ignore_page_extraction, ): """ - Detect Layout (with optional image enhancement and reading order detection) + Detect image regions only """ assert enable_plotting or not save_images, "Plotting with -si also requires -ep" assert not enable_plotting or save_images, "Plotting with -ep also requires -si" diff --git a/src/eynollah/extract_images.py b/src/eynollah/extract_images.py index 40476a3..29ec6b4 100644 --- a/src/eynollah/extract_images.py +++ b/src/eynollah/extract_images.py @@ -16,9 +16,11 @@ from eynollah.utils.contour import filter_contours_area_of_image, return_contour from eynollah.utils.resize import resize_image from .model_zoo.model_zoo import EynollahModelZoo +from .writer import EynollahXmlWriter from .eynollah import Eynollah from .utils import box2rect, is_image_filename from .plot import EynollahPlotter +from .utils import Region class EynollahImageExtractor(Eynollah): @@ -67,22 +69,29 @@ class EynollahImageExtractor(Eynollah): self.setup_models() self.logger.info(f"Model initialization complete ({time.time() - t_start:.1f}s)") - def setup_models(self): + def setup_models(self, device=''): loadable = [ "col_classifier", - "binarization", "page", "extract_images", ] - self.model_zoo.load_models(*loadable) + if self.input_binary: + loadable.append("binarization") + self.model_zoo.load_models(*loadable, device=device) - def get_regions_light_v_extract_only_images(self,img, num_col_classifier): + def get_early_layout( + self, + img, + num_col_classifier, + label_text=1, + label_imgs=2, + label_seps=3, + ): self.logger.debug("enter get_regions_extract_images_only") erosion_hurts = False - img_org = np.copy(img) - img_height_h = img_org.shape[0] - img_width_h = img_org.shape[1] + # already cropped + img_height_h, img_width_h = img.shape[:2] if num_col_classifier == 1: img_w_new = 700 @@ -98,60 +107,47 @@ class EynollahImageExtractor(Eynollah): img_w_new = 2500 else: raise ValueError("num_col_classifier must be in range 1..6") - img_h_new = int(img.shape[0] / float(img.shape[1]) * img_w_new) - img_resized = resize_image(img,img_h_new, img_w_new ) + img_h_new = img_w_new * img_height_h // img_width_h + img_resized = resize_image(img, img_h_new, img_w_new) - prediction_regions_org, _ = self.do_prediction_new_concept(True, img_resized, self.model_zoo.get("extract_images")) + prediction_regions, _ = self.do_prediction_new_concept( + True, img_resized, self.model_zoo.get("extract_images")) + prediction_regions = resize_image(prediction_regions, img_height_h, img_width_h) - prediction_regions_org = resize_image(prediction_regions_org,img_height_h, img_width_h ) - image_page, page_coord, cont_page = self.extract_page() + mask_texts_only = (prediction_regions == label_text).astype(np.uint8) + mask_images_only = (prediction_regions == label_imgs).astype(np.uint8) + mask_seps_only = (prediction_regions == label_seps).astype(np.uint8) - prediction_regions_org = prediction_regions_org[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] - prediction_regions_org=prediction_regions_org[:,:,0] + 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) - mask_seps_only = (prediction_regions_org[:,:] ==3)*1 - mask_texts_only = (prediction_regions_org[:,:] ==1)*1 - mask_images_only=(prediction_regions_org[:,:] ==2)*1 + text_regions_p = np.zeros_like(prediction_regions) + 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=texts_only_cont, color=label_text) - 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) + # rs: why? + text_regions_p[-15:] = 0 + text_regions_p[:, -15:] = 0 - polygons_of_only_texts = return_contours_of_interested_region(mask_texts_only,1,0.00001) - polygons_of_only_seps = return_contours_of_interested_region(mask_seps_only,1,0.00001) + images_cont = return_contours_of_interested_region(text_regions_p, label_imgs, 0.001) - text_regions_p_true = np.zeros(prediction_regions_org.shape) - text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts = polygons_of_only_seps, color=(3,3,3)) - - text_regions_p_true[:,:][mask_images_only[:,:] == 1] = 2 - text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts=polygons_of_only_texts, color=(1,1,1)) - - text_regions_p_true[text_regions_p_true.shape[0]-15:text_regions_p_true.shape[0], :] = 0 - text_regions_p_true[:, text_regions_p_true.shape[1]-15:text_regions_p_true.shape[1]] = 0 - - ##polygons_of_images = return_contours_of_interested_region(text_regions_p_true, 2, 0.0001) - polygons_of_images = return_contours_of_interested_region(text_regions_p_true, 2, 0.001) - - polygons_of_images_fin = [] - for ploy_img_ind in polygons_of_images: - box = _, _, w, h = cv2.boundingRect(ploy_img_ind) + images_cont_fin = [] + for cont in images_cont: + _, _, w, h = box = cv2.boundingRect(cont) if h < 150 or w < 150: pass else: - page_coord_img = box2rect(box) # type: ignore - polygons_of_images_fin.append(np.array([[page_coord_img[2], page_coord_img[0]], - [page_coord_img[3], page_coord_img[0]], - [page_coord_img[3], page_coord_img[1]], - [page_coord_img[2], page_coord_img[1]]])) + y1, y2, x1, x2 = box2rect(box) # type: ignore + images_cont_fin.append(np.array([[[x1, y1]], + [[x2, y1]], + [[x2, y2]], + [[x1, y2]]])) self.logger.debug("exit get_regions_extract_images_only") - return (text_regions_p_true, + return (text_regions_p, erosion_hurts, - polygons_seplines, - polygons_of_images_fin, - image_page, - page_coord, - cont_page) + images_cont_fin) def run(self, overwrite: bool = False, @@ -159,16 +155,12 @@ class EynollahImageExtractor(Eynollah): dir_in: Optional[str] = None, dir_out: Optional[str] = None, dir_of_cropped_images: Optional[str] = None, - dir_of_layout: Optional[str] = None, - dir_of_deskewed: Optional[str] = None, - dir_of_all: Optional[str] = None, - dir_save_page: Optional[str] = None, + **kwargs ): """ Get image and scales, then extract the page of scanned image """ self.logger.debug("enter run") - t0_tot = time.time() # Log enabled features directly enabled_modes = [] if self.full_layout: @@ -181,12 +173,12 @@ class EynollahImageExtractor(Eynollah): self.logger.info("Saving debug plots") if dir_of_cropped_images: self.logger.info(f"Saving cropped images to: {dir_of_cropped_images}") - if dir_of_layout: - self.logger.info(f"Saving layout plots to: {dir_of_layout}") - if dir_of_deskewed: - self.logger.info(f"Saving deskewed images to: {dir_of_deskewed}") - + self.plotter = EynollahPlotter( + dir_out=dir_out, + dir_of_cropped_images=dir_of_cropped_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))] @@ -196,79 +188,72 @@ class EynollahImageExtractor(Eynollah): raise ValueError("run requires either a single image filename or a directory") for img_filename in ls_imgs: - self.logger.info(img_filename) - t0 = time.time() - - self.reset_file_name_dir(img_filename, dir_out) - if self.enable_plotting: - self.plotter = EynollahPlotter(dir_out=dir_out, - dir_of_all=dir_of_all, - dir_save_page=dir_save_page, - dir_of_deskewed=dir_of_deskewed, - dir_of_cropped_images=dir_of_cropped_images, - dir_of_layout=dir_of_layout, - image_filename_stem=Path(img_filename).stem) - #print("text region early -11 in %.1fs", time.time() - t0) - if os.path.exists(self.writer.output_filename): - if overwrite: - self.logger.warning("will overwrite existing output file '%s'", self.writer.output_filename) - else: - self.logger.warning("will skip input for existing output file '%s'", self.writer.output_filename) - continue - - pcgts = self.run_single() - self.logger.info("Job done in %.1fs", time.time() - t0) - self.writer.write_pagexml(pcgts) + 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) - def run_single(self): + def run_single(self, + img_filename: str, + dir_out: Optional[str] = None, + overwrite: bool = False + ) -> 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_coord, cont_page, image_page, mask_page = self.extract_page(image) - _, _, _, 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_coord = np.array(page_coord) + images_cont = [cont - page_coord[::2][::-1][np.newaxis, np.newaxis] + 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( + num_col=num_col_classifier, + page_coord=page_coord, + page_contour=cont_page[0], + images=images, + ) + writer.write_pagexml(pcgts) + self.logger.info("Job done in %.1fs", time.time() - t0) diff --git a/src/eynollah/model_zoo/model_zoo.py b/src/eynollah/model_zoo/model_zoo.py index 5291bf8..8016457 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: 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