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,
):
"""
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"

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 .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)

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

@ -25,6 +25,7 @@ MODEL_VRAM_LIMITS = {
"table": 1818,
"reading_order": 632,
"ocr": 2600, # 850 for bs 8
"extract_images": 954,
}
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-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