qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-07-26 05:29:16 +02:00
Code Issues Projects Releases Wiki Activity

extract_images: make it work (again), adapt to dataclasses…

Browse source 
- 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
This commit is contained in:
Robert Sachunsky 2026-07-19 14:48:35 +02:00
parent beb5545899
commit 12e1c7aea8
4 changed files with 107 additions and 120 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/eynollah/cli/cli_extract_images.py
View file

@ -82,7 +82,7 @@ def extract_images_cli(
ignore_page_extraction, 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 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" assert not enable_plotting or save_images, "Plotting with -ep also requires -si"

223
src/eynollah/extract_images.py
View file

@ -16,9 +16,11 @@ from eynollah.utils.contour import filter_contours_area_of_image, return_contour
from eynollah.utils.resize import resize_image from eynollah.utils.resize import resize_image
from .model_zoo.model_zoo import EynollahModelZoo from .model_zoo.model_zoo import EynollahModelZoo
from .writer import EynollahXmlWriter
from .eynollah import Eynollah from .eynollah import Eynollah
from .utils import box2rect, is_image_filename from .utils import box2rect, is_image_filename
from .plot import EynollahPlotter from .plot import EynollahPlotter
from .utils import Region
class EynollahImageExtractor(Eynollah): class EynollahImageExtractor(Eynollah):
@ -67,22 +69,29 @@ class EynollahImageExtractor(Eynollah):
self.setup_models() self.setup_models()
self.logger.info(f"Model initialization complete ({time.time() - t_start:.1f}s)") self.logger.info(f"Model initialization complete ({time.time() - t_start:.1f}s)")
def setup_models(self): def setup_models(self, device=''):
loadable = [ loadable = [
"col_classifier", "col_classifier",
"binarization",
"page", "page",
"extract_images", "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") self.logger.debug("enter get_regions_extract_images_only")
erosion_hurts = False erosion_hurts = False
img_org = np.copy(img) # already cropped
img_height_h = img_org.shape[0] img_height_h, img_width_h = img.shape[:2]
img_width_h = img_org.shape[1]
if num_col_classifier == 1: if num_col_classifier == 1:
img_w_new = 700 img_w_new = 700
@ -98,60 +107,47 @@ class EynollahImageExtractor(Eynollah):
img_w_new = 2500 img_w_new = 2500
else: else:
raise ValueError("num_col_classifier must be in range 1..6") 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_h_new = img_w_new * img_height_h // img_width_h
img_resized = resize_image(img,img_h_new, img_w_new ) 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 ) mask_texts_only = (prediction_regions == label_text).astype(np.uint8)
image_page, page_coord, cont_page = self.extract_page() 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]] texts_only_cont = return_contours_of_interested_region(mask_texts_only,1,0.00001)
prediction_regions_org=prediction_regions_org[:,:,0] seps_only_cont = return_contours_of_interested_region(mask_seps_only,1,0.00001)
mask_seps_only = (prediction_regions_org[:,:] ==3)*1 text_regions_p = np.zeros_like(prediction_regions)
mask_texts_only = (prediction_regions_org[:,:] ==1)*1 text_regions_p = cv2.fillPoly(text_regions_p, pts=seps_only_cont, color=label_seps)
mask_images_only=(prediction_regions_org[:,:] ==2)*1 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) # rs: why?
polygons_seplines = filter_contours_area_of_image( text_regions_p[-15:] = 0
mask_seps_only, polygons_seplines, hir_seplines, max_area=1, min_area=0.00001, dilate=1) text_regions_p[:, -15:] = 0
polygons_of_only_texts = return_contours_of_interested_region(mask_texts_only,1,0.00001) images_cont = return_contours_of_interested_region(text_regions_p, label_imgs, 0.001)
polygons_of_only_seps = return_contours_of_interested_region(mask_seps_only,1,0.00001)
text_regions_p_true = np.zeros(prediction_regions_org.shape) images_cont_fin = []
text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts = polygons_of_only_seps, color=(3,3,3)) for cont in images_cont:
_, _, w, h = box = cv2.boundingRect(cont)
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)
if h < 150 or w < 150: if h < 150 or w < 150:
pass pass
else: else:
page_coord_img = box2rect(box) # type: ignore y1, y2, x1, x2 = box2rect(box) # type: ignore
polygons_of_images_fin.append(np.array([[page_coord_img[2], page_coord_img[0]], images_cont_fin.append(np.array([[[x1, y1]],
[page_coord_img[3], page_coord_img[0]], [[x2, y1]],
[page_coord_img[3], page_coord_img[1]], [[x2, y2]],
[page_coord_img[2], page_coord_img[1]]])) [[x1, y2]]]))
self.logger.debug("exit get_regions_extract_images_only") self.logger.debug("exit get_regions_extract_images_only")
return (text_regions_p_true, return (text_regions_p,
erosion_hurts, erosion_hurts,
polygons_seplines, images_cont_fin)
polygons_of_images_fin,
image_page,
page_coord,
cont_page)
def run(self, def run(self,
overwrite: bool = False, overwrite: bool = False,
@ -159,16 +155,12 @@ class EynollahImageExtractor(Eynollah):
dir_in: Optional[str] = None, dir_in: Optional[str] = None,
dir_out: Optional[str] = None, dir_out: Optional[str] = None,
dir_of_cropped_images: Optional[str] = None, dir_of_cropped_images: Optional[str] = None,
dir_of_layout: Optional[str] = None, **kwargs
dir_of_deskewed: Optional[str] = None,
dir_of_all: Optional[str] = None,
dir_save_page: Optional[str] = None,
): ):
""" """
Get image and scales, then extract the page of scanned image Get image and scales, then extract the page of scanned image
""" """
self.logger.debug("enter run") self.logger.debug("enter run")
t0_tot = time.time()
# Log enabled features directly # Log enabled features directly
enabled_modes = [] enabled_modes = []
if self.full_layout: if self.full_layout:
@ -181,12 +173,12 @@ class EynollahImageExtractor(Eynollah):
self.logger.info("Saving debug plots") self.logger.info("Saving debug plots")
if dir_of_cropped_images: if dir_of_cropped_images:
self.logger.info(f"Saving cropped images to: {dir_of_cropped_images}") self.logger.info(f"Saving cropped images to: {dir_of_cropped_images}")
if dir_of_layout: self.plotter = EynollahPlotter(
self.logger.info(f"Saving layout plots to: {dir_of_layout}") dir_out=dir_out,
if dir_of_deskewed: dir_of_cropped_images=dir_of_cropped_images,
self.logger.info(f"Saving deskewed images to: {dir_of_deskewed}") )
if dir_in: if dir_in:
t0_tot = time.time()
ls_imgs = [os.path.join(dir_in, image_filename) ls_imgs = [os.path.join(dir_in, image_filename)
for image_filename in filter(is_image_filename, for image_filename in filter(is_image_filename,
os.listdir(dir_in))] os.listdir(dir_in))]
@ -196,79 +188,72 @@ class EynollahImageExtractor(Eynollah):
raise ValueError("run requires either a single image filename or a directory") raise ValueError("run requires either a single image filename or a directory")
for img_filename in ls_imgs: for img_filename in ls_imgs:
self.logger.info(img_filename) self.run_single(img_filename, dir_out=dir_out, overwrite=overwrite)
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)
if dir_in: if dir_in:
self.logger.info("All jobs done in %.1fs", time.time() - t0_tot) 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() t0 = time.time()
self.logger.info(img_filename)
self.logger.info(f"Processing file: {self.writer.image_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") 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.logger.info(f"Image: {image['img_res'].shape[1]}x{image['img_res'].shape[0]}, "
self.run_enhancement() f"scale {image['scale_x']:.1f}x{image['scale_y']:.1f}, "
f"{image['dpi']} DPI, {num_col_classifier} columns")
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"Enhancement complete ({time.time() - t0:.1f}s)") self.logger.info(f"Enhancement complete ({time.time() - t0:.1f}s)")
# Image Extraction Mode # Image Extraction Mode
self.logger.info("Step 2/5: 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, \ _, _, images_cont = self.get_early_layout(
image_page, page_coord, cont_page = \ image['img_res'], num_col_classifier)
self.get_regions_light_v_extract_only_images(img_res, num_col_classifier) self.logger.debug("Found %d images", len(images_cont))
pcgts = self.writer.build_pagexml_no_full_layout( # FIXME: post-hoc cropping (remove when models support it, and replace image['img_res'] with image_page)
found_polygons_text_region=[], page_coord = np.array(page_coord)
page_coord=page_coord, images_cont = [cont - page_coord[::2][::-1][np.newaxis, np.newaxis]
order_of_texts=[], for cont in images_cont]
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=[],
)
if self.plotter: 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") 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)

1
src/eynollah/model_zoo/model_zoo.py
View file

@ -25,6 +25,7 @@ MODEL_VRAM_LIMITS = {
"table": 1818, "table": 1818,
"reading_order": 632, "reading_order": 632,
"ocr": 2600, # 850 for bs 8 "ocr": 2600, # 850 for bs 8
"extract_images": 954,
} }
class EynollahModelZoo: class EynollahModelZoo:

1
src/eynollah/training/reload-models-v0.8.mk
View file

@ -22,6 +22,7 @@ CURRENT_MODELS += model_eynollah_ocr_cnnrnn_20250930
CURRENT_MODELS += eynollah-binarization_20210425 CURRENT_MODELS += eynollah-binarization_20210425
CURRENT_MODELS += eynollah-column-classifier_20210425 CURRENT_MODELS += eynollah-column-classifier_20210425
CURRENT_MODELS += eynollah-enhancement_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 # tf (SavedModel format) for training
# onnx conversion for fast inference # onnx conversion for fast inference