diff --git a/docs/train.md b/docs/train.md index ffa39a9..82bb77c 100644 --- a/docs/train.md +++ b/docs/train.md @@ -47,9 +47,9 @@ on how to generate the corresponding training dataset. The following three tasks can all be accomplished using the code in the [`train`](https://github.com/qurator-spk/eynollah/tree/main/train) directory: -* generate training dataset -* train a model -* inference with the trained model +* [Generate training dataset](#generate-training-dataset) +* [Train a model](#train-a-model) +* [Inference with the trained model](#inference-with-the-trained-model) ## Training, evaluation and output @@ -101,7 +101,7 @@ serve as labels. The enhancement model can be trained with this generated datase For machine-based reading order, we aim to determine the reading priority between two sets of text regions. The model's input is a three-channel image: the first and last channels contain information about each of the two text regions, while the middle channel encodes prominent layout elements necessary for reading order, such as separators and headers. -To generate the training dataset, our script requires a page XML file that specifies the image layout with the correct +To generate the training dataset, our script requires a PAGE XML file that specifies the image layout with the correct reading order. For output images, it is necessary to specify the width and height. Additionally, a minimum text region size can be set @@ -120,8 +120,14 @@ eynollah-training generate-gt machine-based-reading-order \ ### pagexml2label -pagexml2label is designed to generate labels from GT page XML files for various pixel-wise segmentation use cases, -including 'layout,' 'textline,' 'printspace,' 'glyph,' and 'word' segmentation. +`pagexml2label` is designed to generate labels from PAGE XML GT files for various pixel-wise segmentation use cases, +including: +- `printspace` (i.e. page frame), +- `layout` (i.e. regions), +- `textline`, +- `word`, and +- `glyph`. + To train a pixel-wise segmentation model, we require images along with their corresponding labels. Our training script expects a PNG image where each pixel corresponds to a label, represented by an integer. The background is always labeled as zero, while other elements are assigned different integers. For instance, if we have ground truth data with four @@ -131,7 +137,7 @@ In binary segmentation scenarios such as textline or page extraction, the backgr element is automatically encoded as 1 in the PNG label. To specify the desired use case and the elements to be extracted in the PNG labels, a custom JSON file can be passed. -For example, in the case of 'textline' detection, the JSON file would resemble this: +For example, in the case of textline detection, the JSON contents could be this: ```yaml { @@ -139,61 +145,77 @@ For example, in the case of 'textline' detection, the JSON file would resemble t } ``` -In the case of layout segmentation a custom config json file can look like this: +In the case of layout segmentation, the config JSON file might look like this: ```yaml { "use_case": "layout", -"textregions":{"rest_as_paragraph":1 , "drop-capital": 1, "header":2, "heading":2, "marginalia":3}, -"imageregion":4, -"separatorregion":5, -"graphicregions" :{"rest_as_decoration":6 ,"stamp":7} +"textregions": {"rest_as_paragraph": 1, "drop-capital": 1, "header": 2, "heading": 2, "marginalia": 3}, +"imageregion": 4, +"separatorregion": 5, +"graphicregions": {"rest_as_decoration": 6, "stamp": 7} } ``` -A possible custom config json file for layout segmentation where the "printspace" is a class: +The same example if `PrintSpace` (or `Border`) should be represented as a unique class: ```yaml { "use_case": "layout", -"textregions":{"rest_as_paragraph":1 , "drop-capital": 1, "header":2, "heading":2, "marginalia":3}, -"imageregion":4, -"separatorregion":5, -"graphicregions" :{"rest_as_decoration":6 ,"stamp":7} -"printspace_as_class_in_layout" : 8 +"textregions": {"rest_as_paragraph": 1, "drop-capital": 1, "header": 2, "heading": 2, "marginalia": 3}, +"imageregion": 4, +"separatorregion": 5, +"graphicregions": {"rest_as_decoration": 6, "stamp": 7} +"printspace_as_class_in_layout": 8 } ``` -For the layout use case, it is beneficial to first understand the structure of the page XML file and its elements. -In a given image, the annotations of elements are recorded in a page XML file, including their contours and classes. -For an image document, the known regions are 'textregion', 'separatorregion', 'imageregion', 'graphicregion', -'noiseregion', and 'tableregion'. +In the `layout` use-case, it is beneficial to first understand the structure of the PAGE XML file and its elements. +For a given page image, the visible segments are annotated in XML with their polygon coordinates and types. +On the region level, available segment types include `TextRegion`, `SeparatorRegion`, `ImageRegion`, `GraphicRegion`, +`NoiseRegion` and `TableRegion`. -Text regions and graphic regions also have their own specific types. The known types for text regions are 'paragraph', -'header', 'heading', 'marginalia', 'drop-capital', 'footnote', 'footnote-continued', 'signature-mark', 'page-number', -and 'catch-word'. The known types for graphic regions are 'handwritten-annotation', 'decoration', 'stamp', and -'signature'. -Since we don't know all types of text and graphic regions, unknown cases can arise. To handle these, we have defined -two additional types, "rest_as_paragraph" and "rest_as_decoration", to ensure that no unknown types are missed. -This way, users can extract all known types from the labels and be confident that no unknown types are overlooked. +Moreover, text regions and graphic regions in particular are subdivided via `@type`: +- The allowed subtypes for text regions are `paragraph`, `heading`, `marginalia`, `drop-capital`, `header`, `footnote`, +`footnote-continued`, `signature-mark`, `page-number` and `catch-word`. +- The known subtypes for graphic regions are `handwritten-annotation`, `decoration`, `stamp` and `signature`. -In the custom JSON file shown above, "header" and "heading" are extracted as the same class, while "marginalia" is shown -as a different class. All other text region types, including "drop-capital," are grouped into the same class. For the -graphic region, "stamp" has its own class, while all other types are classified together. "Image region" and "separator -region" are also present in the label. However, other regions like "noise region" and "table region" will not be -included in the label PNG file, even if they have information in the page XML files, as we chose not to include them. +These types and subtypes must be mapped to classes for the segmentation model. However, sometimes these fine-grained +distinctions are not useful or the existing annotations are not very usable (too scarce or too unreliable). +In that case, instead of these subtypes with a specific mapping, they can be pooled together by using the two special +types: +- `rest_as_paragraph` (mapping missing TextRegion subtypes and `paragraph`) +- `rest_as_decoration` (mapping missing GraphicRegion subtypes and `decoration`) + +(That way, users can extract all known types from the labels and be confident that no subtypes are overlooked.) + +In the custom JSON example shown above, `header` and `heading` are extracted as the same class, +while `marginalia` is modelled as a different class. All other text region types, including `drop-capital`, +are grouped into the same class. For graphic regions, `stamp` has its own class, while all other types +are classified together. `ImageRegion` and `SeparatorRegion` will also represented with a class label in the +training data. However, other regions like `NoiseRegion` or `TableRegion` will not be included in the PNG files, +even if they were present in the PAGE XML. + +The tool expects various command-line options: ```sh eynollah-training generate-gt pagexml2label \ - -dx "dir of GT xml files" \ - -do "dir where output label png files will be written" \ - -cfg "custom config json file" \ - -to "output type which has 2d and 3d. 2d is used for training and 3d is just to visualise the labels" + -dx "dir of input PAGE XML files" \ + -do "dir of output label PNG files" \ + -cfg "custom config JSON file" \ + -to "output type (2d or 3d)" ``` -We have also defined an artificial class that can be added to the boundary of text region types or text lines. This key -is called "artificial_class_on_boundary." If users want to apply this to certain text regions in the layout use case, -the example JSON config file should look like this: +As output type, use +- `2d` for training, +- `3d` to just visualise the labels. + +We have also defined an artificial class that can be added to (rendered around) the boundary +of text region types or text lines in order to make separation of neighbouring segments more +reliable. The key is called `artificial_class_on_boundary`, and it takes a list of text region +types to be applied to. + +Our example JSON config file could then look like this: ```yaml { @@ -215,14 +237,15 @@ the example JSON config file should look like this: } ``` -This implies that the artificial class label, denoted by 7, will be present on PNG files and will only be added to the -elements labeled as "paragraph," "header," "heading," and "marginalia." +This implies that the artificial class label (denoted by 7) will be present in the generated PNG files +and will only be added around segments labeled `paragraph`, `header`, `heading` or `marginalia`. (This +class will be handled specially during decoding at inference, and not show up in final results.) -For "textline", "word", and "glyph", the artificial class on the boundaries will be activated only if the -"artificial_class_label" key is specified in the config file. Its value should be set as 2 since these elements -represent binary cases. For example, if the background and textline are denoted as 0 and 1 respectively, then the -artificial class should be assigned the value 2. The example JSON config file should look like this for "textline" use -case: +For `printspace`, `textline`, `word`, and `glyph` segmentation use-cases, there is no `artificial_class_on_boundary` key, +but `artificial_class_label` is available. If specified in the config file, then its value should be set at 2, because +these elements represent binary classification problems (with background represented as 0, and segments as 1, respectively). + +For example, the JSON config for textline detection could look as follows: ```yaml { @@ -231,33 +254,33 @@ case: } ``` -If the coordinates of "PrintSpace" or "Border" are present in the page XML ground truth files, and the user wishes to -crop only the print space area, this can be achieved by activating the "-ps" argument. However, it should be noted that -in this scenario, since cropping will be applied to the label files, the directory of the original images must be -provided to ensure that they are cropped in sync with the labels. This ensures that the correct images and labels -required for training are obtained. The command should resemble the following: +If the coordinates of `PrintSpace` (or `Border`) are present in the PAGE XML ground truth files, +and one wishes to crop images to only cover the print space bounding box, this can be achieved +by passing the `-ps` option. Note that in this scenario, the directory of the original images +must also be provided, to ensure that the images are cropped in sync with the labels. The command +line would then resemble this: ```sh eynollah-training generate-gt pagexml2label \ - -dx "dir of GT xml files" \ - -do "dir where output label png files will be written" \ - -cfg "custom config json file" \ - -to "output type which has 2d and 3d. 2d is used for training and 3d is just to visualise the labels" \ + -dx "dir of input PAGE XML files" \ + -do "dir of output label PNG files" \ + -cfg "custom config JSON file" \ + -to "output type (2d or 3d)" \ -ps \ - -di "dir where the org images are located" \ - -doi "dir where the cropped output images will be written" + -di "dir of input original images" \ + -doi "dir of output cropped images" ``` ## Train a model ### classification -For the classification use case, we haven't provided a ground truth generator, as it's unnecessary. For classification, -all we require is a training directory with subdirectories, each containing images of its respective classes. We need +For the image classification use-case, we have not provided a ground truth generator, as it is unnecessary. +All we require is a training directory with subdirectories, each containing images of its respective classes. We need separate directories for training and evaluation, and the class names (subdirectories) must be consistent across both directories. Additionally, the class names should be specified in the config JSON file, as shown in the following example. If, for instance, we aim to classify "apple" and "orange," with a total of 2 classes, the -"classification_classes_name" key in the config file should appear as follows: +`classification_classes_name` key in the config file should appear as follows: ```yaml { @@ -279,7 +302,7 @@ example. If, for instance, we aim to classify "apple" and "orange," with a total } ``` -The "dir_train" should be like this: +Then `dir_train` should be like this: ``` . @@ -288,7 +311,7 @@ The "dir_train" should be like this: └── orange # directory of images for orange class ``` -And the "dir_eval" the same structure as train directory: +And `dir_eval` analogously: ``` . @@ -348,7 +371,7 @@ And the "dir_eval" the same structure as train directory: └── labels # directory of labels ``` -The classification model can be trained like the classification case command line. +The reading-order model can be trained like the classification case command line. ### Segmentation (Textline, Binarization, Page extraction and layout) and enhancement @@ -358,51 +381,17 @@ The following parameter configuration can be applied to all segmentation use cas its sub-parameters, and continued training are defined only for segmentation use cases and enhancements, not for classification and machine-based reading order, as you can see in their example config files. -* `backbone_type`: For segmentation tasks (such as text line, binarization, and layout detection) and enhancement, we - offer two backbone options: a "nontransformer" and a "transformer" backbone. For the "transformer" backbone, we first - apply a CNN followed by a transformer. In contrast, the "nontransformer" backbone utilizes only a CNN ResNet-50. -* `task`: The task parameter can have values such as "segmentation", "enhancement", "classification", and "reading_order". -* `patches`: If you want to break input images into smaller patches (input size of the model) you need to set this -* parameter to `true`. In the case that the model should see the image once, like page extraction, patches should be - set to ``false``. -* `n_batch`: Number of batches at each iteration. -* `n_classes`: Number of classes. In the case of binary classification this should be 2. In the case of reading_order it - should set to 1. And for the case of layout detection just the unique number of classes should be given. -* `n_epochs`: Number of epochs. -* `input_height`: This indicates the height of model's input. -* `input_width`: This indicates the width of model's input. -* `weight_decay`: Weight decay of l2 regularization of model layers. -* `pretraining`: Set to `true` to load pretrained weights of ResNet50 encoder. The downloaded weights should be saved - in a folder named "pretrained_model" in the same directory of "train.py" script. -* `augmentation`: If you want to apply any kind of augmentation this parameter should first set to `true`. -* `flip_aug`: If `true`, different types of filp will be applied on image. Type of flips is given with "flip_index" parameter. -* `blur_aug`: If `true`, different types of blurring will be applied on image. Type of blurrings is given with "blur_k" parameter. -* `scaling`: If `true`, scaling will be applied on image. Scale of scaling is given with "scales" parameter. -* `degrading`: If `true`, degrading will be applied to the image. The amount of degrading is defined with "degrade_scales" parameter. -* `brightening`: If `true`, brightening will be applied to the image. The amount of brightening is defined with "brightness" parameter. -* `rotation_not_90`: If `true`, rotation (not 90 degree) will be applied on image. Rotation angles are given with "thetha" parameter. -* `rotation`: If `true`, 90 degree rotation will be applied on image. -* `binarization`: If `true`,Otsu thresholding will be applied to augment the input data with binarized images. -* `scaling_bluring`: If `true`, combination of scaling and blurring will be applied on image. -* `scaling_binarization`: If `true`, combination of scaling and binarization will be applied on image. -* `scaling_flip`: If `true`, combination of scaling and flip will be applied on image. -* `flip_index`: Type of flips. -* `blur_k`: Type of blurrings. -* `scales`: Scales of scaling. -* `brightness`: The amount of brightenings. -* `thetha`: Rotation angles. -* `degrade_scales`: The amount of degradings. -* `continue_training`: If `true`, it means that you have already trained a model and you would like to continue the - training. So it is needed to providethe dir of trained model with "dir_of_start_model" and index for naming - themodels. For example if you have already trained for 3 epochs then your lastindex is 2 and if you want to continue - from model_1.h5, you can set `index_start` to 3 to start naming model with index 3. -* `weighted_loss`: If `true`, this means that you want to apply weighted categorical_crossentropy as loss fucntion. Be carefull if you set to `true`the parameter "is_loss_soft_dice" should be ``false`` -* `data_is_provided`: If you have already provided the input data you can set this to `true`. Be sure that the train - and eval data are in"dir_output".Since when once we provide training data we resize and augmentthem and then wewrite - them in sub-directories train and eval in "dir_output". -* `dir_train`: This is the directory of "images" and "labels" (dir_train should include two subdirectories with names of images and labels ) for raw images and labels. Namely they are not prepared (not resized and not augmented) yet for training the model. When we run this tool these raw data will be transformed to suitable size needed for the model and they will be written in "dir_output" in train and eval directories. Each of train and eval include "images" and "labels" sub-directories. -* `index_start`: Starting index for saved models in the case that "continue_training" is `true`. -* `dir_of_start_model`: Directory containing pretrained model to continue training the model in the case that "continue_training" is `true`. +* `task`: The task parameter must be one of the following values: + - `binarization`, + - `enhancement`, + - `segmentation`, + - `classification`, + - `reading_order`. +* `backbone_type`: For the tasks `segmentation` (such as text line, and region layout detection), + `binarization` and `enhancement`, we offer two backbone options: + - `nontransformer` (only a CNN ResNet-50). + - `transformer` (first apply a CNN, followed by a transformer) +* `transformer_cnn_first`: Whether to apply the CNN first (followed by the transformer) when using `transformer` backbone. * `transformer_num_patches_xy`: Number of patches for vision transformer in x and y direction respectively. * `transformer_patchsize_x`: Patch size of vision transformer patches in x direction. * `transformer_patchsize_y`: Patch size of vision transformer patches in y direction. @@ -410,11 +399,63 @@ classification and machine-based reading order, as you can see in their example * `transformer_mlp_head_units`: Transformer Multilayer Perceptron (MLP) head units. Default value is [128, 64]. * `transformer_layers`: transformer layers. Default value is 8. * `transformer_num_heads`: Transformer number of heads. Default value is 4. -* `transformer_cnn_first`: We have two types of vision transformers. In one type, a CNN is applied first, followed by a transformer. In the other type, this order is reversed. If transformer_cnn_first is true, it means the CNN will be applied before the transformer. Default value is true. +* `patches`: Whether to break up (tile) input images into smaller patches (input size of the model). + If `false`, the model will see the image once (resized to the input size of the model). + Should be set to `false` for cases like page extraction. +* `n_batch`: Number of batches at each iteration. +* `n_classes`: Number of classes. In the case of binary classification this should be 2. In the case of reading_order it + should set to 1. And for the case of layout detection just the unique number of classes should be given. +* `n_epochs`: Number of epochs (iterations over the data) to train. +* `input_height`: the image height for the model's input. +* `input_width`: the image width for the model's input. +* `weight_decay`: Weight decay of l2 regularization of model layers. +* `weighted_loss`: If `true`, this means that you want to apply weighted categorical crossentropy as loss function. + (Mutually exclusive with `is_loss_soft_dice`, and only applies for `segmentation` and `binarization` tasks.) +* `pretraining`: Set to `true` to (download and) initialise pretrained weights of ResNet50 encoder. +* `dir_train`: Path to directory of raw training data (as extracted via `pagexml2labels`, i.e. with subdirectories + `images` and `labels` for input images and output labels. + (These are not prepared for training the model, yet. Upon first run, the raw data will be transformed to suitable size + needed for the model, and written in `dir_output` under `train` and `eval` subdirectories. See `data_is_provided`.) +* `dir_eval`: Ditto for raw evaluation data. +* `dir_output`: Directory to write model checkpoints, logs (for Tensorboard) and precomputed images to. +* `data_is_provided`: If you have already trained at least one complete epoch (using the same data settings) before, + you can set this to `true` to avoid computing the resized / patched / augmented image files again. + Be sure that there are subdirectories `train` and `eval` data are in `dir_output` (each with subdirectories `images` + and `labels`, respectively). +* `continue_training`: If `true`, continue training a model checkpoint from a previous run. + This requires providing the directory of the model checkpoint to load via `dir_of_start_model` + and setting `index_start` counter for naming new checkpoints. + For example if you have already trained for 3 epochs, then your last index is 2, so if you want + to continue with `model_04`, `model_05` etc., set `index_start=3`. +* `index_start`: Starting index for saving models in the case that `continue_training` is `true`. + (Existing checkpoints above this will be overwritten.) +* `dir_of_start_model`: Directory containing existing model checkpoint to initialise model weights from when `continue_training=true`. + (Can be an epoch-interval checkpoint, or batch-interval checkpoint from `save_interval`.) +* `augmentation`: If you want to apply any kind of augmentation this parameter should first set to `true`. + The remaining settings pertain to that... +* `flip_aug`: If `true`, different types of flipping over the image arrays. Requires `flip_index` parameter. +* `flip_index`: List of flip codes (as in `cv2.flip`, i.e. 0 for vertical, positive for horizontal shift, negative for vertical and horizontal shift). +* `blur_aug`: If `true`, different types of blurring will be applied on image. Requires `blur_k` parameter. +* `blur_k`: Method of blurring (`gauss`, `median` or `blur`). +* `scaling`: If `true`, scaling will be applied on image. Requires `scales` parameter. +* `scales`: List of scale factors for scaling. +* `scaling_bluring`: If `true`, combination of scaling and blurring will be applied on image. +* `scaling_binarization`: If `true`, combination of scaling and binarization will be applied on image. +* `scaling_flip`: If `true`, combination of scaling and flip will be applied on image. +* `degrading`: If `true`, degrading will be applied to the image. Requires `degrade_scales` parameter. +* `degrade_scales`: List of intensity factors for degrading. +* `brightening`: If `true`, brightening will be applied to the image. Requires `brightness` parameter. +* `brightness`: List of intensity factors for brightening. +* `binarization`: If `true`, Otsu thresholding will be applied to augment the input data with binarized images. +* `dir_img_bin`: With `binarization`, use this directory to read precomputed binarized images instead of ad-hoc Otsu. + (Base names should correspond to the files in `dir_train/images`.) +* `rotation`: If `true`, 90° rotation will be applied on images. +* `rotation_not_90`: If `true`, random rotation (other than 90°) will be applied on image. Requires `thetha` parameter. +* `thetha`: List of rotation angles (in degrees). -In the case of segmentation and enhancement the train and evaluation directory should be as following. +In case of segmentation and enhancement the train and evaluation data should be organised as follows. -The "dir_train" should be like this: +The "dir_train" directory should be like this: ``` . @@ -432,11 +473,12 @@ And the "dir_eval" the same structure as train directory: └── labels # directory of labels ``` -After configuring the JSON file for segmentation or enhancement, training can be initiated by running the following -command, similar to the process for classification and reading order: +After configuring the JSON file for segmentation or enhancement, +training can be initiated by running the following command line, +similar to classification and reading-order model training: -``` -eynollah-training train with config_classification.json` +```sh +eynollah-training train with config_classification.json ``` #### Binarization @@ -728,7 +770,7 @@ This will straightforwardly return the class of the image. ### machine based reading order -To infer the reading order using a reading order model, we need a page XML file containing layout information but +To infer the reading order using a reading order model, we need a PAGE XML file containing layout information but without the reading order. We simply need to provide the model directory, the XML file, and the output directory. The new XML file with the added reading order will be written to the output directory with the same name. We need to run: diff --git a/requirements.txt b/requirements.txt index bbacd48..53d1e39 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,9 @@ # ocrd includes opencv, numpy, shapely, click ocrd >= 3.3.0 -numpy <1.24.0 +numpy < 2.0 scikit-learn >= 0.23.2 -tensorflow < 2.13 +tensorflow +tf-keras # avoid keras 3 (also needs TF_USE_LEGACY_KERAS=1) numba <= 0.58.1 scikit-image biopython diff --git a/src/eynollah/cli/__init__.py b/src/eynollah/cli/__init__.py index 05dafa1..43ed046 100644 --- a/src/eynollah/cli/__init__.py +++ b/src/eynollah/cli/__init__.py @@ -2,14 +2,12 @@ # this must be the first import of the CLI! from ..eynollah_imports import imported_libs -from .cli_models import models_cli -from .cli_binarize import binarize_cli - from .cli import main from .cli_binarize import binarize_cli from .cli_enhance import enhance_cli from .cli_extract_images import extract_images_cli from .cli_layout import layout_cli +from .cli_models import models_cli from .cli_ocr import ocr_cli from .cli_readingorder import readingorder_cli diff --git a/src/eynollah/cli/cli_binarize.py b/src/eynollah/cli/cli_binarize.py index d4a6e31..aa6cefc 100644 --- a/src/eynollah/cli/cli_binarize.py +++ b/src/eynollah/cli/cli_binarize.py @@ -21,13 +21,20 @@ import click type=click.Path(file_okay=True, dir_okay=True), required=True, ) +@click.option( + "--overwrite", + "-O", + help="overwrite (instead of skipping) if output xml exists", + is_flag=True, +) @click.pass_context def binarize_cli( - ctx, - patches, - input_image, - dir_in, - output, + ctx, + patches, + input_image, + dir_in, + output, + overwrite, ): """ Binarize images with a ML model @@ -39,6 +46,7 @@ def binarize_cli( image_path=input_image, use_patches=patches, output=output, - dir_in=dir_in + dir_in=dir_in, + overwrite=overwrite ) diff --git a/src/eynollah/extract_images.py b/src/eynollah/extract_images.py index 9942cf8..7a7e3f6 100644 --- a/src/eynollah/extract_images.py +++ b/src/eynollah/extract_images.py @@ -116,19 +116,19 @@ class EynollahImageExtractor(Eynollah): 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] - mask_lines_only = (prediction_regions_org[:,:] ==3)*1 + mask_seps_only = (prediction_regions_org[:,:] ==3)*1 mask_texts_only = (prediction_regions_org[:,:] ==1)*1 mask_images_only=(prediction_regions_org[:,:] ==2)*1 - polygons_seplines, hir_seplines = return_contours_of_image(mask_lines_only) + polygons_seplines, hir_seplines = return_contours_of_image(mask_seps_only) polygons_seplines = filter_contours_area_of_image( - mask_lines_only, polygons_seplines, hir_seplines, max_area=1, min_area=0.00001, dilate=1) + mask_seps_only, polygons_seplines, hir_seplines, 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_lines = return_contours_of_interested_region(mask_lines_only,1,0.00001) + 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) - text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts = polygons_of_only_lines, color=(3,3,3)) + 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)) @@ -255,24 +255,24 @@ class EynollahImageExtractor(Eynollah): self.get_regions_light_v_extract_only_images(img_res, num_col_classifier) 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=[], + 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=[], ) if self.plotter: self.plotter.write_images_into_directory(polygons_of_images, image_page) diff --git a/src/eynollah/eynollah.py b/src/eynollah/eynollah.py index 9383c5e..d089511 100644 --- a/src/eynollah/eynollah.py +++ b/src/eynollah/eynollah.py @@ -70,11 +70,7 @@ from .utils.contour import ( join_polygons, make_intersection, ) -from .utils.rotate import ( - rotate_image, - rotation_not_90_func, - rotation_not_90_func_full_layout, -) +from .utils.rotate import rotate_image from .utils.separate_lines import ( return_deskew_slop, do_work_of_slopes_new_curved, @@ -83,6 +79,7 @@ from .utils.marginals import get_marginals from .utils.resize import resize_image from .utils.shm import share_ndarray from .utils import ( + ensure_array, is_image_filename, isNaN, crop_image_inside_box, @@ -97,7 +94,6 @@ from .utils import ( return_boxes_of_images_by_order_of_reading_new ) from .utils.pil_cv2 import pil2cv -from .utils.xml import order_and_id_of_texts from .plot import EynollahPlotter from .writer import EynollahXmlWriter @@ -165,12 +161,12 @@ class Eynollah: # for parallelization of CPU-intensive tasks: self.executor = ProcessPoolExecutor(max_workers=cpu_count()) - + if threshold_art_class_layout: self.threshold_art_class_layout = float(threshold_art_class_layout) else: self.threshold_art_class_layout = 0.1 - + if threshold_art_class_textline: self.threshold_art_class_textline = float(threshold_art_class_textline) else: @@ -178,20 +174,12 @@ class Eynollah: t_start = time.time() - # #gpu_options = tf.compat.v1.GPUOptions(allow_growth=True) - # #gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=7.7, allow_growth=True) - # #session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) - # config = tf.compat.v1.ConfigProto() - # config.gpu_options.allow_growth = True - # #session = tf.InteractiveSession() - # session = tf.compat.v1.Session(config=config) - # set_session(session) try: for device in tf.config.list_physical_devices('GPU'): tf.config.experimental.set_memory_growth(device, True) except: self.logger.warning("no GPU device available") - + self.logger.info("Loading models...") self.setup_models() self.logger.info(f"Model initialization complete ({time.time() - t_start:.1f}s)") @@ -527,7 +515,7 @@ class Eynollah: label_p_pred = self.model_zoo.get("col_classifier").predict(img_in, verbose=0) num_col = np.argmax(label_p_pred[0]) + 1 - + elif (self.num_col_upper and self.num_col_lower) and (self.num_col_upper!=self.num_col_lower): if self.input_binary: img_in = np.copy(img) @@ -659,12 +647,12 @@ class Eynollah: seg_art[seg_art0] =1 - + skeleton_art = skeletonize(seg_art) skeleton_art = skeleton_art*1 seg[skeleton_art==1]=2 - + if thresholding_for_fl_light_version: seg_header = label_p_pred[0,:,:,2] @@ -672,7 +660,7 @@ class Eynollah: seg_header[seg_header>0] =1 seg[seg_header==1]=2 - + seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2) prediction_true = resize_image(seg_color, img_h_page, img_w_page).astype(np.uint8) return prediction_true @@ -761,7 +749,7 @@ class Eynollah: indexer_inside_batch = 0 for i_batch, j_batch in zip(list_i_s, list_j_s): seg_in = seg[indexer_inside_batch] - + if thresholding_for_artificial_class_in_light_version: seg_in_art = seg_art[indexer_inside_batch] @@ -782,7 +770,7 @@ class Eynollah: index_x_d_in + 0:index_x_u_in - margin, 1] = \ seg_in_art[0:-margin or None, 0:-margin or None] - + elif i_batch == nxf - 1 and j_batch == nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - 0, index_x_d_in + margin:index_x_u_in - 0] = \ @@ -794,7 +782,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - 0, 1] = \ seg_in_art[margin:, margin:] - + elif i_batch == 0 and j_batch == nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - 0, index_x_d_in + 0:index_x_u_in - margin] = \ @@ -806,7 +794,7 @@ class Eynollah: index_x_d_in + 0:index_x_u_in - margin, 1] = \ seg_in_art[margin:, 0:-margin or None] - + elif i_batch == nxf - 1 and j_batch == 0: prediction_true[index_y_d_in + 0:index_y_u_in - margin, index_x_d_in + margin:index_x_u_in - 0] = \ @@ -818,7 +806,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - 0, 1] = \ seg_in_art[0:-margin or None, margin:] - + elif i_batch == 0 and j_batch != 0 and j_batch != nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - margin, index_x_d_in + 0:index_x_u_in - margin] = \ @@ -830,7 +818,7 @@ class Eynollah: index_x_d_in + 0:index_x_u_in - margin, 1] = \ seg_in_art[margin:-margin or None, 0:-margin or None] - + elif i_batch == nxf - 1 and j_batch != 0 and j_batch != nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - margin, index_x_d_in + margin:index_x_u_in - 0] = \ @@ -842,7 +830,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - 0, 1] = \ seg_in_art[margin:-margin or None, margin:] - + elif i_batch != 0 and i_batch != nxf - 1 and j_batch == 0: prediction_true[index_y_d_in + 0:index_y_u_in - margin, index_x_d_in + margin:index_x_u_in - margin] = \ @@ -854,7 +842,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - margin, 1] = \ seg_in_art[0:-margin or None, margin:-margin or None] - + elif i_batch != 0 and i_batch != nxf - 1 and j_batch == nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - 0, index_x_d_in + margin:index_x_u_in - margin] = \ @@ -866,7 +854,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - margin, 1] = \ seg_in_art[margin:, margin:-margin or None] - + else: prediction_true[index_y_d_in + margin:index_y_u_in - margin, index_x_d_in + margin:index_x_u_in - margin] = \ @@ -892,16 +880,16 @@ class Eynollah: img_patch[:] = 0 prediction_true = prediction_true.astype(np.uint8) - + if thresholding_for_artificial_class_in_light_version: kernel_min = np.ones((3, 3), np.uint8) prediction_true[:,:,0][prediction_true[:,:,0]==2] = 0 - + skeleton_art = skeletonize(prediction_true[:,:,1]) skeleton_art = skeleton_art*1 - + skeleton_art = skeleton_art.astype('uint8') - + skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) prediction_true[:,:,0][skeleton_art==1]=2 @@ -932,7 +920,7 @@ class Eynollah: seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2) prediction_true = resize_image(seg_color, img_h_page, img_w_page).astype(np.uint8) - + if thresholding_for_artificial_class_in_light_version: kernel_min = np.ones((3, 3), np.uint8) seg_art = label_p_pred[0,:,:,4] @@ -940,18 +928,18 @@ class Eynollah: seg_art[seg_art>0] =1 #seg[seg_art==1]=4 seg_art = resize_image(seg_art, img_h_page, img_w_page).astype(np.uint8) - + prediction_true[:,:,0][prediction_true[:,:,0]==4] = 0 - + skeleton_art = skeletonize(seg_art) skeleton_art = skeleton_art*1 - + skeleton_art = skeleton_art.astype('uint8') - + skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) - + prediction_true[:,:,0][skeleton_art==1] = 4 - + return prediction_true , resize_image(label_p_pred[0, :, :, 1] , img_h_page, img_w_page) if img.shape[0] < img_height_model: @@ -1044,7 +1032,7 @@ class Eynollah: indexer_inside_batch = 0 for i_batch, j_batch in zip(list_i_s, list_j_s): seg_in = seg[indexer_inside_batch] - + if (thresholding_for_artificial_class_in_light_version or thresholding_for_some_classes_in_light_version): seg_in_art = seg_art[indexer_inside_batch] @@ -1072,7 +1060,7 @@ class Eynollah: index_x_d_in + 0:index_x_u_in - margin, 1] = \ seg_in_art[0:-margin or None, 0:-margin or None] - + elif i_batch == nxf - 1 and j_batch == nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - 0, index_x_d_in + margin:index_x_u_in - 0] = \ @@ -1090,7 +1078,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - 0, 1] = \ seg_in_art[margin:, margin:] - + elif i_batch == 0 and j_batch == nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - 0, index_x_d_in + 0:index_x_u_in - margin] = \ @@ -1102,14 +1090,14 @@ class Eynollah: label_p_pred[0, margin:, 0:-margin or None, 1] - + if (thresholding_for_artificial_class_in_light_version or thresholding_for_some_classes_in_light_version): prediction_true[index_y_d_in + margin:index_y_u_in - 0, index_x_d_in + 0:index_x_u_in - margin, 1] = \ seg_in_art[margin:, 0:-margin or None] - + elif i_batch == nxf - 1 and j_batch == 0: prediction_true[index_y_d_in + 0:index_y_u_in - margin, index_x_d_in + margin:index_x_u_in - 0] = \ @@ -1127,7 +1115,7 @@ class Eynollah: index_x_d_in + margin:index_x_u_in - 0, 1] = \ seg_in_art[0:-margin or None, margin:] - + elif i_batch == 0 and j_batch != 0 and j_batch != nyf - 1: prediction_true[index_y_d_in + margin:index_y_u_in - margin, index_x_d_in + 0:index_x_u_in - margin] = \ @@ -1226,29 +1214,29 @@ class Eynollah: img_patch[:] = 0 prediction_true = prediction_true.astype(np.uint8) - + if thresholding_for_artificial_class_in_light_version: kernel_min = np.ones((3, 3), np.uint8) prediction_true[:,:,0][prediction_true[:,:,0]==2] = 0 - + skeleton_art = skeletonize(prediction_true[:,:,1]) skeleton_art = skeleton_art*1 - + skeleton_art = skeleton_art.astype('uint8') - + skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) prediction_true[:,:,0][skeleton_art==1]=2 - + if thresholding_for_some_classes_in_light_version: kernel_min = np.ones((3, 3), np.uint8) prediction_true[:,:,0][prediction_true[:,:,0]==4] = 0 - + skeleton_art = skeletonize(prediction_true[:,:,1]) skeleton_art = skeleton_art*1 - + skeleton_art = skeleton_art.astype('uint8') - + skeleton_art = cv2.dilate(skeleton_art, kernel_min, iterations=1) prediction_true[:,:,0][skeleton_art==1]=4 @@ -1294,10 +1282,10 @@ class Eynollah: else: box = [0, 0, self.image.shape[1], self.image.shape[0]] cropped_page, page_coord = crop_image_inside_box(box, self.image) - cont_page.append(np.array([[page_coord[2], page_coord[0]], - [page_coord[3], page_coord[0]], - [page_coord[3], page_coord[1]], - [page_coord[2], page_coord[1]]])) + cont_page.append(np.array([[[page_coord[2], page_coord[0]]], + [[page_coord[3], page_coord[0]]], + [[page_coord[3], page_coord[1]]], + [[page_coord[2], page_coord[1]]]])) return cropped_page, page_coord, cont_page def early_page_for_num_of_column_classification(self,img_bin): @@ -1374,16 +1362,16 @@ class Eynollah: prediction_regions = resize_image(prediction_regions, img_height_h, img_width_h) self.logger.debug("exit extract_text_regions") return prediction_regions, None - + def get_textlines_of_a_textregion_sorted(self, textlines_textregion, cx_textline, cy_textline, w_h_textline): N = len(cy_textline) if N==0: return [] - + diff_cy = np.abs( np.diff(sorted(cy_textline)) ) diff_cx = np.abs(np.diff(sorted(cx_textline)) ) - + if len(diff_cy)>0: mean_y_diff = np.mean(diff_cy) mean_x_diff = np.mean(diff_cx) @@ -1395,13 +1383,13 @@ class Eynollah: mean_x_diff = 0 count_hor = 1 count_ver = 0 - + if count_hor >= count_ver: row_threshold = mean_y_diff / 1.5 if mean_y_diff > 0 else 10 indices_sorted_by_y = sorted(range(N), key=lambda i: cy_textline[i]) - + rows = [] current_row = [indices_sorted_by_y[0]] for i in range(1, N): @@ -1470,7 +1458,7 @@ class Eynollah: cx_textline_in, cy_textline_in, w_h_textlines_in) - + all_found_textline_polygons.append(textlines_ins)#[::-1]) slopes.append(slope_deskew) @@ -1523,27 +1511,24 @@ class Eynollah: textline_mask_tot_ea_art = (prediction_textline[:,:]==2)*1 old_art = np.copy(textline_mask_tot_ea_art) - + textline_mask_tot_ea_lines = (prediction_textline[:,:]==1)*1 textline_mask_tot_ea_lines = textline_mask_tot_ea_lines.astype('uint8') prediction_textline[:,:][textline_mask_tot_ea_lines[:,:]==1]=1 - + #cv2.imwrite('prediction_textline2.png', prediction_textline[:,:,0]) prediction_textline_longshot = self.do_prediction(False, img, self.model_zoo.get("textline")) prediction_textline_longshot_true_size = resize_image(prediction_textline_longshot, img_h, img_w) - - + + #cv2.imwrite('prediction_textline.png', prediction_textline[:,:,0]) #sys.exit() self.logger.debug('exit textline_contours') return ((prediction_textline[:, :, 0]==1).astype(np.uint8), (prediction_textline_longshot_true_size[:, :, 0]==1).astype(np.uint8)) - - - def get_regions_light_v(self,img,is_image_enhanced, num_col_classifier): self.logger.debug("enter get_regions_light_v") t_in = time.time() @@ -1649,7 +1634,7 @@ class Eynollah: img_bin = resize_image(img_bin, img_height_h, img_width_h ) prediction_regions_org=prediction_regions_org[:,:,0] - mask_lines_only = (prediction_regions_org[:,:] ==3)*1 + mask_seps_only = (prediction_regions_org[:,:] == 3)*1 mask_texts_only = (prediction_regions_org[:,:] ==1)*1 mask_texts_only = mask_texts_only.astype('uint8') @@ -1660,7 +1645,7 @@ class Eynollah: mask_texts_only = cv2.dilate(mask_texts_only, kernel=np.ones((2,2), np.uint8), iterations=1) mask_images_only=(prediction_regions_org[:,:] ==2)*1 - polygons_seplines, hir_seplines = return_contours_of_image(mask_lines_only) + polygons_seplines, hir_seplines = return_contours_of_image(mask_seps_only) test_khat = np.zeros(prediction_regions_org.shape) test_khat = cv2.fillPoly(test_khat, pts=polygons_seplines, color=(1,1,1)) @@ -1674,7 +1659,7 @@ class Eynollah: #plt.show() polygons_seplines = filter_contours_area_of_image( - mask_lines_only, polygons_seplines, hir_seplines, max_area=1, min_area=0.00001, dilate=1) + mask_seps_only, polygons_seplines, hir_seplines, max_area=1, min_area=0.00001, dilate=1) test_khat = np.zeros(prediction_regions_org.shape) test_khat = cv2.fillPoly(test_khat, pts = polygons_seplines, color=(1,1,1)) @@ -1685,10 +1670,10 @@ class Eynollah: 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_lines = return_contours_of_interested_region(mask_lines_only,1,0.00001) + 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) - text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts=polygons_of_only_lines, color=(3,3,3)) + 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)) @@ -1710,8 +1695,8 @@ class Eynollah: self, contours_only_text_parent, contours_only_text_parent_h, boxes, textline_mask_tot): self.logger.debug("enter do_order_of_regions") - contours_only_text_parent = np.array(contours_only_text_parent) - contours_only_text_parent_h = np.array(contours_only_text_parent_h) + contours_only_text_parent = ensure_array(contours_only_text_parent) + contours_only_text_parent_h = ensure_array(contours_only_text_parent_h) 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))) @@ -1719,11 +1704,15 @@ class Eynollah: contours_only_text_parent) cx_head, cy_head, mx_head, Mx_head, my_head, My_head, mxy_head = find_new_features_of_contours( contours_only_text_parent_h) + cx_main = np.array(cx_main, dtype=int) + cy_main = np.array(cy_main, dtype=int) + cx_head = np.array(cx_head, dtype=int) + cy_head = np.array(cy_head, dtype=int) def match_boxes(only_centers: bool): arg_text_con_main = np.zeros(len(contours_only_text_parent), dtype=int) for ii in range(len(contours_only_text_parent)): - check_if_textregion_located_in_a_box = False + box_found = False for jj, box in enumerate(boxes): if ((cx_main[ii] >= box[0] and cx_main[ii] < box[1] and @@ -1734,20 +1723,23 @@ class Eynollah: my_main[ii] >= box[2] and My_main[ii] < box[3])): arg_text_con_main[ii] = jj - check_if_textregion_located_in_a_box = True + box_found = True + # print("main/matched ", ii, "\t", (mx_main[ii], Mx_main[ii], my_main[ii], My_main[ii]), "\tin", jj, box, only_centers) break - if not check_if_textregion_located_in_a_box: + if not box_found: dists_tr_from_box = np.linalg.norm(c_boxes - np.array([[cy_main[ii]], [cx_main[ii]]]), axis=0) pcontained_in_box = ((boxes[:, 2] <= cy_main[ii]) & (cy_main[ii] < boxes[:, 3]) & (boxes[:, 0] <= cx_main[ii]) & (cx_main[ii] < boxes[:, 1])) + assert pcontained_in_box.any(), (ii, cx_main[ii], cy_main[ii]) ind_min = np.argmin(np.ma.masked_array(dists_tr_from_box, ~pcontained_in_box)) arg_text_con_main[ii] = ind_min + # print("main/fallback ", ii, "\t", (mx_main[ii], Mx_main[ii], my_main[ii], My_main[ii]), "\tin", ind_min, boxes[ind_min], only_centers) args_contours_main = np.arange(len(contours_only_text_parent)) order_by_con_main = np.zeros_like(arg_text_con_main) arg_text_con_head = np.zeros(len(contours_only_text_parent_h), dtype=int) for ii in range(len(contours_only_text_parent_h)): - check_if_textregion_located_in_a_box = False + box_found = False for jj, box in enumerate(boxes): if ((cx_head[ii] >= box[0] and cx_head[ii] < box[1] and @@ -1758,20 +1750,21 @@ class Eynollah: my_head[ii] >= box[2] and My_head[ii] < box[3])): arg_text_con_head[ii] = jj - check_if_textregion_located_in_a_box = True + box_found = True + # print("head/matched ", ii, "\t", (mx_head[ii], Mx_head[ii], my_head[ii], My_head[ii]), "\tin", jj, box, only_centers) break - if not check_if_textregion_located_in_a_box: + if not box_found: dists_tr_from_box = np.linalg.norm(c_boxes - np.array([[cy_head[ii]], [cx_head[ii]]]), axis=0) pcontained_in_box = ((boxes[:, 2] <= cy_head[ii]) & (cy_head[ii] < boxes[:, 3]) & (boxes[:, 0] <= cx_head[ii]) & (cx_head[ii] < boxes[:, 1])) + assert pcontained_in_box.any(), (ii, cx_head[ii], cy_head[ii]) ind_min = np.argmin(np.ma.masked_array(dists_tr_from_box, ~pcontained_in_box)) arg_text_con_head[ii] = ind_min + # print("head/fallback ", ii, "\t", (mx_head[ii], Mx_head[ii], my_head[ii], My_head[ii]), "\tin", ind_min, boxes[ind_min], only_centers) args_contours_head = np.arange(len(contours_only_text_parent_h)) order_by_con_head = np.zeros_like(arg_text_con_head) - ref_point = 0 - order_of_texts_tot = [] - id_of_texts_tot = [] + idx = 0 for iij, box in enumerate(boxes): ys = slice(*box[2:4]) xs = slice(*box[0:2]) @@ -1780,42 +1773,30 @@ class Eynollah: con_inter_box = contours_only_text_parent[args_contours_box_main] con_inter_box_h = contours_only_text_parent_h[args_contours_box_head] - indexes_sorted, kind_of_texts_sorted, index_by_kind_sorted = order_of_regions( - textline_mask_tot[ys, xs], con_inter_box, con_inter_box_h, box[2]) + _, kind_of_texts_sorted, index_by_kind_sorted = order_of_regions( + textline_mask_tot[ys, xs], con_inter_box, con_inter_box_h, box[2], box[0]) - order_of_texts, id_of_texts = order_and_id_of_texts( - con_inter_box, con_inter_box_h, - indexes_sorted, index_by_kind_sorted, kind_of_texts_sorted, ref_point) + for tidx, kind in zip(index_by_kind_sorted, kind_of_texts_sorted): + if kind == 1: + # print(iij, "main", args_contours_box_main[tidx], "becomes", idx) + order_by_con_main[args_contours_box_main[tidx]] = idx + else: + # print(iij, "head", args_contours_box_head[tidx], "becomes", idx) + order_by_con_head[args_contours_box_head[tidx]] = idx + idx += 1 - indexes_sorted_main = indexes_sorted[kind_of_texts_sorted == 1] - indexes_by_type_main = index_by_kind_sorted[kind_of_texts_sorted == 1] - indexes_sorted_head = indexes_sorted[kind_of_texts_sorted == 2] - indexes_by_type_head = index_by_kind_sorted[kind_of_texts_sorted == 2] - - for zahler, _ in enumerate(args_contours_box_main): - arg_order_v = indexes_sorted_main[zahler] - order_by_con_main[args_contours_box_main[indexes_by_type_main[zahler]]] = \ - np.flatnonzero(indexes_sorted == arg_order_v) + ref_point - - for zahler, _ in enumerate(args_contours_box_head): - arg_order_v = indexes_sorted_head[zahler] - order_by_con_head[args_contours_box_head[indexes_by_type_head[zahler]]] = \ - np.flatnonzero(indexes_sorted == arg_order_v) + ref_point - - for jji in range(len(id_of_texts)): - order_of_texts_tot.append(order_of_texts[jji] + ref_point) - id_of_texts_tot.append(id_of_texts[jji]) - ref_point += len(id_of_texts) - - order_of_texts_tot = np.concatenate((order_by_con_main, - order_by_con_head)) - order_text_new = np.argsort(order_of_texts_tot) - return order_text_new, id_of_texts_tot + # xml writer will create region ids in order of + # - contours_only_text_parent (main text), followed by + # - contours_only_text_parent (headings), + # 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))) + return order_text_new try: results = match_boxes(False) except Exception as why: - self.logger.error(why) + self.logger.exception(why) results = match_boxes(True) self.logger.debug("exit do_order_of_regions") @@ -1893,45 +1874,35 @@ class Eynollah: return layout_org, contours_new - def delete_separator_around(self, spliter_y,peaks_neg,image_by_region, pixel_line, pixel_table): + def delete_separator_around(self, splitter_y, peaks_neg, image_by_region, label_seps, label_table): # format of subboxes: box=[x1, x2 , y1, y2] pix_del = 100 - if len(image_by_region.shape)==3: - for i in range(len(spliter_y)-1): - for j in range(1,len(peaks_neg[i])-1): - ys = slice(int(spliter_y[i]), - int(spliter_y[i+1])) - xs = slice(peaks_neg[i][j] - pix_del, - peaks_neg[i][j] + pix_del) - image_by_region[ys,xs,0][image_by_region[ys,xs,0]==pixel_line] = 0 - image_by_region[ys,xs,0][image_by_region[ys,xs,1]==pixel_line] = 0 - image_by_region[ys,xs,0][image_by_region[ys,xs,2]==pixel_line] = 0 - - image_by_region[ys,xs,0][image_by_region[ys,xs,0]==pixel_table] = 0 - image_by_region[ys,xs,0][image_by_region[ys,xs,1]==pixel_table] = 0 - image_by_region[ys,xs,0][image_by_region[ys,xs,2]==pixel_table] = 0 - else: - for i in range(len(spliter_y)-1): - for j in range(1,len(peaks_neg[i])-1): - ys = slice(int(spliter_y[i]), - int(spliter_y[i+1])) - xs = slice(peaks_neg[i][j] - pix_del, - peaks_neg[i][j] + pix_del) - image_by_region[ys,xs][image_by_region[ys,xs]==pixel_line] = 0 - image_by_region[ys,xs][image_by_region[ys,xs]==pixel_table] = 0 + for i in range(len(splitter_y)-1): + for j in range(1,len(peaks_neg[i])-1): + where = np.index_exp[splitter_y[i]: + splitter_y[i+1], + peaks_neg[i][j] - pix_del: + peaks_neg[i][j] + pix_del, + :] + if image_by_region.ndim < 3: + where = where[:2] + else: + print("image_by_region ndim is 3!") # rs + image_by_region[where][image_by_region[where] == label_seps] = 0 + image_by_region[where][image_by_region[where] == label_table] = 0 return image_by_region def add_tables_heuristic_to_layout( self, image_regions_eraly_p, boxes, - slope_mean_hor, spliter_y, peaks_neg_tot, image_revised, - num_col_classifier, min_area, pixel_line): + slope_mean_hor, splitter_y, peaks_neg_tot, image_revised, + num_col_classifier, min_area, label_seps): - pixel_table =10 - image_revised_1 = self.delete_separator_around(spliter_y, peaks_neg_tot, image_revised, pixel_line, pixel_table) + label_table =10 + image_revised_1 = self.delete_separator_around(splitter_y, peaks_neg_tot, image_revised, label_seps, label_table) try: - image_revised_1[:,:30][image_revised_1[:,:30]==pixel_line] = 0 - image_revised_1[:,-30:][image_revised_1[:,-30:]==pixel_line] = 0 + image_revised_1[:,:30][image_revised_1[:,:30]==label_seps] = 0 + image_revised_1[:,-30:][image_revised_1[:,-30:]==label_seps] = 0 except: pass boxes = np.array(boxes, dtype=int) # to be on the safe side @@ -1942,7 +1913,7 @@ class Eynollah: _, thresh = cv2.threshold(image_col, 0, 255, 0) contours,hirarchy=cv2.findContours(thresh.copy(), cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) - if indiv==pixel_table: + if indiv==label_table: main_contours = filter_contours_area_of_image_tables(thresh, contours, hirarchy, max_area=1, min_area=0.001) else: @@ -1958,11 +1929,11 @@ class Eynollah: box_xs = slice(*boxes[i][0:2]) image_box = img_comm[box_ys, box_xs] try: - image_box_tabels_1 = (image_box == pixel_table) * 1 + image_box_tabels_1 = (image_box == label_table) * 1 contours_tab,_=return_contours_of_image(image_box_tabels_1) contours_tab=filter_contours_area_of_image_tables(image_box_tabels_1,contours_tab,_,1,0.003) - image_box_tabels_1 = (image_box == pixel_line).astype(np.uint8) * 1 - image_box_tabels_and_m_text = ( (image_box == pixel_table) | + image_box_tabels_1 = (image_box == label_seps).astype(np.uint8) * 1 + image_box_tabels_and_m_text = ( (image_box == label_table) | (image_box == 1) ).astype(np.uint8) * 1 image_box_tabels_1 = cv2.dilate(image_box_tabels_1, KERNEL, iterations=5) @@ -2024,7 +1995,7 @@ class Eynollah: y_up_tabs=[] for ii in range(len(y_up_tabs)): - image_box[y_up_tabs[ii]:y_down_tabs[ii]] = pixel_table + image_box[y_up_tabs[ii]:y_down_tabs[ii]] = label_table image_revised_last[box_ys, box_xs] = image_box else: @@ -2035,14 +2006,14 @@ class Eynollah: image_revised_last[box_ys, box_xs] = image_box if num_col_classifier==1: - img_tables_col_1 = (image_revised_last == pixel_table).astype(np.uint8) + img_tables_col_1 = (image_revised_last == label_table).astype(np.uint8) contours_table_col1, _ = return_contours_of_image(img_tables_col_1) _,_ ,_ , _, y_min_tab_col1 ,y_max_tab_col1, _= find_new_features_of_contours(contours_table_col1) if len(y_min_tab_col1)>0: for ijv in range(len(y_min_tab_col1)): - image_revised_last[int(y_min_tab_col1[ijv]):int(y_max_tab_col1[ijv])] = pixel_table + image_revised_last[int(y_min_tab_col1[ijv]):int(y_max_tab_col1[ijv])] = label_table return image_revised_last def get_tables_from_model(self, img, num_col_classifier): @@ -2078,18 +2049,18 @@ class Eynollah: if self.plotter: self.plotter.save_page_image(image_page) - + if not self.ignore_page_extraction: mask_page = np.zeros((text_regions_p_1.shape[0], text_regions_p_1.shape[1])).astype(np.int8) mask_page = cv2.fillPoly(mask_page, pts=[cont_page[0]], color=(1,)) - + text_regions_p_1[mask_page==0] = 0 textline_mask_tot_ea[mask_page==0] = 0 - + text_regions_p_1 = text_regions_p_1[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] textline_mask_tot_ea = textline_mask_tot_ea[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] img_bin_light = img_bin_light[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] - + ###text_regions_p_1 = text_regions_p_1[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] ###textline_mask_tot_ea = textline_mask_tot_ea[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] ###img_bin_light = img_bin_light[page_coord[0] : page_coord[1], page_coord[2] : page_coord[3]] @@ -2097,8 +2068,8 @@ class Eynollah: mask_images = (text_regions_p_1[:, :] == 2) * 1 mask_images = mask_images.astype(np.uint8) mask_images = cv2.erode(mask_images[:, :], KERNEL, iterations=10) - mask_lines = (text_regions_p_1[:, :] == 3) * 1 - mask_lines = mask_lines.astype(np.uint8) + mask_seps = (text_regions_p_1[:, :] == 3) * 1 + mask_seps = mask_seps.astype(np.uint8) img_only_regions_with_sep = ((text_regions_p_1[:, :] != 3) & (text_regions_p_1[:, :] != 0)) * 1 img_only_regions_with_sep = img_only_regions_with_sep.astype(np.uint8) @@ -2121,10 +2092,10 @@ class Eynollah: max(self.num_col_lower or num_col_classifier, num_col_classifier)) except Exception as why: - self.logger.error(why) + self.logger.exception(why) num_col = None #print("inside graphics 3 ", time.time() - t_in_gr) - return (num_col, num_col_classifier, img_only_regions, page_coord, image_page, mask_images, mask_lines, + return (num_col, num_col_classifier, img_only_regions, page_coord, image_page, mask_images, mask_seps, text_regions_p_1, cont_page, table_prediction, textline_mask_tot_ea, img_bin_light) def run_graphics_and_columns_without_layout(self, textline_mask_tot_ea, img_bin_light): @@ -2196,13 +2167,13 @@ class Eynollah: return slope_deskew def run_marginals( - self, textline_mask_tot_ea, mask_images, mask_lines, + self, textline_mask_tot_ea, mask_images, mask_seps, num_col_classifier, slope_deskew, text_regions_p_1, table_prediction): textline_mask_tot = textline_mask_tot_ea[:, :] textline_mask_tot[mask_images[:, :] == 1] = 0 - text_regions_p_1[mask_lines[:, :] == 1] = 3 + text_regions_p_1[mask_seps[:, :] == 1] = 3 text_regions_p = text_regions_p_1[:, :] text_regions_p = np.array(text_regions_p) if num_col_classifier in (1, 2): @@ -2226,12 +2197,10 @@ class Eynollah: self.logger.debug('enter run_boxes_no_full_layout') t_0_box = time.time() if np.abs(slope_deskew) >= SLOPE_THRESHOLD: - _, textline_mask_tot_d, text_regions_p_1_n, table_prediction_n = rotation_not_90_func( - image_page, textline_mask_tot, text_regions_p, table_prediction, slope_deskew) - text_regions_p_1_n = resize_image(text_regions_p_1_n, text_regions_p.shape[0], text_regions_p.shape[1]) - textline_mask_tot_d = resize_image(textline_mask_tot_d, text_regions_p.shape[0], text_regions_p.shape[1]) - table_prediction_n = resize_image(table_prediction_n, text_regions_p.shape[0], text_regions_p.shape[1]) - regions_without_separators_d = (text_regions_p_1_n[:, :] == 1) * 1 + textline_mask_tot_d = rotate_image(textline_mask_tot, slope_deskew) + text_regions_p_d = rotate_image(text_regions_p, slope_deskew) + table_prediction_n = rotate_image(table_prediction, slope_deskew) + regions_without_separators_d = (text_regions_p_d[:, :] == 1) * 1 if self.tables: regions_without_separators_d[table_prediction_n[:,:] == 1] = 1 regions_without_separators = (text_regions_p[:, :] == 1) * 1 @@ -2241,21 +2210,21 @@ class Eynollah: if self.tables: regions_without_separators[table_prediction ==1 ] = 1 if np.abs(slope_deskew) < SLOPE_THRESHOLD: - text_regions_p_1_n = None + text_regions_p_d = None textline_mask_tot_d = None regions_without_separators_d = None - pixel_lines = 3 + label_seps = 3 if np.abs(slope_deskew) < SLOPE_THRESHOLD: - _, _, matrix_of_lines_ch, splitter_y_new, _ = find_number_of_columns_in_document( - text_regions_p, num_col_classifier, self.tables, pixel_lines) + _, _, matrix_of_seps_ch, splitter_y_new, _ = find_number_of_columns_in_document( + text_regions_p, num_col_classifier, self.tables, label_seps) if np.abs(slope_deskew) >= SLOPE_THRESHOLD: - _, _, matrix_of_lines_ch_d, splitter_y_new_d, _ = find_number_of_columns_in_document( - text_regions_p_1_n, num_col_classifier, self.tables, pixel_lines) + _, _, matrix_of_seps_ch_d, splitter_y_new_d, _ = find_number_of_columns_in_document( + text_regions_p_d, num_col_classifier, self.tables, label_seps) #print(time.time()-t_0_box,'time box in 2') self.logger.info("num_col_classifier: %s", num_col_classifier) - if num_col_classifier >= 3: + if not erosion_hurts: if np.abs(slope_deskew) < SLOPE_THRESHOLD: regions_without_separators = regions_without_separators.astype(np.uint8) regions_without_separators = cv2.erode(regions_without_separators[:, :], KERNEL, iterations=6) @@ -2266,7 +2235,7 @@ class Eynollah: t1 = time.time() if np.abs(slope_deskew) < SLOPE_THRESHOLD: boxes, peaks_neg_tot_tables = return_boxes_of_images_by_order_of_reading_new( - splitter_y_new, regions_without_separators, matrix_of_lines_ch, + splitter_y_new, regions_without_separators, text_regions_p, matrix_of_seps_ch, num_col_classifier, erosion_hurts, self.tables, self.right2left) boxes_d = None self.logger.debug("len(boxes): %s", len(boxes)) @@ -2274,7 +2243,7 @@ class Eynollah: else: boxes_d, peaks_neg_tot_tables_d = return_boxes_of_images_by_order_of_reading_new( - splitter_y_new_d, regions_without_separators_d, matrix_of_lines_ch_d, + splitter_y_new_d, regions_without_separators_d, text_regions_p_d, matrix_of_seps_ch_d, num_col_classifier, erosion_hurts, self.tables, self.right2left) boxes = None self.logger.debug("len(boxes): %s", len(boxes_d)) @@ -2290,19 +2259,19 @@ class Eynollah: #img_revised_tab = text_regions_p[:, :] polygons_of_images = return_contours_of_interested_region(text_regions_p, 2) - pixel_img = 4 + label_marginalia = 4 min_area_mar = 0.00001 - marginal_mask = (text_regions_p[:,:]==pixel_img)*1 + marginal_mask = (text_regions_p[:,:]==label_marginalia)*1 marginal_mask = marginal_mask.astype('uint8') marginal_mask = cv2.dilate(marginal_mask, KERNEL, iterations=2) polygons_of_marginals = return_contours_of_interested_region(marginal_mask, 1, min_area_mar) - pixel_img = 10 - contours_tables = return_contours_of_interested_region(text_regions_p, pixel_img, min_area_mar) + label_tables = 10 + contours_tables = return_contours_of_interested_region(text_regions_p, label_tables, min_area_mar) #print(time.time()-t_0_box,'time box in 5') self.logger.debug('exit run_boxes_no_full_layout') - return (polygons_of_images, img_revised_tab, text_regions_p_1_n, textline_mask_tot_d, + return (polygons_of_images, img_revised_tab, text_regions_p_d, textline_mask_tot_d, regions_without_separators_d, boxes, boxes_d, polygons_of_marginals, contours_tables) @@ -2317,24 +2286,13 @@ class Eynollah: text_regions_p[:,:][table_prediction[:,:]==1] = 10 img_revised_tab = text_regions_p[:,:] if np.abs(slope_deskew) >= SLOPE_THRESHOLD: - _, textline_mask_tot_d, text_regions_p_1_n, table_prediction_n = \ - rotation_not_90_func(image_page, textline_mask_tot, text_regions_p, - table_prediction, slope_deskew) - - text_regions_p_1_n = resize_image(text_regions_p_1_n, - text_regions_p.shape[0], - text_regions_p.shape[1]) - textline_mask_tot_d = resize_image(textline_mask_tot_d, - text_regions_p.shape[0], - text_regions_p.shape[1]) - table_prediction_n = resize_image(table_prediction_n, - text_regions_p.shape[0], - text_regions_p.shape[1]) - - regions_without_separators_d = (text_regions_p_1_n[:,:] == 1)*1 + textline_mask_tot_d = rotate_image(textline_mask_tot, slope_deskew) + text_regions_p_d = rotate_image(text_regions_p, slope_deskew) + table_prediction_n = rotate_image(table_prediction, slope_deskew) + regions_without_separators_d = (text_regions_p_d[:,:] == 1)*1 regions_without_separators_d[table_prediction_n[:,:] == 1] = 1 else: - text_regions_p_1_n = None + text_regions_p_d = None textline_mask_tot_d = None regions_without_separators_d = None # regions_without_separators = ( text_regions_p[:,:]==1 | text_regions_p[:,:]==2 )*1 @@ -2343,17 +2301,18 @@ class Eynollah: regions_without_separators[table_prediction == 1] = 1 - pixel_img = 4 + + label_marginalia = 4 min_area_mar = 0.00001 - marginal_mask = (text_regions_p[:,:]==pixel_img)*1 + marginal_mask = (text_regions_p[:,:]==label_marginalia)*1 marginal_mask = marginal_mask.astype('uint8') marginal_mask = cv2.dilate(marginal_mask, KERNEL, iterations=2) polygons_of_marginals = return_contours_of_interested_region(marginal_mask, 1, min_area_mar) - pixel_img = 10 - contours_tables = return_contours_of_interested_region(text_regions_p, pixel_img, min_area_mar) + label_tables = 10 + contours_tables = return_contours_of_interested_region(text_regions_p, label_tables, min_area_mar) # set first model with second model text_regions_p[:, :][text_regions_p[:, :] == 2] = 5 @@ -2369,7 +2328,7 @@ class Eynollah: # 6 is the separators lable in old full layout model # 4 is the drop capital class in old full layout model # in the new full layout drop capital is 3 and separators are 5 - + # the separators in full layout will not be written on layout if not self.reading_order_machine_based: text_regions_p[:,:][regions_fully[:,:,0]==5]=6 @@ -2406,16 +2365,13 @@ class Eynollah: #plt.show() ####if not self.tables: if np.abs(slope_deskew) >= SLOPE_THRESHOLD: - _, textline_mask_tot_d, text_regions_p_1_n, regions_fully_n = rotation_not_90_func_full_layout( - image_page, textline_mask_tot, text_regions_p, regions_fully, slope_deskew) - - text_regions_p_1_n = resize_image(text_regions_p_1_n, text_regions_p.shape[0], text_regions_p.shape[1]) - textline_mask_tot_d = resize_image(textline_mask_tot_d, text_regions_p.shape[0], text_regions_p.shape[1]) - regions_fully_n = resize_image(regions_fully_n, text_regions_p.shape[0], text_regions_p.shape[1]) + textline_mask_tot_d = rotate_image(textline_mask_tot, slope_deskew) + text_regions_p_d = rotate_image(text_regions_p, slope_deskew) + regions_fully_n = rotate_image(regions_fully, slope_deskew) if not self.tables: - regions_without_separators_d = (text_regions_p_1_n[:, :] == 1) * 1 + regions_without_separators_d = (text_regions_p_d[:, :] == 1) * 1 else: - text_regions_p_1_n = None + text_regions_p_d = None textline_mask_tot_d = None regions_without_separators_d = None if not self.tables: @@ -2425,12 +2381,12 @@ class Eynollah: self.logger.debug('exit run_boxes_full_layout') #print("full inside 3", time.time()- t_full0) - return (polygons_of_images, img_revised_tab, text_regions_p_1_n, textline_mask_tot_d, + return (polygons_of_images, img_revised_tab, text_regions_p_d, textline_mask_tot_d, regions_without_separators_d, regions_fully, regions_without_separators, polygons_of_marginals, contours_tables) 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 @@ -2439,13 +2395,13 @@ class Eynollah: 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, @@ -2453,19 +2409,19 @@ class Eynollah: 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] #contours_only_text_parent[diff_x_ratio>=1.3] for ind in range(len(contours_only_text_parent)) @@ -2474,7 +2430,7 @@ class Eynollah: #contours_only_text_parent[diff_x_ratio<1.3] for ind in range(len(contours_only_text_parent)) if diff_x_ratio[ind]<1.3] - + cx_conts_excluded = [cx_conts[ind] #cx_conts[diff_x_ratio>=1.3] for ind in range(len(cx_conts)) @@ -2491,43 +2447,43 @@ class Eynollah: #cy_conts[diff_x_ratio<1.3] for ind in range(len(cy_conts)) if diff_x_ratio[ind]<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_located_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(range(len(contours_only_text_parent))) - set(flattened_array) ) #print(missing_textregions, 'missing_textregions') @@ -2535,15 +2491,15 @@ class Eynollah: 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] @@ -2552,7 +2508,7 @@ class Eynollah: img_poly[text_regions_p[:,:]==2] = 2 img_poly[text_regions_p[:,:]==3] = 4 img_poly[text_regions_p[:,:]==6] = 5 - + 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( @@ -2573,15 +2529,15 @@ class Eynollah: co_text_all = contours_only_text_parent if not len(co_text_all): - return [], [] + 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 @@ -2589,9 +2545,9 @@ class Eynollah: 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 @@ -2641,7 +2597,7 @@ class Eynollah: break ordered = [i[0] for i in ordered] - + if len(contours_only_text_parent)>min_cont_size_to_be_dilated: org_contours_indexes = [] for ind in range(len(ordered)): @@ -2655,12 +2611,10 @@ class Eynollah: np.array(indexes_of_located_cont[region_with_curr_order])[arg_sort_located_cont]) else: org_contours_indexes.extend([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 + + return org_contours_indexes else: - region_ids = ['region_%04d' % i for i in range(len(co_text_all_org))] - return ordered, region_ids + return ordered def filter_contours_inside_a_bigger_one(self, contours, contours_d_ordered, image, marginal_cnts=None, type_contour="textregion"): @@ -2741,7 +2695,7 @@ class Eynollah: axis=0)) return contours - + def return_indexes_of_contours_located_inside_another_list_of_contours( self, contours, contours_loc, cx_main_loc, cy_main_loc, indexes_loc): indexes_of_located_cont = [] @@ -2751,7 +2705,7 @@ class Eynollah: #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)) ] @@ -2763,9 +2717,9 @@ class Eynollah: 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 filter_contours_without_textline_inside( self, contours_par, contours_textline, @@ -2778,7 +2732,7 @@ class Eynollah: def filterfun(lis): if len(lis) == 0: return [] - return list(np.array(lis)[indices]) + return list(ensure_array(lis)[indices]) return (filterfun(contours_par), filterfun(contours_textline), @@ -2806,15 +2760,15 @@ class Eynollah: (all_found_textline_polygons_marginals_left, all_found_textline_polygons_marginals_right) = \ split(all_found_textline_polygons_marginals) - + (all_box_coord_marginals_left, all_box_coord_marginals_right) = \ split(all_box_coord_marginals) - + (slopes_marg_left, slopes_marg_right) = \ split(slopes_marginals) - + (cy_marg_left, cy_marg_right) = \ split(cy_marg) @@ -2825,19 +2779,19 @@ class Eynollah: return list(np.array(lis)[order_left]) def sort_right(lis): return list(np.array(lis)[order_right]) - + ordered_left_marginals = sort_left(poly_marg_left) ordered_right_marginals = sort_right(poly_marg_right) - + ordered_left_marginals_textline = sort_left(all_found_textline_polygons_marginals_left) ordered_right_marginals_textline = sort_right(all_found_textline_polygons_marginals_right) - + ordered_left_marginals_bbox = sort_left(all_box_coord_marginals_left) ordered_right_marginals_bbox = sort_right(all_box_coord_marginals_right) - + ordered_left_slopes_marginals = sort_left(slopes_marg_left) ordered_right_slopes_marginals = sort_right(slopes_marg_right) - + return (ordered_left_marginals, ordered_right_marginals, ordered_left_marginals_textline, @@ -2920,26 +2874,27 @@ class Eynollah: def run_single(self): t0 = time.time() - + self.logger.info(f"Processing file: {self.writer.image_filename}") self.logger.info("Step 1/5: Image Enhancement") - + img_res, is_image_enhanced, num_col_classifier, num_column_is_classified = \ self.run_enhancement() - + self.logger.info(f"Image: {self.image.shape[1]}x{self.image.shape[0]}, " + f"scale {self.scale_x:.1f}x{self.scale_y:.1f}, " 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)") - + # Basic Processing Mode if self.skip_layout_and_reading_order: self.logger.info("Step 2/5: Basic Processing Mode") self.logger.info("Skipping layout analysis and reading order detection") - + _ ,_, _, _, textline_mask_tot_ea, img_bin_light, _ = \ self.get_regions_light_v(img_res, is_image_enhanced, num_col_classifier,) @@ -2951,7 +2906,7 @@ class Eynollah: cnt_clean_rot_raw, hir_on_cnt_clean_rot = return_contours_of_image(textline_mask_tot_ea) all_found_textline_polygons = filter_contours_area_of_image( textline_mask_tot_ea, cnt_clean_rot_raw, hir_on_cnt_clean_rot, max_area=1, min_area=0.00001) - + cx_main_tot, cy_main_tot = find_center_of_contours(all_found_textline_polygons) w_h_textlines = [cv2.boundingRect(polygon)[2:] for polygon in all_found_textline_polygons] @@ -2964,30 +2919,31 @@ class Eynollah: 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, textline_mask_tot_ea, type_contour="textline") - + order_text_new = [0] slopes =[0] conf_contours_textregions =[0] - + pcgts = self.writer.build_pagexml_no_full_layout( - found_polygons_text_region=cont_page, - page_coord=page_coord, - order_of_texts=order_text_new, - all_found_textline_polygons=all_found_textline_polygons, - all_box_coord=page_coord, - found_polygons_text_region_img=[], - 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, - slopes_marginals_left=[], - slopes_marginals_right=[], - cont_page=cont_page, - polygons_seplines=[], - found_polygons_tables=[], + found_polygons_text_region=cont_page, + page_coord=page_coord, + order_of_texts=order_text_new, + all_found_textline_polygons=all_found_textline_polygons, + all_box_coord=page_coord, + found_polygons_text_region_img=[], + 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, + slopes_marginals_left=[], + slopes_marginals_right=[], + cont_page=cont_page, + polygons_seplines=[], + found_polygons_tables=[], + skip_layout_reading_order=True ) self.logger.info("Basic processing complete") return pcgts @@ -2995,7 +2951,7 @@ class Eynollah: #print("text region early -1 in %.1fs", time.time() - t0) t1 = time.time() self.logger.info("Step 2/5: Layout Analysis") - + self.logger.info("Using light version processing") text_regions_p_1 ,erosion_hurts, polygons_seplines, polygons_text_early, \ textline_mask_tot_ea, img_bin_light, confidence_matrix = \ @@ -3015,7 +2971,7 @@ class Eynollah: slope_deskew = self.run_deskew(textline_mask_tot_ea) #print("text region early -2,5 in %.1fs", time.time() - t0) #self.logger.info("Textregion detection took %.1fs ", time.time() - t1t) - num_col, num_col_classifier, img_only_regions, page_coord, image_page, mask_images, mask_lines, \ + num_col, num_col_classifier, img_only_regions, page_coord, image_page, mask_images, mask_seps, \ text_regions_p_1, cont_page, table_prediction, textline_mask_tot_ea, img_bin_light = \ self.run_graphics_and_columns_light(text_regions_p_1, textline_mask_tot_ea, num_col_classifier, num_column_is_classified, @@ -3030,7 +2986,7 @@ class Eynollah: if not num_col and len(polygons_text_early) == 0: self.logger.info("No columns detected - generating empty PAGE-XML") - + pcgts = self.writer.build_pagexml_no_full_layout( found_polygons_text_region=[], page_coord=page_coord, @@ -3067,12 +3023,12 @@ class Eynollah: image_page = resize_image(image_page,img_h_new, img_w_new ) textline_mask_tot_ea = resize_image(textline_mask_tot_ea,img_h_new, img_w_new ) mask_images = resize_image(mask_images,img_h_new, img_w_new ) - mask_lines = resize_image(mask_lines,img_h_new, img_w_new ) + mask_seps = resize_image(mask_seps, img_h_new, img_w_new) text_regions_p_1 = resize_image(text_regions_p_1,img_h_new, img_w_new ) table_prediction = resize_image(table_prediction,img_h_new, img_w_new ) textline_mask_tot, text_regions_p = \ - self.run_marginals(textline_mask_tot_ea, mask_images, mask_lines, + self.run_marginals(textline_mask_tot_ea, mask_images, mask_seps, num_col_classifier, slope_deskew, text_regions_p_1, table_prediction) if self.plotter: self.plotter.save_plot_of_layout_main_all(text_regions_p, image_page) @@ -3091,7 +3047,7 @@ class Eynollah: text_regions_p[text_regions_p == 4] = 1 self.logger.info("Step 3/5: Text Line Detection") - + if self.curved_line: self.logger.info("Mode: Curved line detection") @@ -3107,14 +3063,14 @@ class Eynollah: ## birdan sora chock chakir t1 = time.time() if not self.full_layout: - polygons_of_images, img_revised_tab, text_regions_p_1_n, \ + polygons_of_images, img_revised_tab, text_regions_p_d, \ textline_mask_tot_d, regions_without_separators_d, \ boxes, boxes_d, polygons_of_marginals, contours_tables = \ self.run_boxes_no_full_layout(image_page, textline_mask_tot, text_regions_p, slope_deskew, num_col_classifier, table_prediction, erosion_hurts) ###polygons_of_marginals = dilate_textregion_contours(polygons_of_marginals) else: - polygons_of_images, img_revised_tab, text_regions_p_1_n, \ + polygons_of_images, img_revised_tab, text_regions_p_d, \ textline_mask_tot_d, regions_without_separators_d, \ regions_fully, regions_without_separators, polygons_of_marginals, contours_tables = \ self.run_boxes_full_layout(image_page, textline_mask_tot, text_regions_p, slope_deskew, @@ -3127,7 +3083,7 @@ class Eynollah: text_only = (img_revised_tab[:, :] == 1) * 1 if np.abs(slope_deskew) >= SLOPE_THRESHOLD: - text_only_d = (text_regions_p_1_n[:, :] == 1) * 1 + text_only_d = ((text_regions_p_d[:, :] == 1)) * 1 #print("text region early 2 in %.1fs", time.time() - t0) ###min_con_area = 0.000005 @@ -3141,7 +3097,8 @@ class Eynollah: areas_cnt_text = np.array([cv2.contourArea(c) for c in contours_only_text_parent]) areas_cnt_text = areas_cnt_text / float(areas_tot_text) #self.logger.info('areas_cnt_text %s', areas_cnt_text) - contours_only_text_parent = np.array(contours_only_text_parent)[areas_cnt_text > MIN_AREA_REGION] + contours_only_text_parent = ensure_array(contours_only_text_parent) + contours_only_text_parent = contours_only_text_parent[areas_cnt_text > MIN_AREA_REGION] areas_cnt_text_parent = areas_cnt_text[areas_cnt_text > MIN_AREA_REGION] index_con_parents = np.argsort(areas_cnt_text_parent) @@ -3160,12 +3117,13 @@ class Eynollah: areas_cnt_text_d = np.array([cv2.contourArea(c) for c in contours_only_text_parent_d]) areas_cnt_text_d = areas_cnt_text_d / float(areas_tot_text_d) - contours_only_text_parent_d = np.array(contours_only_text_parent_d)[areas_cnt_text_d > MIN_AREA_REGION] + contours_only_text_parent_d = ensure_array(contours_only_text_parent_d) + contours_only_text_parent_d = contours_only_text_parent_d[areas_cnt_text_d > MIN_AREA_REGION] areas_cnt_text_d = areas_cnt_text_d[areas_cnt_text_d > MIN_AREA_REGION] if len(contours_only_text_parent_d): index_con_parents_d = np.argsort(areas_cnt_text_d) - contours_only_text_parent_d = np.array(contours_only_text_parent_d)[index_con_parents_d] + contours_only_text_parent_d = contours_only_text_parent_d[index_con_parents_d] areas_cnt_text_d = areas_cnt_text_d[index_con_parents_d] centers_d = np.stack(find_center_of_contours(contours_only_text_parent_d)) # [2, N] @@ -3201,42 +3159,42 @@ class Eynollah: dists[i] = np.linalg.norm(centers[:, i:i + 1] - centers_d, axis=0) corresp = np.zeros(dists.shape, dtype=bool) # keep searching next-closest until at least one correspondence on each side - while not np.all(corresp.sum(axis=1)) and not np.all(corresp.sum(axis=0)): + while not np.all(corresp.sum(axis=1)) or not np.all(corresp.sum(axis=0)): idx = np.nanargmin(dists) i, j = np.unravel_index(idx, dists.shape) dists[i, j] = np.nan corresp[i, j] = True - #print("original/deskewed adjacency", corresp.nonzero()) + # print("original/deskewed adjacency", corresp.nonzero()) contours_only_text_parent_d_ordered = np.zeros_like(contours_only_text_parent) contours_only_text_parent_d_ordered = contours_only_text_parent_d[np.argmax(corresp, axis=1)] # img1 = np.zeros(text_only_d.shape[:2], dtype=np.uint8) # for i in range(len(contours_only_text_parent)): # cv2.fillPoly(img1, pts=[contours_only_text_parent_d_ordered[i]], color=i + 1) - # plt.subplot(2, 2, 1, title="direct corresp contours") + # plt.subplot(1, 4, 1, title="direct corresp contours") # plt.imshow(img1) # img2 = np.zeros(text_only_d.shape[:2], dtype=np.uint8) # join deskewed regions mapping to single original ones for i in range(len(contours_only_text_parent)): if np.count_nonzero(corresp[i]) > 1: indices = np.flatnonzero(corresp[i]) - #print("joining", indices) + # print("joining", indices) polygons_d = [contour2polygon(contour) for contour in contours_only_text_parent_d[indices]] contour_d = polygon2contour(join_polygons(polygons_d)) contours_only_text_parent_d_ordered[i] = contour_d # cv2.fillPoly(img2, pts=[contour_d], color=i + 1) - # plt.subplot(2, 2, 3, title="joined contours") + # plt.subplot(1, 4, 2, title="joined contours") # plt.imshow(img2) # img3 = np.zeros(text_only_d.shape[:2], dtype=np.uint8) # split deskewed regions mapping to multiple original ones def deskew(polygon): polygon = shapely.affinity.rotate(polygon, -slope_deskew, origin=center) - polygon = shapely.affinity.translate(polygon, *offset.squeeze()) + #polygon = shapely.affinity.translate(polygon, *offset.squeeze()) return polygon for j in range(len(contours_only_text_parent_d)): if np.count_nonzero(corresp[:, j]) > 1: indices = np.flatnonzero(corresp[:, j]) - #print("splitting along", indices) + # print("splitting along", indices) polygons = [deskew(contour2polygon(contour)) for contour in contours_only_text_parent[indices]] polygon_d = contour2polygon(contours_only_text_parent_d[j]) @@ -3249,14 +3207,38 @@ class Eynollah: if polygon_d] contours_only_text_parent_d_ordered[indices] = contours_d # cv2.fillPoly(img3, pts=contours_d, color=j + 1) - # plt.subplot(2, 2, 4, title="split contours") + # plt.subplot(1, 4, 3, title="split contours") # plt.imshow(img3) # img4 = np.zeros(text_only_d.shape[:2], dtype=np.uint8) # for i in range(len(contours_only_text_parent)): # cv2.fillPoly(img4, pts=[contours_only_text_parent_d_ordered[i]], color=i + 1) - # plt.subplot(2, 2, 2, title="result contours") + # plt.subplot(1, 4, 4, title="result contours") # plt.imshow(img4) # plt.show() + # from matplotlib import patches as ptchs + # plt.subplot(1, 2, 1, title="undeskewed") + # plt.imshow(text_only) + # centers = np.stack(find_center_of_contours(contours_only_text_parent)) # [2, N] + # for i in range(len(contours_only_text_parent)): + # cnt = contours_only_text_parent[i] + # ctr = centers[:, i] + # plt.gca().add_patch(ptchs.Polygon(cnt[:, 0], closed=False, fill=False, color='blue')) + # plt.gca().scatter(ctr[0], ctr[1], 20, c='blue', marker='x') + # plt.gca().text(ctr[0], ctr[1], str(i), c='blue') + # plt.subplot(1, 2, 2, title="deskewed") + # plt.imshow(text_only_d) + # centers_d = np.stack(find_center_of_contours(contours_only_text_parent_d_ordered)) # [2, N] + # for i in range(len(contours_only_text_parent)): + # cnt = contours_only_text_parent[i] + # cnt = polygon2contour(deskew(contour2polygon(cnt))) + # plt.gca().add_patch(ptchs.Polygon(cnt[:, 0], closed=False, fill=False, color='blue')) + # for i in range(len(contours_only_text_parent_d_ordered)): + # cnt = contours_only_text_parent_d_ordered[i] + # ctr = centers_d[:, i] + # plt.gca().add_patch(ptchs.Polygon(cnt[:, 0], closed=False, fill=False, color='red')) + # plt.gca().scatter(ctr[0], ctr[1], 20, c='red', marker='x') + # plt.gca().text(ctr[0], ctr[1], str(i), c='red') + # plt.show() if not len(contours_only_text_parent): # stop early @@ -3363,7 +3345,7 @@ class Eynollah: num_col_classifier, scale_param, slope_deskew) all_found_textline_polygons_marginals = small_textlines_to_parent_adherence2( all_found_textline_polygons_marginals, textline_mask_tot_ea, num_col_classifier) - + mid_point_of_page_width = text_regions_p.shape[1] / 2. (polygons_of_marginals_left, polygons_of_marginals_right, all_found_textline_polygons_marginals_left, all_found_textline_polygons_marginals_right, @@ -3372,7 +3354,7 @@ class Eynollah: self.separate_marginals_to_left_and_right_and_order_from_top_to_down( polygons_of_marginals, all_found_textline_polygons_marginals, all_box_coord_marginals, slopes_marginals, mid_point_of_page_width) - + #print(len(polygons_of_marginals), len(ordered_left_marginals), len(ordered_right_marginals), 'marginals ordred') if self.full_layout: @@ -3401,20 +3383,20 @@ class Eynollah: label_seps = 6 if not self.headers_off: if np.abs(slope_deskew) < SLOPE_THRESHOLD: - num_col, _, matrix_of_lines_ch, splitter_y_new, _ = find_number_of_columns_in_document( + num_col, _, matrix_of_seps_ch, splitter_y_new, _ = find_number_of_columns_in_document( text_regions_p, num_col_classifier, self.tables, label_seps, contours_only_text_parent_h) else: - _, _, matrix_of_lines_ch_d, splitter_y_new_d, _ = find_number_of_columns_in_document( - text_regions_p_1_n, num_col_classifier, self.tables, label_seps, contours_only_text_parent_h_d_ordered) + _, _, matrix_of_seps_ch_d, splitter_y_new_d, _ = find_number_of_columns_in_document( + text_regions_p_d, num_col_classifier, self.tables, label_seps, contours_only_text_parent_h_d_ordered) elif self.headers_off: if np.abs(slope_deskew) < SLOPE_THRESHOLD: - num_col, _, matrix_of_lines_ch, splitter_y_new, _ = find_number_of_columns_in_document( + num_col, _, matrix_of_seps_ch, splitter_y_new, _ = find_number_of_columns_in_document( text_regions_p, num_col_classifier, self.tables, label_seps) else: - _, _, matrix_of_lines_ch_d, splitter_y_new_d, _ = find_number_of_columns_in_document( - text_regions_p_1_n, num_col_classifier, self.tables, label_seps) + _, _, matrix_of_seps_ch_d, splitter_y_new_d, _ = find_number_of_columns_in_document( + text_regions_p_d, num_col_classifier, self.tables, label_seps) - if num_col_classifier >= 3: + if not erosion_hurts: if np.abs(slope_deskew) < SLOPE_THRESHOLD: regions_without_separators = regions_without_separators.astype(np.uint8) regions_without_separators = cv2.erode(regions_without_separators[:, :], KERNEL, iterations=6) @@ -3423,13 +3405,13 @@ class Eynollah: regions_without_separators_d = cv2.erode(regions_without_separators_d[:, :], KERNEL, iterations=6) if np.abs(slope_deskew) < SLOPE_THRESHOLD: - boxes, peaks_neg_tot_tables = return_boxes_of_images_by_order_of_reading_new( - splitter_y_new, regions_without_separators, matrix_of_lines_ch, + boxes, _ = return_boxes_of_images_by_order_of_reading_new( + splitter_y_new, regions_without_separators, text_regions_p, matrix_of_seps_ch, num_col_classifier, erosion_hurts, self.tables, self.right2left, logger=self.logger) else: - boxes_d, peaks_neg_tot_tables_d = return_boxes_of_images_by_order_of_reading_new( - splitter_y_new_d, regions_without_separators_d, matrix_of_lines_ch_d, + boxes_d, _ = return_boxes_of_images_by_order_of_reading_new( + splitter_y_new_d, regions_without_separators_d, text_regions_p_d, matrix_of_seps_ch_d, num_col_classifier, erosion_hurts, self.tables, self.right2left, logger=self.logger) else: @@ -3450,14 +3432,14 @@ class Eynollah: self.logger.info("Headers ignored in reading order") if self.reading_order_machine_based: - order_text_new, id_of_texts_tot = self.do_order_of_regions_with_model( + order_text_new = self.do_order_of_regions_with_model( contours_only_text_parent, contours_only_text_parent_h, text_regions_p) else: if np.abs(slope_deskew) < SLOPE_THRESHOLD: - order_text_new, id_of_texts_tot = self.do_order_of_regions( + order_text_new = self.do_order_of_regions( contours_only_text_parent, contours_only_text_parent_h, boxes, textline_mask_tot) else: - order_text_new, id_of_texts_tot = self.do_order_of_regions( + order_text_new = self.do_order_of_regions( contours_only_text_parent_d_ordered, contours_only_text_parent_h_d_ordered, boxes_d, textline_mask_tot_d) self.logger.info(f"Detection of reading order took {time.time() - t_order:.1f}s") @@ -3513,5 +3495,5 @@ class Eynollah: polygons_seplines=polygons_seplines, found_polygons_tables=contours_tables, ) - + return pcgts diff --git a/src/eynollah/eynollah_imports.py b/src/eynollah/eynollah_imports.py index f04cfdc..496406c 100644 --- a/src/eynollah/eynollah_imports.py +++ b/src/eynollah/eynollah_imports.py @@ -1,6 +1,9 @@ """ Load libraries with possible race conditions once. This must be imported as the first module of eynollah. """ +import os +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 + from ocrd_utils import tf_disable_interactive_logs from torch import * tf_disable_interactive_logs() diff --git a/src/eynollah/image_enhancer.py b/src/eynollah/image_enhancer.py index babbd55..67145a3 100644 --- a/src/eynollah/image_enhancer.py +++ b/src/eynollah/image_enhancer.py @@ -15,11 +15,13 @@ from pathlib import Path import gc import cv2 -from keras.models import Model import numpy as np -import tensorflow as tf # type: ignore from skimage.morphology import skeletonize +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 +import tensorflow as tf # type: ignore +from tensorflow.keras.models import Model + from .model_zoo import EynollahModelZoo from .utils.resize import resize_image from .utils.pil_cv2 import pil2cv diff --git a/src/eynollah/mb_ro_on_layout.py b/src/eynollah/mb_ro_on_layout.py index eec544c..22fe97b 100644 --- a/src/eynollah/mb_ro_on_layout.py +++ b/src/eynollah/mb_ro_on_layout.py @@ -14,10 +14,12 @@ from pathlib import Path import xml.etree.ElementTree as ET import cv2 -from keras.models import Model import numpy as np import statistics + +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 import tensorflow as tf +from tensorflow.keras.models import Model from .model_zoo import EynollahModelZoo from .utils.resize import resize_image diff --git a/src/eynollah/model_zoo/model_zoo.py b/src/eynollah/model_zoo/model_zoo.py index 80d0aa7..83068ff 100644 --- a/src/eynollah/model_zoo/model_zoo.py +++ b/src/eynollah/model_zoo/model_zoo.py @@ -1,16 +1,19 @@ +import os import json import logging from copy import deepcopy from pathlib import Path from typing import Dict, List, Optional, Tuple, Type, Union +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() -from keras.layers import StringLookup -from keras.models import Model as KerasModel -from keras.models import load_model +from tensorflow.keras.layers import StringLookup +from tensorflow.keras.models import Model as KerasModel +from tensorflow.keras.models import load_model from tabulate import tabulate + from ..patch_encoder import PatchEncoder, Patches from .specs import EynollahModelSpecSet from .default_specs import DEFAULT_MODEL_SPECS diff --git a/src/eynollah/ocrd-tool.json b/src/eynollah/ocrd-tool.json index 3b500fc..fc61af7 100644 --- a/src/eynollah/ocrd-tool.json +++ b/src/eynollah/ocrd-tool.json @@ -28,7 +28,19 @@ "full_layout": { "type": "boolean", "default": true, - "description": "Try to detect all element subtypes, including drop-caps and headings" + "description": "Try to detect all region subtypes, including drop-capital and heading" + }, + "light_version": { + "type": "boolean", + "default": true, + "enum": [true], + "description": "ignored (only for backwards-compatibility)" + }, + "textline_light": { + "type": "boolean", + "default": true, + "enum": [true], + "description": "ignored (only for backwards-compatibility)" }, "tables": { "type": "boolean", @@ -38,12 +50,12 @@ "curved_line": { "type": "boolean", "default": false, - "description": "try to return contour of textlines instead of just rectangle bounding box. Needs more processing time" + "description": "retrieve textline polygons independent of each other (needs more processing time)" }, "ignore_page_extraction": { "type": "boolean", "default": false, - "description": "if this parameter set to true, this tool would ignore page extraction" + "description": "if true, do not attempt page frame detection (cropping)" }, "allow_scaling": { "type": "boolean", @@ -58,7 +70,7 @@ "right_to_left": { "type": "boolean", "default": false, - "description": "if this parameter set to true, this tool will extract right-to-left reading order." + "description": "if true, return reading order in right-to-left reading direction." }, "headers_off": { "type": "boolean", @@ -123,13 +135,22 @@ } }, "resources": [ + { + "url": "https://zenodo.org/records/17580627/files/models_all_v0_7_0.zip?download=1", + "name": "models_layout_v0_7_0", + "type": "archive", + "size": 6119874002, + "description": "Models for layout detection, reading order detection, textline detection, page extraction, column classification, table detection, binarization, image enhancement and OCR", + "version_range": ">= v0.7.0" + }, { "url": "https://github.com/qurator-spk/sbb_binarization/releases/download/v0.0.11/saved_model_2020_01_16.zip", "name": "default", "type": "archive", "path_in_archive": "saved_model_2020_01_16", "size": 563147331, - "description": "default models provided by github.com/qurator-spk (SavedModel format)" + "description": "default models provided by github.com/qurator-spk (SavedModel format)", + "version_range": "< v0.7.0" }, { "url": "https://github.com/qurator-spk/sbb_binarization/releases/download/v0.0.11/saved_model_2021_03_09.zip", @@ -137,7 +158,8 @@ "type": "archive", "path_in_archive": ".", "size": 133230419, - "description": "updated default models provided by github.com/qurator-spk (SavedModel format)" + "description": "updated default models provided by github.com/qurator-spk (SavedModel format)", + "version_range": "< v0.7.0" } ] } diff --git a/src/eynollah/ocrd_cli_binarization.py b/src/eynollah/ocrd_cli_binarization.py index f234520..e9059df 100644 --- a/src/eynollah/ocrd_cli_binarization.py +++ b/src/eynollah/ocrd_cli_binarization.py @@ -75,7 +75,7 @@ class SbbBinarizeProcessor(Processor): if oplevel == 'page': self.logger.info("Binarizing on 'page' level in page '%s'", page_id) - page_image_bin = cv2pil(self.binarizer.run(image=pil2cv(page_image), use_patches=True)) + page_image_bin = cv2pil(self.binarizer.run_single(image=pil2cv(page_image), use_patches=True)) # update PAGE (reference the image file): page_image_ref = AlternativeImageType(comments=page_xywh['features'] + ',binarized,clipped') page.add_AlternativeImage(page_image_ref) @@ -88,7 +88,7 @@ class SbbBinarizeProcessor(Processor): for region in regions: region_image, region_xywh = self.workspace.image_from_segment( region, page_image, page_xywh, feature_filter='binarized') - region_image_bin = cv2pil(self.binarizer.run(image=pil2cv(region_image), use_patches=True)) + region_image_bin = cv2pil(self.binarizer.run_single(image=pil2cv(region_image), use_patches=True)) # update PAGE (reference the image file): region_image_ref = AlternativeImageType(comments=region_xywh['features'] + ',binarized') region.add_AlternativeImage(region_image_ref) @@ -100,7 +100,7 @@ class SbbBinarizeProcessor(Processor): self.logger.warning("Page '%s' contains no text lines", page_id) for line in lines: line_image, line_xywh = self.workspace.image_from_segment(line, page_image, page_xywh, feature_filter='binarized') - line_image_bin = cv2pil(self.binarizer.run(image=pil2cv(line_image), use_patches=True)) + line_image_bin = cv2pil(self.binarizer.run_single(image=pil2cv(line_image), use_patches=True)) # update PAGE (reference the image file): line_image_ref = AlternativeImageType(comments=line_xywh['features'] + ',binarized') line.add_AlternativeImage(line_image_ref) diff --git a/src/eynollah/patch_encoder.py b/src/eynollah/patch_encoder.py index 939ad7b..07b843d 100644 --- a/src/eynollah/patch_encoder.py +++ b/src/eynollah/patch_encoder.py @@ -1,52 +1,46 @@ -from keras import layers +import os +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 import tensorflow as tf - -projection_dim = 64 -patch_size = 1 -num_patches =21*21#14*14#28*28#14*14#28*28 +from tensorflow.keras import layers class PatchEncoder(layers.Layer): - def __init__(self): + # 441=21*21 # 14*14 # 28*28 + def __init__(self, num_patches=441, projection_dim=64): super().__init__() - self.projection = layers.Dense(units=projection_dim) - self.position_embedding = layers.Embedding(input_dim=num_patches, output_dim=projection_dim) + self.num_patches = num_patches + self.projection_dim = projection_dim + self.projection = layers.Dense(self.projection_dim) + self.position_embedding = layers.Embedding(self.num_patches, self.projection_dim) def call(self, patch): - positions = tf.range(start=0, limit=num_patches, delta=1) - encoded = self.projection(patch) + self.position_embedding(positions) - return encoded + positions = tf.range(start=0, limit=self.num_patches, delta=1) + return self.projection(patch) + self.position_embedding(positions) def get_config(self): - config = super().get_config().copy() - config.update({ - 'num_patches': num_patches, - 'projection': self.projection, - 'position_embedding': self.position_embedding, - }) - return config + return dict(num_patches=self.num_patches, + projection_dim=self.projection_dim, + **super().get_config()) class Patches(layers.Layer): - def __init__(self, **kwargs): - super(Patches, self).__init__() - self.patch_size = patch_size + def __init__(self, patch_size_x=1, patch_size_y=1): + super().__init__() + self.patch_size_x = patch_size_x + self.patch_size_y = patch_size_y def call(self, images): batch_size = tf.shape(images)[0] patches = tf.image.extract_patches( images=images, - sizes=[1, self.patch_size, self.patch_size, 1], - strides=[1, self.patch_size, self.patch_size, 1], + sizes=[1, self.patch_size_y, self.patch_size_x, 1], + strides=[1, self.patch_size_y, self.patch_size_x, 1], rates=[1, 1, 1, 1], padding="VALID", ) patch_dims = patches.shape[-1] - patches = tf.reshape(patches, [batch_size, -1, patch_dims]) - return patches - def get_config(self): + return tf.reshape(patches, [batch_size, -1, patch_dims]) - config = super().get_config().copy() - config.update({ - 'patch_size': self.patch_size, - }) - return config + def get_config(self): + return dict(patch_size_x=self.patch_size_x, + patch_size_y=self.patch_size_y, + **super().get_config()) diff --git a/src/eynollah/sbb_binarize.py b/src/eynollah/sbb_binarize.py index 37ac7c3..fe044c9 100644 --- a/src/eynollah/sbb_binarize.py +++ b/src/eynollah/sbb_binarize.py @@ -9,17 +9,18 @@ Tool to load model and binarize a given image. import os import logging +from pathlib import Path from typing import Optional import numpy as np import cv2 -from ocrd_utils import tf_disable_interactive_logs -from eynollah.model_zoo import EynollahModelZoo +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() import tensorflow as tf -from tensorflow.python.keras import backend as tensorflow_backend -from pathlib import Path + +from .model_zoo import EynollahModelZoo from .utils import is_image_filename def resize_image(img_in, input_height, input_width): @@ -34,21 +35,13 @@ class SbbBinarizer: logger: Optional[logging.Logger] = None, ): self.logger = logger if logger else logging.getLogger('eynollah.binarization') + try: + for device in tf.config.list_physical_devices('GPU'): + tf.config.experimental.set_memory_growth(device, True) + except: + self.logger.warning("no GPU device available") self.models = (model_zoo.model_path('binarization'), model_zoo.load_model('binarization')) - self.session = self.start_new_session() - - def start_new_session(self): - config = tf.compat.v1.ConfigProto() - config.gpu_options.allow_growth = True - - session = tf.compat.v1.Session(config=config) # tf.InteractiveSession() - tensorflow_backend.set_session(session) - return session - - def end_session(self): - tensorflow_backend.clear_session() - self.session.close() - del self.session + self.logger.info('Loaded model %s [%s]', self.models[1], self.models[0]) def predict(self, model, img, use_patches, n_batch_inference=5): model_height = model.layers[len(model.layers)-1].output_shape[1] @@ -311,34 +304,20 @@ class SbbBinarizer: prediction_true = prediction_true.astype(np.uint8) return prediction_true[:,:,0] - def run(self, image=None, image_path=None, output=None, use_patches=False, dir_in=None): - # print(dir_in,'dir_in') + def run(self, image=None, image_path=None, output=None, use_patches=False, dir_in=None, overwrite=False): if not dir_in: - if (image is not None and image_path is not None) or \ - (image is None and image_path is None): + if (image is None) == (image_path is None): raise ValueError("Must pass either a opencv2 image or an image_path") if image_path is not None: image = cv2.imread(image_path) - img_last = 0 - model_file, model = self.models - self.logger.info('Predicting %s with model %s', image_path if image_path else '[image]', model_file) - res = self.predict(model, image, use_patches) - - img_fin = np.zeros((res.shape[0], res.shape[1], 3)) - res[:, :][res[:, :] == 0] = 2 - res = res - 1 - res = res * 255 - img_fin[:, :, 0] = res - img_fin[:, :, 1] = res - img_fin[:, :, 2] = res - - img_fin = img_fin.astype(np.uint8) - img_fin = (res[:, :] == 0) * 255 - img_last = img_last + img_fin - - img_last[:, :][img_last[:, :] > 0] = 255 - img_last = (img_last[:, :] == 0) * 255 + img_last = self.run_single(image, use_patches) if output: + if os.path.exists(output): + if overwrite: + self.logger.warning("will overwrite existing output file '%s'", output) + else: + self.logger.warning("output file already exists '%s'", output) + return img_last self.logger.info('Writing binarized image to %s', output) cv2.imwrite(output, img_last) return img_last @@ -346,29 +325,38 @@ class SbbBinarizer: ls_imgs = list(filter(is_image_filename, os.listdir(dir_in))) self.logger.info("Found %d image files to binarize in %s", len(ls_imgs), dir_in) for i, image_path in enumerate(ls_imgs): + image_stem = os.path.splitext(image_path)[0] + output_path = os.path.join(output, image_stem + '.png') + if os.path.exists(output_path): + if overwrite: + self.logger.warning("will overwrite existing output file '%s'", output_path) + else: + self.logger.warning("will skip input for existing output file '%s'", output_path) + continue self.logger.info('Binarizing [%3d/%d] %s', i + 1, len(ls_imgs), image_path) - image_stem = Path(image_path).stem - image = cv2.imread(os.path.join(dir_in,image_path) ) - img_last = 0 - model_file, model = self.models - self.logger.info('Predicting %s with model %s', image_path if image_path else '[image]', model_file) - res = self.predict(model, image, use_patches) + image = cv2.imread(os.path.join(dir_in, image_path)) + img_last = self.run_single(image, use_patches) + self.logger.info('Writing binarized image to %s', output_path) + cv2.imwrite(output_path, img_last) - img_fin = np.zeros((res.shape[0], res.shape[1], 3)) - res[:, :][res[:, :] == 0] = 2 - res = res - 1 - res = res * 255 - img_fin[:, :, 0] = res - img_fin[:, :, 1] = res - img_fin[:, :, 2] = res + def run_single(self, image: np.ndarray, use_patches=False): + img_last = 0 + model_file, model = self.models + res = self.predict(model, image, use_patches) - img_fin = img_fin.astype(np.uint8) - img_fin = (res[:, :] == 0) * 255 - img_last = img_last + img_fin + img_fin = np.zeros((res.shape[0], res.shape[1], 3)) + res[:, :][res[:, :] == 0] = 2 + res = res - 1 + res = res * 255 + img_fin[:, :, 0] = res + img_fin[:, :, 1] = res + img_fin[:, :, 2] = res - img_last[:, :][img_last[:, :] > 0] = 255 - img_last = (img_last[:, :] == 0) * 255 - - output_filename = os.path.join(output, image_stem + '.png') - self.logger.info('Writing binarized image to %s', output_filename) - cv2.imwrite(output_filename, img_last) + img_fin = img_fin.astype(np.uint8) + img_fin = (res[:, :] == 0) * 255 + img_last = img_last + img_fin + + kernel = np.ones((5, 5), np.uint8) + img_last[:, :][img_last[:, :] > 0] = 255 + img_last = (img_last[:, :] == 0) * 255 + return img_last diff --git a/src/eynollah/training/build_model_load_pretrained_weights_and_save.py b/src/eynollah/training/build_model_load_pretrained_weights_and_save.py index 40fc1fe..15eaf64 100644 --- a/src/eynollah/training/build_model_load_pretrained_weights_and_save.py +++ b/src/eynollah/training/build_model_load_pretrained_weights_and_save.py @@ -1,13 +1,9 @@ +import sys import click -import tensorflow as tf from .models import resnet50_unet -def configuration(): - gpu_options = tf.compat.v1.GPUOptions(allow_growth=True) - session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) - @click.command() def build_model_load_pretrained_weights_and_save(): n_classes = 2 @@ -17,8 +13,6 @@ def build_model_load_pretrained_weights_and_save(): pretraining = False dir_of_weights = 'model_bin_sbb_ens.h5' - # configuration() - model = resnet50_unet(n_classes, input_height, input_width, weight_decay, pretraining) model.load_weights(dir_of_weights) model.save('./name_in_another_python_version.h5') diff --git a/src/eynollah/training/cli.py b/src/eynollah/training/cli.py index 3718275..ae14f04 100644 --- a/src/eynollah/training/cli.py +++ b/src/eynollah/training/cli.py @@ -9,7 +9,7 @@ from .generate_gt_for_training import main as generate_gt_cli from .inference import main as inference_cli from .train import ex from .extract_line_gt import linegt_cli -from .weights_ensembling import main as ensemble_cli +from .weights_ensembling import ensemble_cli @click.command(context_settings=dict( ignore_unknown_options=True, diff --git a/src/eynollah/training/generate_gt_for_training.py b/src/eynollah/training/generate_gt_for_training.py index 30abd04..2422cc2 100644 --- a/src/eynollah/training/generate_gt_for_training.py +++ b/src/eynollah/training/generate_gt_for_training.py @@ -7,7 +7,7 @@ from PIL import Image, ImageDraw, ImageFont import cv2 import numpy as np -from eynollah.training.gt_gen_utils import ( +from .gt_gen_utils import ( filter_contours_area_of_image, find_format_of_given_filename_in_dir, find_new_features_of_contours, @@ -26,6 +26,9 @@ from eynollah.training.gt_gen_utils import ( @click.group() def main(): + """ + extract GT data suitable for model training for various tasks + """ pass @main.command() @@ -74,6 +77,9 @@ def main(): ) def pagexml2label(dir_xml,dir_out,type_output,config, printspace, dir_images, dir_out_images): + """ + extract PAGE-XML GT data suitable for model training for segmentation tasks + """ if config: with open(config) as f: config_params = json.load(f) @@ -110,6 +116,9 @@ def pagexml2label(dir_xml,dir_out,type_output,config, printspace, dir_images, di type=click.Path(exists=True, dir_okay=False), ) def image_enhancement(dir_imgs, dir_out_images, dir_out_labels, scales): + """ + extract image GT data suitable for model training for image enhancement tasks + """ ls_imgs = os.listdir(dir_imgs) with open(scales) as f: scale_dict = json.load(f) @@ -175,6 +184,9 @@ def image_enhancement(dir_imgs, dir_out_images, dir_out_labels, scales): ) def machine_based_reading_order(dir_xml, dir_out_modal_image, dir_out_classes, input_height, input_width, min_area_size, min_area_early): + """ + extract PAGE-XML GT data suitable for model training for reading-order task + """ xml_files_ind = os.listdir(dir_xml) xml_files_ind = [ind_xml for ind_xml in xml_files_ind if ind_xml.endswith('.xml')] input_height = int(input_height) @@ -205,14 +217,20 @@ def machine_based_reading_order(dir_xml, dir_out_modal_image, dir_out_classes, i img_header_and_sep = np.zeros((y_len,x_len), dtype='uint8') 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 + 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 - texts_corr_order_index = [index_tot_regions[tot_region_ref.index(i)] for i in id_all_text ] - texts_corr_order_index_int = [int(x) for x in texts_corr_order_index] - + try: + texts_corr_order_index_int = [int(index_tot_regions[tot_region_ref.index(i)]) + for i in id_all_text] + except ValueError as e: + print("incomplete ReadingOrder in", xml_file, "- skipping:", str(e)) + continue - co_text_all, texts_corr_order_index_int, regions_ar_less_than_early_min = filter_contours_area_of_image(img_poly, co_text_all, texts_corr_order_index_int, max_area, min_area, min_area_early) + co_text_all, texts_corr_order_index_int, regions_ar_less_than_early_min = \ + filter_contours_area_of_image(img_poly, co_text_all, texts_corr_order_index_int, + max_area, min_area, min_area_early) arg_array = np.array(range(len(texts_corr_order_index_int))) diff --git a/src/eynollah/training/gt_gen_utils.py b/src/eynollah/training/gt_gen_utils.py index 50bb6fa..d5ad4d9 100644 --- a/src/eynollah/training/gt_gen_utils.py +++ b/src/eynollah/training/gt_gen_utils.py @@ -1,15 +1,18 @@ import os import numpy as np import warnings -import xml.etree.ElementTree as ET +from lxml import etree as ET from tqdm import tqdm import cv2 from shapely import geometry from pathlib import Path from PIL import ImageFont +from ocrd_utils import bbox_from_points KERNEL = np.ones((5, 5), np.uint8) +NS = { 'pc': 'http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15' +} with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -235,12 +238,11 @@ def update_region_contours(co_text, img_boundary, erosion_rate, dilation_rate, y con_eroded = return_contours_of_interested_region(img_boundary_in,pixel, min_size ) try: - if len(con_eroded)>1: - cnt_size = np.array([cv2.contourArea(con_eroded[j]) for j in range(len(con_eroded))]) - cnt = contours[np.argmax(cnt_size)] - co_text_eroded.append(cnt) + if len(con_eroded) > 1: + largest = np.argmax(list(map(cv2.contourArea, con_eroded))) else: - co_text_eroded.append(con_eroded[0]) + largest = 0 + co_text_eroded.append(con_eroded[largest]) except: co_text_eroded.append(con) @@ -664,7 +666,10 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if dir_images: ls_org_imgs = os.listdir(dir_images) - ls_org_imgs_stem = [os.path.splitext(item)[0] for item in ls_org_imgs] + ls_org_imgs = {os.path.splitext(item)[0]: item + for item in ls_org_imgs + if not item.endswith('.xml')} + for index in tqdm(range(len(gt_list))): #try: print(gt_list[index]) @@ -681,6 +686,7 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if 'columns_width' in list(config_params.keys()): columns_width_dict = config_params['columns_width'] + # FIXME: look in /Page/@custom as well metadata_element = root1.find(link+'Metadata') num_col = None for child in metadata_element: @@ -694,55 +700,13 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ y_new = int ( x_new * (y_len / float(x_len)) ) if printspace or "printspace_as_class_in_layout" in list(config_params.keys()): - region_tags = np.unique([x for x in alltags if x.endswith('PrintSpace') or x.endswith('Border')]) - co_use_case = [] - - for tag in region_tags: - tag_endings = ['}PrintSpace','}Border'] - - if tag.endswith(tag_endings[0]) or tag.endswith(tag_endings[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_use_case.append(np.array(c_t_in)) - - img = np.zeros((y_len, x_len, 3)) - - img_poly = cv2.fillPoly(img, pts=co_use_case, color=(1, 1, 1)) - - img_poly = img_poly.astype(np.uint8) - - imgray = cv2.cvtColor(img_poly, 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))]) - - try: - cnt = contours[np.argmax(cnt_size)] - x, y, w, h = cv2.boundingRect(cnt) - except: - x, y , w, h = 0, 0, x_len, y_len - - bb_xywh = [x, y, w, h] + ps = (root1.xpath('/pc:PcGts/pc:Page/pc:Border', namespaces=NS) + + root1.xpath('/pc:PcGts/pc:Page/pc:PrintSpace', namespaces=NS)) + if len(ps): + points = ps[0].find('pc:Coords', NS).get('points') + ps_bbox = bbox_from_points(points) + else: + ps_bbox = [0, 0, None, None] if config_file and (config_params['use_case']=='textline' or config_params['use_case']=='word' or config_params['use_case']=='glyph' or config_params['use_case']=='printspace'): @@ -824,7 +788,8 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if printspace and config_params['use_case']!='printspace': - img_poly = img_poly[bb_xywh[1]:bb_xywh[1]+bb_xywh[3], bb_xywh[0]:bb_xywh[0]+bb_xywh[2], :] + img_poly = img_poly[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2], :] if 'columns_width' in list(config_params.keys()) and num_col and config_params['use_case']!='printspace': @@ -838,11 +803,18 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ cv2.imwrite(os.path.join(output_dir, xml_file_stem + '.png'), img_poly) if dir_images: - org_image_name = ls_org_imgs[ls_org_imgs_stem.index(xml_file_stem)] + org_image_name = ls_org_imgs[xml_file_stem] + if not org_image_name: + print("image file for XML stem", xml_file_stem, "is missing") + continue + if not os.path.isfile(os.path.join(dir_images, org_image_name)): + print("image file for XML stem", xml_file_stem, "is not readable") + continue img_org = cv2.imread(os.path.join(dir_images, org_image_name)) if printspace and config_params['use_case']!='printspace': - img_org = img_org[bb_xywh[1]:bb_xywh[1]+bb_xywh[3], bb_xywh[0]:bb_xywh[0]+bb_xywh[2], :] + img_org = img_org[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2], :] if 'columns_width' in list(config_params.keys()) and num_col and config_params['use_case']!='printspace': img_org = resize_image(img_org, y_new, x_new) @@ -1254,7 +1226,8 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if "printspace_as_class_in_layout" in list(config_params.keys()): printspace_mask = np.zeros((img_poly.shape[0], img_poly.shape[1])) - printspace_mask[bb_xywh[1]:bb_xywh[1]+bb_xywh[3], bb_xywh[0]:bb_xywh[0]+bb_xywh[2]] = 1 + printspace_mask[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2]] = 1 img_poly[:,:,0][printspace_mask[:,:] == 0] = printspace_class_rgb_color[0] img_poly[:,:,1][printspace_mask[:,:] == 0] = printspace_class_rgb_color[1] @@ -1315,7 +1288,8 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if "printspace_as_class_in_layout" in list(config_params.keys()): printspace_mask = np.zeros((img_poly.shape[0], img_poly.shape[1])) - printspace_mask[bb_xywh[1]:bb_xywh[1]+bb_xywh[3], bb_xywh[0]:bb_xywh[0]+bb_xywh[2]] = 1 + printspace_mask[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2]] = 1 img_poly[:,:,0][printspace_mask[:,:] == 0] = printspace_class_label img_poly[:,:,1][printspace_mask[:,:] == 0] = printspace_class_label @@ -1324,7 +1298,8 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if printspace: - img_poly = img_poly[bb_xywh[1]:bb_xywh[1]+bb_xywh[3], bb_xywh[0]:bb_xywh[0]+bb_xywh[2], :] + img_poly = img_poly[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2], :] if 'columns_width' in list(config_params.keys()) and num_col: img_poly = resize_image(img_poly, y_new, x_new) @@ -1338,11 +1313,18 @@ def get_images_of_ground_truth(gt_list, dir_in, output_dir, output_type, config_ if dir_images: - org_image_name = ls_org_imgs[ls_org_imgs_stem.index(xml_file_stem)] + org_image_name = ls_org_imgs[xml_file_stem] + if not org_image_name: + print("image file for XML stem", xml_file_stem, "is missing") + continue + if not os.path.isfile(os.path.join(dir_images, org_image_name)): + print("image file for XML stem", xml_file_stem, "is not readable") + continue img_org = cv2.imread(os.path.join(dir_images, org_image_name)) if printspace: - img_org = img_org[bb_xywh[1]:bb_xywh[1]+bb_xywh[3], bb_xywh[0]:bb_xywh[0]+bb_xywh[2], :] + img_org = img_org[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2], :] if 'columns_width' in list(config_params.keys()) and num_col: img_org = resize_image(img_org, y_new, x_new) @@ -1383,6 +1365,7 @@ def find_new_features_of_contours(contours_main): y_max_main = np.array([np.max(contours_main[j][:, 1]) for j in range(len(contours_main))]) return cx_main, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, y_corr_x_min_from_argmin + def read_xml(xml_file): file_name = Path(xml_file).stem tree1 = ET.parse(xml_file, parser = ET.XMLParser(encoding='utf-8')) @@ -1401,57 +1384,13 @@ def read_xml(xml_file): 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')]) - elif link+'Border' in alltags: - 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'] - elif link+'Border' in alltags: - 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] - + ps = (root1.xpath('/pc:PcGts/pc:Page/pc:Border', namespaces=NS) + + root1.xpath('/pc:PcGts/pc:Page/pc:PrintSpace', namespaces=NS)) + if len(ps): + points = ps[0].find('pc:Coords', NS).get('points') + ps_bbox = bbox_from_points(points) else: - bb_coord_printspace = None - + ps_bbox = [0, 0, None, None] region_tags=np.unique([x for x in alltags if x.endswith('Region')]) co_text_paragraph=[] @@ -1806,11 +1745,19 @@ def read_xml(xml_file): 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, file_name, 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 - - - + return (tree1, + root1, + ps_bbox, + file_name, + 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 bounding_box(cnt,color, corr_order_index ): # x, y, w, h = cv2.boundingRect(cnt) diff --git a/src/eynollah/training/inference.py b/src/eynollah/training/inference.py index f74e9e1..2be937d 100644 --- a/src/eynollah/training/inference.py +++ b/src/eynollah/training/inference.py @@ -1,19 +1,24 @@ +""" +Tool to load model and predict for given image. +""" + import sys import os from typing import Tuple import warnings import json -import numpy as np -import cv2 -from numpy._typing import NDArray -import tensorflow as tf -from keras.models import Model, load_model -from keras import backend as K import click -from tensorflow.python.keras import backend as tensorflow_backend +import numpy as np +from numpy._typing import NDArray +import cv2 import xml.etree.ElementTree as ET +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 +import tensorflow as tf +from tensorflow.keras.models import Model, load_model +from tensorflow.keras.layers import StringLookup + from .gt_gen_utils import ( filter_contours_area_of_image, find_new_features_of_contours, @@ -21,24 +26,37 @@ from .gt_gen_utils import ( resize_image, update_list_and_return_first_with_length_bigger_than_one ) -from .models import ( +from ..patch_encoder import ( PatchEncoder, Patches ) +from .metrics import ( + soft_dice_loss, + weighted_categorical_crossentropy, +) +from.utils import scale_padd_image_for_ocr +from ..utils.utils_ocr import decode_batch_predictions -from.utils import (scale_padd_image_for_ocr) -from eynollah.utils.utils_ocr import (decode_batch_predictions) with warnings.catch_warnings(): warnings.simplefilter("ignore") -__doc__=\ -""" -Tool to load model and predict for given image. -""" - -class sbb_predict: - def __init__(self,image, dir_in, model, task, config_params_model, patches, save, save_layout, ground_truth, xml_file, cpu, out, min_area): +class SBBPredict: + def __init__(self, + image, + dir_in, + model, + task, + config_params_model, + patches, + save, + save_layout, + ground_truth, + xml_file, + cpu, + out, + min_area, + ): self.image=image self.dir_in=dir_in self.patches=patches @@ -57,8 +75,9 @@ class sbb_predict: self.min_area = 0 def resize_image(self,img_in,input_height,input_width): - return cv2.resize( img_in, ( input_width,input_height) ,interpolation=cv2.INTER_NEAREST) - + return cv2.resize(img_in, (input_width, + input_height), + interpolation=cv2.INTER_NEAREST) def color_images(self,seg): ann_u=range(self.n_classes) @@ -74,68 +93,6 @@ class sbb_predict: seg_img[:,:,2][seg==c]=c return seg_img - def otsu_copy_binary(self,img): - img_r=np.zeros((img.shape[0],img.shape[1],3)) - img1=img[:,:,0] - - #print(img.min()) - #print(img[:,:,0].min()) - #blur = cv2.GaussianBlur(img,(5,5)) - #ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU) - _, threshold1 = cv2.threshold(img1, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - - - - img_r[:,:,0]=threshold1 - img_r[:,:,1]=threshold1 - img_r[:,:,2]=threshold1 - #img_r=img_r/float(np.max(img_r))*255 - return img_r - - def otsu_copy(self,img): - img_r=np.zeros((img.shape[0],img.shape[1],3)) - #img1=img[:,:,0] - - #print(img.min()) - #print(img[:,:,0].min()) - #blur = cv2.GaussianBlur(img,(5,5)) - #ret3,th3 = cv2.threshold(blur,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU) - _, threshold1 = cv2.threshold(img[:,:,0], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - _, threshold2 = cv2.threshold(img[:,:,1], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - _, threshold3 = cv2.threshold(img[:,:,2], 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU) - - - - img_r[:,:,0]=threshold1 - img_r[:,:,1]=threshold2 - img_r[:,:,2]=threshold3 - ###img_r=img_r/float(np.max(img_r))*255 - return img_r - - def soft_dice_loss(self,y_true, y_pred, epsilon=1e-6): - - axes = tuple(range(1, len(y_pred.shape)-1)) - - numerator = 2. * K.sum(y_pred * y_true, axes) - - denominator = K.sum(K.square(y_pred) + K.square(y_true), axes) - return 1.00 - K.mean(numerator / (denominator + epsilon)) # average over classes and batch - - # def weighted_categorical_crossentropy(self,weights=None): - # - # def loss(y_true, y_pred): - # labels_floats = tf.cast(y_true, tf.float32) - # per_pixel_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels_floats,logits=y_pred) - # - # if weights is not None: - # weight_mask = tf.maximum(tf.reduce_max(tf.constant( - # np.array(weights, dtype=np.float32)[None, None, None]) - # * labels_floats, axis=-1), 1.0) - # per_pixel_loss = per_pixel_loss * weight_mask[:, :, :, None] - # return tf.reduce_mean(per_pixel_loss) - # return self.loss - - def IoU(self,Yi,y_predi): ## mean Intersection over Union ## Mean IoU = TP/(FN + TP + FP) @@ -162,29 +119,33 @@ class sbb_predict: return mIoU def start_new_session_and_model(self): - if self.task == "cnn-rnn-ocr": - if self.cpu: - os.environ['CUDA_VISIBLE_DEVICES']='-1' - self.model = load_model(self.model_dir) - self.model = tf.keras.models.Model( - self.model.get_layer(name = "image").input, - self.model.get_layer(name = "dense2").output) + if self.cpu: + tf.config.set_visible_devices([], 'GPU') else: - config = tf.compat.v1.ConfigProto() - config.gpu_options.allow_growth = True + try: + for device in tf.config.list_physical_devices('GPU'): + tf.config.experimental.set_memory_growth(device, True) + except: + print("no GPU device available", file=sys.stderr) - session = tf.compat.v1.Session(config=config) # tf.InteractiveSession() - tensorflow_backend.set_session(session) + if self.task == "cnn-rnn-ocr": + self.model = Model( + self.model.get_layer(name = "image").input, + self.model.get_layer(name = "dense2").output) + else: + self.model = load_model(self.model_dir, compile=False, + custom_objects={"PatchEncoder": PatchEncoder, + "Patches": Patches}) - ##if self.weights_dir!=None: ##self.model.load_weights(self.weights_dir) assert isinstance(self.model, Model) if self.task != 'classification' and self.task != 'reading_order': - self.img_height=self.model.layers[len(self.model.layers)-1].output_shape[1] - self.img_width=self.model.layers[len(self.model.layers)-1].output_shape[2] - self.n_classes=self.model.layers[len(self.model.layers)-1].output_shape[3] + last = self.model.layers[-1] + self.img_height = last.output_shape[1] + self.img_width = last.output_shape[2] + self.n_classes = last.output_shape[3] def visualize_model_output(self, prediction, img, task) -> Tuple[NDArray, NDArray]: if task == "binarization": @@ -212,21 +173,16 @@ class sbb_predict: '15' : [255, 0, 255]} layout_only = np.zeros(prediction.shape) - for unq_class in unique_classes: + where = prediction[:,:,0]==unq_class rgb_class_unique = rgb_colors[str(int(unq_class))] - layout_only[:,:,0][prediction[:,:,0]==unq_class] = rgb_class_unique[0] - layout_only[:,:,1][prediction[:,:,0]==unq_class] = rgb_class_unique[1] - layout_only[:,:,2][prediction[:,:,0]==unq_class] = rgb_class_unique[2] - - + layout_only[:,:,0][where] = rgb_class_unique[0] + layout_only[:,:,1][where] = rgb_class_unique[1] + layout_only[:,:,2][where] = rgb_class_unique[2] + layout_only = layout_only.astype(np.int32) img = self.resize_image(img, layout_only.shape[0], layout_only.shape[1]) - - layout_only = layout_only.astype(np.int32) img = img.astype(np.int32) - - added_image = cv2.addWeighted(img,0.5,layout_only,0.1,0) @@ -238,10 +194,10 @@ class sbb_predict: assert isinstance(self.model, Model) if self.task == 'classification': classes_names = self.config_params_model['classification_classes_name'] - img_1ch = img=cv2.imread(image_dir, 0) - - img_1ch = img_1ch / 255.0 - img_1ch = cv2.resize(img_1ch, (self.config_params_model['input_height'], self.config_params_model['input_width']), interpolation=cv2.INTER_NEAREST) + img_1ch = cv2.imread(image_dir, 0) / 255.0 + img_1ch = cv2.resize(img_1ch, (self.config_params_model['input_height'], + self.config_params_model['input_width']), + interpolation=cv2.INTER_NEAREST) img_in = np.zeros((1, img_1ch.shape[0], img_1ch.shape[1], 3)) img_in[0, :, :, 0] = img_1ch[:, :] img_in[0, :, :, 1] = img_1ch[:, :] @@ -251,6 +207,7 @@ class sbb_predict: index_class = np.argmax(label_p_pred[0]) print("Predicted Class: {}".format(classes_names[str(int(index_class))])) + elif self.task == "cnn-rnn-ocr": img=cv2.imread(image_dir) img = scale_padd_image_for_ocr(img, self.config_params_model['input_height'], self.config_params_model['input_width']) @@ -279,19 +236,22 @@ class sbb_predict: img_height = self.config_params_model['input_height'] img_width = self.config_params_model['input_width'] - tree_xml, root_xml, bb_coord_printspace, file_name, id_paragraph, id_header, co_text_paragraph, co_text_header, tot_region_ref, x_len, y_len, index_tot_regions, img_poly = read_xml(self.xml_file) - _, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, _ = find_new_features_of_contours(co_text_header) + tree_xml, root_xml, ps_bbox, file_name, \ + id_paragraph, id_header, \ + co_text_paragraph, co_text_header, \ + tot_region_ref, x_len, y_len, index_tot_regions, \ + img_poly = read_xml(self.xml_file) + _, cy_main, x_min_main, x_max_main, y_min_main, y_max_main, _ = \ + find_new_features_of_contours(co_text_header) img_header_and_sep = np.zeros((y_len,x_len), dtype='uint8') - - 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 - + 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 = co_text_paragraph + co_text_header id_all_text = id_paragraph + id_header - ##texts_corr_order_index = [index_tot_regions[tot_region_ref.index(i)] for i in id_all_text ] ##texts_corr_order_index_int = [int(x) for x in texts_corr_order_index] texts_corr_order_index_int = list(np.array(range(len(co_text_all)))) @@ -302,7 +262,8 @@ class sbb_predict: #print(np.shape(co_text_all[0]), len( np.shape(co_text_all[0]) ),'co_text_all') #co_text_all = filter_contours_area_of_image_tables(img_poly, co_text_all, _, max_area, min_area) #print(co_text_all,'co_text_all') - co_text_all, texts_corr_order_index_int, _ = filter_contours_area_of_image(img_poly, co_text_all, texts_corr_order_index_int, max_area, self.min_area) + co_text_all, texts_corr_order_index_int, _ = filter_contours_area_of_image( + img_poly, co_text_all, texts_corr_order_index_int, max_area, self.min_area) #print(texts_corr_order_index_int) @@ -315,15 +276,13 @@ class sbb_predict: img_label=cv2.fillPoly(img_label, pts =[co_text_all[i]], color=(1,1,1)) labels_con[:,:,i] = img_label[:,:,0] - if bb_coord_printspace: - #bb_coord_printspace[x,y,w,h,_,_] - x = bb_coord_printspace[0] - y = bb_coord_printspace[1] - w = bb_coord_printspace[2] - h = bb_coord_printspace[3] - labels_con = labels_con[y:y+h, x:x+w, :] - img_poly = img_poly[y:y+h, x:x+w, :] - img_header_and_sep = img_header_and_sep[y:y+h, x:x+w] + if ps_bbox: + labels_con = labels_con[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2], :] + img_poly = img_poly[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2], :] + img_header_and_sep = img_header_and_sep[ps_bbox[1]:ps_bbox[3], + ps_bbox[0]:ps_bbox[2]] @@ -709,17 +668,15 @@ class sbb_predict: help="min area size of regions considered for reading order detection. The default value is zero and means that all text regions are considered for reading order.", ) def main(image, dir_in, model, patches, save, save_layout, ground_truth, xml_file, cpu, out, min_area): - assert image or dir_in, "Either a single image -i or a dir_in -di is required" + assert image or dir_in, "Either a single image -i or a dir_in -di input is required" with open(os.path.join(model,'config.json')) as f: config_params_model = json.load(f) task = config_params_model['task'] - if task != 'classification' and task != 'reading_order' and task != "cnn-rnn-ocr": - if image and not save: - print("Error: You used one of segmentation or binarization task with image input but not set -s, you need a filename to save visualized output with -s") - sys.exit(1) - if dir_in and not out: - print("Error: You used one of segmentation or binarization task with dir_in but not set -out") - sys.exit(1) - x=sbb_predict(image, dir_in, model, task, config_params_model, patches, save, save_layout, ground_truth, xml_file, cpu, out, min_area) + if task not in ['classification', 'reading_order', "cnn-rnn-ocr"]: + assert not image or save, "For segmentation or binarization, an input single image -i also requires an output filename -s" + assert not dir_in or out, "For segmentation or binarization, an input directory -di also requires an output directory -o" + x = SBBPredict(image, dir_in, model, task, config_params_model, + patches, save, save_layout, ground_truth, xml_file, + cpu, out, min_area) x.run() diff --git a/src/eynollah/training/models.py b/src/eynollah/training/models.py index a1d52ce..ba61764 100644 --- a/src/eynollah/training/models.py +++ b/src/eynollah/training/models.py @@ -1,9 +1,14 @@ -from tensorflow import keras -from keras.layers import ( +import os + +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 +import tensorflow as tf +from tensorflow.keras.layers import ( Activation, Add, AveragePooling2D, BatchNormalization, + Bidirectional, + Conv1D, Conv2D, Dense, Dropout, @@ -13,30 +18,33 @@ from keras.layers import ( Lambda, Layer, LayerNormalization, + LSTM, MaxPooling2D, MultiHeadAttention, + Reshape, UpSampling2D, ZeroPadding2D, add, concatenate ) -from keras.models import Model -import tensorflow as tf -# from keras import layers, models -from keras.regularizers import l2 +from tensorflow.keras.models import Model +from tensorflow.keras.regularizers import l2 -from eynollah.patch_encoder import Patches, PatchEncoder +from ..patch_encoder import Patches, PatchEncoder ##mlp_head_units = [512, 256]#[2048, 1024] ###projection_dim = 64 ##transformer_layers = 2#8 ##num_heads = 1#4 -resnet50_Weights_path = './pretrained_model/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5' +RESNET50_WEIGHTS_PATH = './pretrained_model/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5' +RESNET50_WEIGHTS_URL = ('https://github.com/fchollet/deep-learning-models/releases/download/v0.2/' + 'resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5') + IMAGE_ORDERING = 'channels_last' MERGE_AXIS = -1 -class CTCLayer(tf.keras.layers.Layer): +class CTCLayer(Layer): def __init__(self, name=None): super().__init__(name=name) self.loss_fn = tf.keras.backend.ctc_batch_cost @@ -61,14 +69,9 @@ def mlp(x, hidden_units, dropout_rate): return x def one_side_pad(x): - x = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(x) - if IMAGE_ORDERING == 'channels_first': - x = Lambda(lambda x: x[:, :, :-1, :-1])(x) - elif IMAGE_ORDERING == 'channels_last': - x = Lambda(lambda x: x[:, :-1, :-1, :])(x) + x = ZeroPadding2D(((1, 0), (1, 0)), data_format=IMAGE_ORDERING)(x) return x - def identity_block(input_tensor, kernel_size, filters, stage, block): """The identity block is the block that has no conv layer at shortcut. # Arguments @@ -151,19 +154,13 @@ def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)) x = Activation('relu')(x) return x - -def resnet50_unet_light(n_classes, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False): - assert input_height % 32 == 0 - assert input_width % 32 == 0 - - img_input = Input(shape=(input_height, input_width, 3)) - +def resnet50(inputs, weight_decay=1e-6, pretraining=False): if IMAGE_ORDERING == 'channels_last': bn_axis = 3 else: bn_axis = 1 - x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(img_input) + x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(inputs) x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2), kernel_regularizer=l2(weight_decay), name='conv1')(x) f1 = x @@ -197,61 +194,86 @@ def resnet50_unet_light(n_classes, input_height=224, input_width=224, task="segm f5 = x if pretraining: - model = Model(img_input, x).load_weights(resnet50_Weights_path) + model = Model(inputs, x).load_weights(RESNET50_WEIGHTS_PATH) - v512_2048 = Conv2D(512, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f5) - v512_2048 = (BatchNormalization(axis=bn_axis))(v512_2048) - v512_2048 = Activation('relu')(v512_2048) + return f1, f2, f3, f4, f5 - v512_1024 = Conv2D(512, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f4) - v512_1024 = (BatchNormalization(axis=bn_axis))(v512_1024) - v512_1024 = Activation('relu')(v512_1024) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(v512_2048) - o = (concatenate([o, v512_1024], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f3], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f2], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f1], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, img_input], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(32, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = Conv2D(n_classes, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) - if task == "segmentation": - o = (BatchNormalization(axis=bn_axis))(o) - o = (Activation('softmax'))(o) +def unet_decoder(img, f1, f2, f3, f4, f5, n_classes, light=False, task="segmentation", weight_decay=1e-6): + if IMAGE_ORDERING == 'channels_last': + bn_axis = 3 else: - o = (Activation('sigmoid'))(o) + bn_axis = 1 - model = Model(img_input, o) - return model + o = Conv2D(512 if light else 1024, (1, 1), padding='same', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f5) + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('relu')(o) + if light: + f4 = Conv2D(512, (1, 1), padding='same', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(f4) + f4 = BatchNormalization(axis=bn_axis)(f4) + f4 = Activation('relu')(f4) + + o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING)(o) + o = concatenate([o, f4], axis=MERGE_AXIS) + o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o) + o = Conv2D(512, (3, 3), padding='valid', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('relu')(o) + + o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING)(o) + o = concatenate([o, f3], axis=MERGE_AXIS) + o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o) + o = Conv2D(256, (3, 3), padding='valid', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('relu')(o) + + o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING)(o) + o = concatenate([o, f2], axis=MERGE_AXIS) + o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o) + o = Conv2D(128, (3, 3), padding='valid', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('relu')(o) + + o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING)(o) + o = concatenate([o, f1], axis=MERGE_AXIS) + o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o) + o = Conv2D(64, (3, 3), padding='valid', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('relu')(o) + + o = UpSampling2D((2, 2), data_format=IMAGE_ORDERING)(o) + o = concatenate([o, img], axis=MERGE_AXIS) + o = ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING)(o) + o = Conv2D(32, (3, 3), padding='valid', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('relu')(o) + + o = Conv2D(n_classes, (1, 1), padding='same', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) + if task == "segmentation": + o = BatchNormalization(axis=bn_axis)(o) + o = Activation('softmax')(o) + else: + o = Activation('sigmoid')(o) + + return Model(img, o) + +def resnet50_unet_light(n_classes, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False): + assert input_height % 32 == 0 + assert input_width % 32 == 0 + + img_input = Input(shape=(input_height, input_width, 3)) + + features = resnet50(img_input, weight_decay=weight_decay, pretraining=pretraining) + + return unet_decoder(img_input, *features, n_classes, light=True, task=task, weight_decay=weight_decay) def resnet50_unet(n_classes, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False): assert input_height % 32 == 0 @@ -259,162 +281,29 @@ def resnet50_unet(n_classes, input_height=224, input_width=224, task="segmentati img_input = Input(shape=(input_height, input_width, 3)) - if IMAGE_ORDERING == 'channels_last': - bn_axis = 3 - else: - bn_axis = 1 + features = resnet50(img_input, weight_decay=weight_decay, pretraining=pretraining) - x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(img_input) - x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2), kernel_regularizer=l2(weight_decay), - name='conv1')(x) - f1 = x + return unet_decoder(img_input, *features, n_classes, light=False, task=task, weight_decay=weight_decay) - x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x) - x = Activation('relu')(x) - x = MaxPooling2D((3, 3), data_format=IMAGE_ORDERING, strides=(2, 2))(x) - - x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) - x = identity_block(x, 3, [64, 64, 256], stage=2, block='b') - x = identity_block(x, 3, [64, 64, 256], stage=2, block='c') - f2 = one_side_pad(x) - - x = conv_block(x, 3, [128, 128, 512], stage=3, block='a') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='b') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='c') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='d') - f3 = x - - x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f') - f4 = x - - x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c') - f5 = x - - if pretraining: - Model(img_input, x).load_weights(resnet50_Weights_path) - - v1024_2048 = Conv2D(1024, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))( - f5) - v1024_2048 = (BatchNormalization(axis=bn_axis))(v1024_2048) - v1024_2048 = Activation('relu')(v1024_2048) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(v1024_2048) - o = (concatenate([o, f4], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f3], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f2], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f1], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, img_input], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(32, (3, 3), padding='valid', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = Conv2D(n_classes, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay))(o) - if task == "segmentation": - o = (BatchNormalization(axis=bn_axis))(o) - o = (Activation('softmax'))(o) - else: - o = (Activation('sigmoid'))(o) - - model = Model(img_input, o) - - return model - - -def vit_resnet50_unet(n_classes, patch_size_x, patch_size_y, num_patches, mlp_head_units=None, transformer_layers=8, num_heads =4, projection_dim = 64, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False): - if mlp_head_units is None: - mlp_head_units = [128, 64] - inputs = Input(shape=(input_height, input_width, 3)) - - #transformer_units = [ - #projection_dim * 2, - #projection_dim, - #] # Size of the transformer layers - IMAGE_ORDERING = 'channels_last' - bn_axis=3 - - x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(inputs) - x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2),kernel_regularizer=l2(weight_decay), name='conv1')(x) - f1 = x - - x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x) - x = Activation('relu')(x) - x = MaxPooling2D((3, 3), data_format=IMAGE_ORDERING, strides=(2, 2))(x) - - x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) - x = identity_block(x, 3, [64, 64, 256], stage=2, block='b') - x = identity_block(x, 3, [64, 64, 256], stage=2, block='c') - f2 = one_side_pad(x) - - x = conv_block(x, 3, [128, 128, 512], stage=3, block='a') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='b') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='c') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='d') - f3 = x - - x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f') - f4 = x - - x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c') - f5 = x - - if pretraining: - model = Model(inputs, x).load_weights(resnet50_Weights_path) - - #num_patches = x.shape[1]*x.shape[2] - - #patch_size_y = input_height / x.shape[1] - #patch_size_x = input_width / x.shape[2] - #patch_size = patch_size_x * patch_size_y - patches = Patches(patch_size_x, patch_size_y)(x) +def transformer_block(img, + num_patches, + patchsize_x, + patchsize_y, + mlp_head_units, + n_layers, + num_heads, + projection_dim): + patches = Patches(patchsize_x, patchsize_y)(img) # Encode patches. encoded_patches = PatchEncoder(num_patches, projection_dim)(patches) - - for _ in range(transformer_layers): + + for _ in range(n_layers): # Layer normalization 1. x1 = LayerNormalization(epsilon=1e-6)(encoded_patches) # Create a multi-head attention layer. - attention_output = MultiHeadAttention( - num_heads=num_heads, key_dim=projection_dim, dropout=0.1 - )(x1, x1) + attention_output = MultiHeadAttention(num_heads=num_heads, + key_dim=projection_dim, + dropout=0.1)(x1, x1) # Skip connection 1. x2 = Add()([attention_output, encoded_patches]) # Layer normalization 2. @@ -423,180 +312,80 @@ def vit_resnet50_unet(n_classes, patch_size_x, patch_size_y, num_patches, mlp_he x3 = mlp(x3, hidden_units=mlp_head_units, dropout_rate=0.1) # Skip connection 2. encoded_patches = Add()([x3, x2]) - - assert isinstance(x, Layer) - encoded_patches = tf.reshape(encoded_patches, [-1, x.shape[1], x.shape[2] , int( projection_dim / (patch_size_x * patch_size_y) )]) - v1024_2048 = Conv2D( 1024 , (1, 1), padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay))(encoded_patches) - v1024_2048 = (BatchNormalization(axis=bn_axis))(v1024_2048) - v1024_2048 = Activation('relu')(v1024_2048) - - o = (UpSampling2D( (2, 2), data_format=IMAGE_ORDERING))(v1024_2048) - o = (concatenate([o, f4],axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o ,f3], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f2], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f1], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, inputs],axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(32, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = Conv2D(n_classes, (1, 1), padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay))(o) - if task == "segmentation": - o = (BatchNormalization(axis=bn_axis))(o) - o = (Activation('softmax'))(o) - else: - o = (Activation('sigmoid'))(o) + encoded_patches = tf.reshape(encoded_patches, + [-1, + img.shape[1], + img.shape[2], + projection_dim // (patchsize_x * patchsize_y)]) + return encoded_patches - model = Model(inputs=inputs, outputs=o) - - return model - -def vit_resnet50_unet_transformer_before_cnn(n_classes, patch_size_x, patch_size_y, num_patches, mlp_head_units=None, transformer_layers=8, num_heads =4, projection_dim = 64, input_height=224, input_width=224, task="segmentation", weight_decay=1e-6, pretraining=False): - if mlp_head_units is None: - mlp_head_units = [128, 64] +def vit_resnet50_unet(num_patches, + n_classes, + transformer_patchsize_x, + transformer_patchsize_y, + transformer_mlp_head_units=None, + transformer_layers=8, + transformer_num_heads=4, + transformer_projection_dim=64, + input_height=224, + input_width=224, + task="segmentation", + weight_decay=1e-6, + pretraining=False): + if transformer_mlp_head_units is None: + transformer_mlp_head_units = [128, 64] inputs = Input(shape=(input_height, input_width, 3)) - ##transformer_units = [ - ##projection_dim * 2, - ##projection_dim, - ##] # Size of the transformer layers - IMAGE_ORDERING = 'channels_last' - bn_axis=3 - - patches = Patches(patch_size_x, patch_size_y)(inputs) - # Encode patches. - encoded_patches = PatchEncoder(num_patches, projection_dim)(patches) - - for _ in range(transformer_layers): - # Layer normalization 1. - x1 = LayerNormalization(epsilon=1e-6)(encoded_patches) - # Create a multi-head attention layer. - attention_output = MultiHeadAttention( - num_heads=num_heads, key_dim=projection_dim, dropout=0.1 - )(x1, x1) - # Skip connection 1. - x2 = Add()([attention_output, encoded_patches]) - # Layer normalization 2. - x3 = LayerNormalization(epsilon=1e-6)(x2) - # MLP. - x3 = mlp(x3, hidden_units=mlp_head_units, dropout_rate=0.1) - # Skip connection 2. - encoded_patches = Add()([x3, x2]) - - encoded_patches = tf.reshape(encoded_patches, [-1, input_height, input_width , int( projection_dim / (patch_size_x * patch_size_y) )]) - - encoded_patches = Conv2D(3, (1, 1), padding='same', data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay), name='convinput')(encoded_patches) - - x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(encoded_patches) - x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2),kernel_regularizer=l2(weight_decay), name='conv1')(x) - f1 = x + features = resnet50(inputs, weight_decay=weight_decay, pretraining=pretraining) - x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x) - x = Activation('relu')(x) - x = MaxPooling2D((3, 3), data_format=IMAGE_ORDERING, strides=(2, 2))(x) - - x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) - x = identity_block(x, 3, [64, 64, 256], stage=2, block='b') - x = identity_block(x, 3, [64, 64, 256], stage=2, block='c') - f2 = one_side_pad(x) + features[-1] = transformer_block(features[-1], + num_patches, + transformer_patchsize_x, + transformer_patchsize_y, + transformer_mlp_head_units, + transformer_layers, + transformer_num_heads, + transformer_projection_dim) - x = conv_block(x, 3, [128, 128, 512], stage=3, block='a') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='b') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='c') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='d') - f3 = x + o = unet_decoder(inputs, *features, n_classes, task=task, weight_decay=weight_decay) - x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f') - f4 = x + return Model(inputs, o) - x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c') - f5 = x - - if pretraining: - model = Model(encoded_patches, x).load_weights(resnet50_Weights_path) +def vit_resnet50_unet_transformer_before_cnn(num_patches, + n_classes, + transformer_patchsize_x, + transformer_patchsize_y, + transformer_mlp_head_units=None, + transformer_layers=8, + transformer_num_heads=4, + transformer_projection_dim=64, + input_height=224, + input_width=224, + task="segmentation", + weight_decay=1e-6, + pretraining=False): + if transformer_mlp_head_units is None: + transformer_mlp_head_units = [128, 64] + inputs = Input(shape=(input_height, input_width, 3)) - v1024_2048 = Conv2D( 1024 , (1, 1), padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay))(x) - v1024_2048 = (BatchNormalization(axis=bn_axis))(v1024_2048) - v1024_2048 = Activation('relu')(v1024_2048) - - o = (UpSampling2D( (2, 2), data_format=IMAGE_ORDERING))(v1024_2048) - o = (concatenate([o, f4],axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o ,f3], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f2], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, f1], axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = (UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(o) - o = (concatenate([o, inputs],axis=MERGE_AXIS)) - o = (ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(o) - o = (Conv2D(32, (3, 3), padding='valid', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay)))(o) - o = (BatchNormalization(axis=bn_axis))(o) - o = Activation('relu')(o) - - o = Conv2D(n_classes, (1, 1), padding='same', data_format=IMAGE_ORDERING,kernel_regularizer=l2(weight_decay))(o) - if task == "segmentation": - o = (BatchNormalization(axis=bn_axis))(o) - o = (Activation('softmax'))(o) - else: - o = (Activation('sigmoid'))(o) + encoded_patches = transformer_block(inputs, + num_patches, + transformer_patchsize_x, + transformer_patchsize_y, + transformer_mlp_head_units, + transformer_layers, + transformer_num_heads, + transformer_projection_dim) + encoded_patches = Conv2D(3, (1, 1), padding='same', + data_format=IMAGE_ORDERING, kernel_regularizer=l2(weight_decay), + name='convinput')(encoded_patches) - model = Model(inputs=inputs, outputs=o) + features = resnet50(encoded_patches, weight_decay=weight_decay, pretraining=pretraining) + + o = unet_decoder(inputs, *features, n_classes, task=task, weight_decay=weight_decay) - return model + return Model(inputs, o) def resnet50_classifier(n_classes,input_height=224,input_width=224,weight_decay=1e-6,pretraining=False): include_top=True @@ -606,47 +395,7 @@ def resnet50_classifier(n_classes,input_height=224,input_width=224,weight_decay= img_input = Input(shape=(input_height,input_width , 3 )) - if IMAGE_ORDERING == 'channels_last': - bn_axis = 3 - else: - bn_axis = 1 - - x = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(img_input) - x = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2),kernel_regularizer=l2(weight_decay), name='conv1')(x) - f1 = x - - x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x) - x = Activation('relu')(x) - x = MaxPooling2D((3, 3) , data_format=IMAGE_ORDERING , strides=(2, 2))(x) - - - x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) - x = identity_block(x, 3, [64, 64, 256], stage=2, block='b') - x = identity_block(x, 3, [64, 64, 256], stage=2, block='c') - f2 = one_side_pad(x ) - - - x = conv_block(x, 3, [128, 128, 512], stage=3, block='a') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='b') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='c') - x = identity_block(x, 3, [128, 128, 512], stage=3, block='d') - f3 = x - - x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e') - x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f') - f4 = x - - x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b') - x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c') - f5 = x - - if pretraining: - Model(img_input, x).load_weights(resnet50_Weights_path) + _, _, _, _, x = resnet50(img_input, weight_decay, pretraining) x = AveragePooling2D((7, 7), name='avg_pool')(x) x = Flatten()(x) @@ -658,9 +407,6 @@ def resnet50_classifier(n_classes,input_height=224,input_width=224,weight_decay= x = Dense(n_classes, activation='softmax', name='fc1000')(x) model = Model(img_input, x) - - - return model def machine_based_reading_order_model(n_classes,input_height=224,input_width=224,weight_decay=1e-6,pretraining=False): @@ -669,43 +415,10 @@ def machine_based_reading_order_model(n_classes,input_height=224,input_width=224 img_input = Input(shape=(input_height,input_width , 3 )) - if IMAGE_ORDERING == 'channels_last': - bn_axis = 3 - else: - bn_axis = 1 - - x1 = ZeroPadding2D((3, 3), data_format=IMAGE_ORDERING)(img_input) - x1 = Conv2D(64, (7, 7), data_format=IMAGE_ORDERING, strides=(2, 2),kernel_regularizer=l2(weight_decay), name='conv1')(x1) - - x1 = BatchNormalization(axis=bn_axis, name='bn_conv1')(x1) - x1 = Activation('relu')(x1) - x1 = MaxPooling2D((3, 3) , data_format=IMAGE_ORDERING , strides=(2, 2))(x1) + _, _, _, _, x = resnet50(img_input, weight_decay, pretraining) - x1 = conv_block(x1, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) - x1 = identity_block(x1, 3, [64, 64, 256], stage=2, block='b') - x1 = identity_block(x1, 3, [64, 64, 256], stage=2, block='c') - - x1 = conv_block(x1, 3, [128, 128, 512], stage=3, block='a') - x1 = identity_block(x1, 3, [128, 128, 512], stage=3, block='b') - x1 = identity_block(x1, 3, [128, 128, 512], stage=3, block='c') - x1 = identity_block(x1, 3, [128, 128, 512], stage=3, block='d') - - x1 = conv_block(x1, 3, [256, 256, 1024], stage=4, block='a') - x1 = identity_block(x1, 3, [256, 256, 1024], stage=4, block='b') - x1 = identity_block(x1, 3, [256, 256, 1024], stage=4, block='c') - x1 = identity_block(x1, 3, [256, 256, 1024], stage=4, block='d') - x1 = identity_block(x1, 3, [256, 256, 1024], stage=4, block='e') - x1 = identity_block(x1, 3, [256, 256, 1024], stage=4, block='f') - - x1 = conv_block(x1, 3, [512, 512, 2048], stage=5, block='a') - x1 = identity_block(x1, 3, [512, 512, 2048], stage=5, block='b') - x1 = identity_block(x1, 3, [512, 512, 2048], stage=5, block='c') - - if pretraining: - Model(img_input , x1).load_weights(resnet50_Weights_path) - - x1 = AveragePooling2D((7, 7), name='avg_pool1')(x1) - flattened = Flatten()(x1) + x = AveragePooling2D((7, 7), name='avg_pool1')(x) + flattened = Flatten()(x) o = Dense(256, activation='relu', name='fc512')(flattened) o=Dropout(0.2)(o) @@ -719,83 +432,79 @@ def machine_based_reading_order_model(n_classes,input_height=224,input_width=224 return model def cnn_rnn_ocr_model(image_height=None, image_width=None, n_classes=None, max_seq=None): - input_img = tf.keras.Input(shape=(image_height, image_width, 3), name="image") - labels = tf.keras.layers.Input(name="label", shape=(None,)) + input_img = Input(shape=(image_height, image_width, 3), name="image") + labels = Input(name="label", shape=(None,)) - x = tf.keras.layers.Conv2D(64,kernel_size=(3,3),padding="same")(input_img) - x = tf.keras.layers.BatchNormalization(name="bn1")(x) - x = tf.keras.layers.Activation("relu", name="relu1")(x) - x = tf.keras.layers.Conv2D(64,kernel_size=(3,3),padding="same")(x) - x = tf.keras.layers.BatchNormalization(name="bn2")(x) - x = tf.keras.layers.Activation("relu", name="relu2")(x) - x = tf.keras.layers.MaxPool2D(pool_size=(1,2),strides=(1,2))(x) + x = Conv2D(64,kernel_size=(3,3),padding="same")(input_img) + x = BatchNormalization(name="bn1")(x) + x = Activation("relu", name="relu1")(x) + x = Conv2D(64,kernel_size=(3,3),padding="same")(x) + x = BatchNormalization(name="bn2")(x) + x = Activation("relu", name="relu2")(x) + x = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x) - x = tf.keras.layers.Conv2D(128,kernel_size=(3,3),padding="same")(x) - x = tf.keras.layers.BatchNormalization(name="bn3")(x) - x = tf.keras.layers.Activation("relu", name="relu3")(x) - x = tf.keras.layers.Conv2D(128,kernel_size=(3,3),padding="same")(x) - x = tf.keras.layers.BatchNormalization(name="bn4")(x) - x = tf.keras.layers.Activation("relu", name="relu4")(x) - x = tf.keras.layers.MaxPool2D(pool_size=(1,2),strides=(1,2))(x) + x = Conv2D(128,kernel_size=(3,3),padding="same")(x) + x = BatchNormalization(name="bn3")(x) + x = Activation("relu", name="relu3")(x) + x = Conv2D(128,kernel_size=(3,3),padding="same")(x) + x = BatchNormalization(name="bn4")(x) + x = Activation("relu", name="relu4")(x) + x = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x) - x = tf.keras.layers.Conv2D(256,kernel_size=(3,3),padding="same")(x) - x = tf.keras.layers.BatchNormalization(name="bn5")(x) - x = tf.keras.layers.Activation("relu", name="relu5")(x) - x = tf.keras.layers.Conv2D(256,kernel_size=(3,3),padding="same")(x) - x = tf.keras.layers.BatchNormalization(name="bn6")(x) - x = tf.keras.layers.Activation("relu", name="relu6")(x) - x = tf.keras.layers.MaxPool2D(pool_size=(2,2),strides=(2,2))(x) + x = Conv2D(256,kernel_size=(3,3),padding="same")(x) + x = BatchNormalization(name="bn5")(x) + x = Activation("relu", name="relu5")(x) + x = Conv2D(256,kernel_size=(3,3),padding="same")(x) + x = BatchNormalization(name="bn6")(x) + x = Activation("relu", name="relu6")(x) + x = MaxPooling2D(pool_size=(2,2),strides=(2,2))(x) - x = tf.keras.layers.Conv2D(image_width,kernel_size=(3,3),padding="same")(x) - x = tf.keras.layers.BatchNormalization(name="bn7")(x) - x = tf.keras.layers.Activation("relu", name="relu7")(x) - x = tf.keras.layers.Conv2D(image_width,kernel_size=(16,1))(x) - x = tf.keras.layers.BatchNormalization(name="bn8")(x) - x = tf.keras.layers.Activation("relu", name="relu8")(x) - x2d = tf.keras.layers.MaxPool2D(pool_size=(1,2),strides=(1,2))(x) - x4d = tf.keras.layers.MaxPool2D(pool_size=(1,2),strides=(1,2))(x2d) + x = Conv2D(image_width,kernel_size=(3,3),padding="same")(x) + x = BatchNormalization(name="bn7")(x) + x = Activation("relu", name="relu7")(x) + x = Conv2D(image_width,kernel_size=(16,1))(x) + x = BatchNormalization(name="bn8")(x) + x = Activation("relu", name="relu8")(x) + x2d = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x) + x4d = MaxPooling2D(pool_size=(1,2),strides=(1,2))(x2d) new_shape = (x.shape[1]*x.shape[2], x.shape[3]) new_shape2 = (x2d.shape[1]*x2d.shape[2], x2d.shape[3]) new_shape4 = (x4d.shape[1]*x4d.shape[2], x4d.shape[3]) - x = tf.keras.layers.Reshape(target_shape=new_shape, name="reshape")(x) - x2d = tf.keras.layers.Reshape(target_shape=new_shape2, name="reshape2")(x2d) - x4d = tf.keras.layers.Reshape(target_shape=new_shape4, name="reshape4")(x4d) + x = Reshape(target_shape=new_shape, name="reshape")(x) + x2d = Reshape(target_shape=new_shape2, name="reshape2")(x2d) + x4d = Reshape(target_shape=new_shape4, name="reshape4")(x4d) + + xrnnorg = Bidirectional(LSTM(image_width, return_sequences=True, dropout=0.25))(x) + xrnn2d = Bidirectional(LSTM(image_width, return_sequences=True, dropout=0.25))(x2d) + xrnn4d = Bidirectional(LSTM(image_width, return_sequences=True, dropout=0.25))(x4d) + + xrnn2d = Reshape(target_shape=(1, xrnn2d.shape[1], xrnn2d.shape[2]), name="reshape6")(xrnn2d) + xrnn4d = Reshape(target_shape=(1, xrnn4d.shape[1], xrnn4d.shape[2]), name="reshape8")(xrnn4d) - xrnnorg = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(image_width, return_sequences=True, dropout=0.25))(x) - xrnn2d = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(image_width, return_sequences=True, dropout=0.25))(x2d) - xrnn4d = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(image_width, return_sequences=True, dropout=0.25))(x4d) - - xrnn2d = tf.keras.layers.Reshape(target_shape=(1, xrnn2d.shape[1], xrnn2d.shape[2]), name="reshape6")(xrnn2d) - xrnn4d = tf.keras.layers.Reshape(target_shape=(1, xrnn4d.shape[1], xrnn4d.shape[2]), name="reshape8")(xrnn4d) + xrnn2dup = UpSampling2D(size=(1, 2), interpolation="nearest")(xrnn2d) + xrnn4dup = UpSampling2D(size=(1, 4), interpolation="nearest")(xrnn4d) + xrnn2dup = Reshape(target_shape=(xrnn2dup.shape[2], xrnn2dup.shape[3]), name="reshape10")(xrnn2dup) + xrnn4dup = Reshape(target_shape=(xrnn4dup.shape[2], xrnn4dup.shape[3]), name="reshape12")(xrnn4dup) - xrnn2dup = tf.keras.layers.UpSampling2D(size=(1, 2), interpolation="nearest")(xrnn2d) - xrnn4dup = tf.keras.layers.UpSampling2D(size=(1, 4), interpolation="nearest")(xrnn4d) + addition = Add()([xrnnorg, xrnn2dup, xrnn4dup]) - xrnn2dup = tf.keras.layers.Reshape(target_shape=(xrnn2dup.shape[2], xrnn2dup.shape[3]), name="reshape10")(xrnn2dup) - xrnn4dup = tf.keras.layers.Reshape(target_shape=(xrnn4dup.shape[2], xrnn4dup.shape[3]), name="reshape12")(xrnn4dup) + addition_rnn = Bidirectional(LSTM(image_width, return_sequences=True, dropout=0.25))(addition) + + out = Conv1D(max_seq, 1, data_format="channels_first")(addition_rnn) + out = BatchNormalization(name="bn9")(out) + out = Activation("relu", name="relu9")(out) + #out = Conv1D(n_classes, 1, activation='relu', data_format="channels_last")(out) - addition = tf.keras.layers.Add()([xrnnorg, xrnn2dup, xrnn4dup]) - - addition_rnn = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(image_width, return_sequences=True, dropout=0.25))(addition) - - out = tf.keras.layers.Conv1D(max_seq, 1, data_format="channels_first")(addition_rnn) - out = tf.keras.layers.BatchNormalization(name="bn9")(out) - out = tf.keras.layers.Activation("relu", name="relu9")(out) - #out = tf.keras.layers.Conv1D(n_classes, 1, activation='relu', data_format="channels_last")(out) - - out = tf.keras.layers.Dense( - n_classes, activation="softmax", name="dense2" - )(out) + out = Dense(n_classes, activation="softmax", name="dense2")(out) # Add CTC layer for calculating CTC loss at each step. output = CTCLayer(name="ctc_loss")(labels, out) - model = tf.keras.models.Model(inputs=[input_img, labels], outputs=output, name="handwriting_recognizer") + model = Model(inputs=[input_img, labels], outputs=output, name="handwriting_recognizer") return model - diff --git a/src/eynollah/training/train.py b/src/eynollah/training/train.py index 7a0cb3d..f6117f7 100644 --- a/src/eynollah/training/train.py +++ b/src/eynollah/training/train.py @@ -2,13 +2,29 @@ import os import sys import json -import click +import requests -from eynollah.training.metrics import ( +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 +import tensorflow as tf +from tensorflow.keras.optimizers import SGD, Adam +from tensorflow.keras.metrics import MeanIoU, F1Score +from tensorflow.keras.models import load_model +from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard +from tensorflow.keras.layers import StringLookup +from tensorflow.keras.utils import image_dataset_from_directory +from tensorflow.keras.backend import one_hot +from sacred import Experiment +from sacred.config import create_captured_function + +import numpy as np +import cv2 + +from .metrics import ( soft_dice_loss, weighted_categorical_crossentropy ) -from eynollah.training.models import ( +from .models import ( PatchEncoder, Patches, machine_based_reading_order_model, @@ -16,62 +32,47 @@ from eynollah.training.models import ( resnet50_unet, vit_resnet50_unet, vit_resnet50_unet_transformer_before_cnn, - cnn_rnn_ocr_model + cnn_rnn_ocr_model, + RESNET50_WEIGHTS_PATH, + RESNET50_WEIGHTS_URL ) -from eynollah.training.utils import ( - data_gen, - data_gen_ocr, - return_multiplier_based_on_augmnentations, +from .utils import ( generate_arrays_from_folder_reading_order, - generate_data_from_folder_evaluation, - generate_data_from_folder_training, get_one_hot, - provide_patches, - return_number_of_total_training_data + preprocess_imgs, ) +from .weights_ensembling import run_ensembling -os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' -import tensorflow as tf -from tensorflow.compat.v1.keras.backend import set_session -from tensorflow.keras.optimizers import SGD, Adam -from sacred import Experiment -from tensorflow.keras.models import load_model -from tqdm import tqdm -from sklearn.metrics import f1_score -from tensorflow.keras.callbacks import Callback -from tensorflow.keras.layers import StringLookup -import numpy as np -import cv2 - -class SaveWeightsAfterSteps(Callback): - def __init__(self, save_interval, save_path, _config): - super(SaveWeightsAfterSteps, self).__init__() - self.save_interval = save_interval - self.save_path = save_path - self.step_count = 0 +class SaveWeightsAfterSteps(ModelCheckpoint): + def __init__(self, save_interval, save_path, _config, **kwargs): + if save_interval: + # batches + super().__init__( + os.path.join(save_path, "model_step_{batch:04d}"), + save_freq=save_interval, + verbose=1, + **kwargs) + else: + super().__init__( + os.path.join(save_path, "model_{epoch:02d}"), + save_freq="epoch", + verbose=1, + **kwargs) self._config = _config - def on_train_batch_end(self, batch, logs=None): - self.step_count += 1 + # overwrite tf-keras (Keras 2) implementation to get our _config JSON in + def _save_handler(self, filepath): + super()._save_handler(filepath) + with open(os.path.join(filepath, "config.json"), "w") as fp: + json.dump(self._config, fp) # encode dict into JSON - if self.step_count % self.save_interval ==0: - save_file = f"{self.save_path}/model_step_{self.step_count}" - #os.system('mkdir '+save_file) - - self.model.save(save_file) - - with open(os.path.join(os.path.join(self.save_path, f"model_step_{self.step_count}"),"config.json"), "w") as fp: - json.dump(self._config, fp) # encode dict into JSON - print(f"saved model as steps {self.step_count} to {save_file}") - - - def configuration(): - config = tf.compat.v1.ConfigProto() - config.gpu_options.allow_growth = True - session = tf.compat.v1.Session(config=config) - set_session(session) + try: + for device in tf.config.list_physical_devices('GPU'): + tf.config.experimental.set_memory_growth(device, True) + except: + print("no GPU device available", file=sys.stderr) def get_dirs_or_files(input_data): @@ -82,241 +83,228 @@ def get_dirs_or_files(input_data): assert os.path.isdir(labels_input), "{} is not a directory".format(labels_input) return image_input, labels_input +def download_file(url, path): + with open(path, 'wb') as f: + with requests.get(url, stream=True) as r: + r.raise_for_status() + for data in r.iter_content(chunk_size=4096): + f.write(data) ex = Experiment(save_git_info=False) @ex.config def config_params(): + task = "segmentation" # This parameter defines task of model which can be segmentation, enhancement or classification. + if task in ["segmentation", "binarization", "enhancement"]: + backbone_type = "nontransformer" # Type of image feature map network backbone. Either a vision transformer alongside a CNN we call "transformer", or only a CNN which we call "nontransformer" + if backbone_type == "transformer": + transformer_patchsize_x = None # Patch size of vision transformer patches in x direction. + transformer_patchsize_y = None # Patch size of vision transformer patches in y direction. + transformer_num_patches_xy = None # Number of patches for vision transformer in x and y direction respectively. + transformer_projection_dim = 64 # Transformer projection dimension. Default value is 64. + transformer_mlp_head_units = [128, 64] # Transformer Multilayer Perceptron (MLP) head units. Default value is [128, 64] + transformer_layers = 8 # transformer layers. Default value is 8. + transformer_num_heads = 4 # Transformer number of heads. Default value is 4. + transformer_cnn_first = True # We have two types of vision transformers: either the CNN is applied first, followed by the transformer, or reversed. n_classes = None # Number of classes. In the case of binary classification this should be 2. - n_epochs = 1 # Number of epochs. + n_epochs = 1 # Number of epochs to train. + n_batch = 1 # Number of images per batch at each iteration. (Try as large as fits on VRAM.) + if task == 'cnn-rnn-ocr': + max_len = None # Maximum sequence length (characters per line) for OCR output. + characters_txt_file = None # Path of JSON file defining character set needed of OCR model. input_height = 224 * 1 # Height of model's input in pixels. input_width = 224 * 1 # Width of model's input in pixels. weight_decay = 1e-6 # Weight decay of l2 regularization of model layers. - n_batch = 1 # Number of batches at each iteration. - max_len = None # max len for ocr output. learning_rate = 1e-4 # Set the learning rate. + if task in ["segmentation", "binarization"]: + is_loss_soft_dice = False # Use soft dice as loss function. When set to true, "weighted_loss" must be false. + weighted_loss = False # Use weighted categorical cross entropy as loss fucntion. When set to true, "is_loss_soft_dice" must be false. + elif task == "classification": + f1_threshold_classification = None # This threshold is used to consider models with an evaluation f1 scores bigger than it. The selected model weights undergo a weights ensembling. And avreage ensembled model will be written to output. + classification_classes_name = None # Dictionary of classification classes names. patches = False # Divides input image into smaller patches (input size of the model) when set to true. For the model to see the full image, like page extraction, set this to false. augmentation = False # To apply any kind of augmentation, this parameter must be set to true. - flip_aug = False # If true, different types of flipping will be applied to the image. Types of flips are defined with "flip_index" in config_params.json. - blur_aug = False # If true, different types of blurring will be applied to the image. Types of blur are defined with "blur_k" in config_params.json. - padding_white = False # If true, white padding will be applied to the image. - padding_black = False # If true, black padding will be applied to the image. - scaling = False # If true, scaling will be applied to the image. The amount of scaling is defined with "scales" in config_params.json. - shifting = False - degrading = False # If true, degrading will be applied to the image. The amount of degrading is defined with "degrade_scales" in config_params.json. - brightening = False # If true, brightening will be applied to the image. The amount of brightening is defined with "brightness" in config_params.json. - binarization = False # If true, Otsu thresholding will be applied to augment the input with binarized images. - image_inversion = False # If true, and if the binarized images are avilable the image inevrsion will be applied. - white_noise_strap = False # If true, white noise will be applied on some straps on the textline image. - textline_skewing = False # If true, textline images will be skewed for augmentation. - textline_skewing_bin = False # If true, textline image skewing augmentation for binarized images will be applied if already are available. - textline_left_in_depth = False # If true, left side of textline image will be displayed in depth. - textline_left_in_depth_bin = False # If true, left side of textline binarized image (if available) will be displayed in depth. - textline_right_in_depth = False # If true, right side of textline image will be displayed in depth. - textline_right_in_depth_bin = False # If true, right side of textline binarized image (if available) will be displayed in depth. - textline_up_in_depth = False # If true, upper side of textline image will be displayed in depth. - textline_up_in_depth_bin = False # If true, upper side of textline binarized image (if available) will be displayed in depth. - textline_down_in_depth = False # If true, lower side of textline image will be displayed in depth. - textline_down_in_depth_bin = False # If true, lower side of textline binarized image (if available) will be displayed in depth. - pepper_bin_aug = False # If true, pepper noise will be added to textline binarized image (if available). - pepper_aug = False # If true, pepper noise will be added to textline image. - adding_rgb_background = False - adding_rgb_foreground = False - add_red_textlines = False - channels_shuffling = False + if augmentation: + flip_aug = False # Whether different types of flipping will be applied to the image. Requires "flip_index" setting. + if flip_aug: + flip_index = None # List of codes (as in cv2.flip) for flip augmentation. + blur_aug = False # Whether images will be blurred. Requires "blur_k" setting. + if blur_aug: + blur_k = None # Method of blurring (gauss, median or blur). + padding_white = False # If true, white padding will be applied to the image. + if padding_white and task == 'cnn-rnn-ocr': + white_padds = None # List of padding sizes. + padd_colors = None # List of padding colors, but only "white" or "black" or both. + padding_black = False # If true, black padding will be applied to the image. + scaling = False # Whether images will be scaled up or down. Requires "scales" setting. + scaling_bluring = False # Whether a combination of scaling and blurring will be applied to the image. + scaling_binarization = False # Whether a combination of scaling and binarization will be applied to the image. + scaling_brightness = False # Whether a combination of scaling and brightening will be applied to the image. + scaling_flip = False # Whether a combination of scaling and flipping will be applied to the image. + if scaling or scaling_brightness or scaling_bluring or scaling_binarization or scaling_flip: + scales = None # Scale patches for augmentation. + shifting = False + brightening = False # Whether images will be brightened. Requires "brightness" setting. + if brightening: + brightness = None # List of intensity factors for brightening. + binarization = False # Whether binary images will be used, too. (Will use Otsu thresholding unless supplying precomputed images in "dir_img_bin".) + if binarization: + dir_img_bin = None # Directory of training dataset subdirectory of binarized images + add_red_textlines = False + adding_rgb_background = False # Whether texture images will be added as artificial background. + if adding_rgb_background: + dir_rgb_backgrounds = None # Directory of texture images for synthetic background + adding_rgb_foreground = False # Whether texture images will be added as artificial foreground. + if adding_rgb_foreground: + dir_rgb_foregrounds = None # Directory of texture images for synthetic foreground + if adding_rgb_background or adding_rgb_foreground: + number_of_backgrounds_per_image = 1 + if task == 'cnn-rnn-ocr': + image_inversion = False # Whether the binarized images will be inverted. + textline_skewing_bin = False # Whether binarized textline images will be rotated. + textline_left_in_depth_bin = False # Whether left side of binary textline image will be displayed in depth. + textline_right_in_depth_bin = False # Whether right side of binary textline image will be displayed in depth. + textline_up_in_depth_bin = False # Whether upper side of binary textline image will be displayed in depth. + textline_down_in_depth_bin = False # Whether lower side of binary textline image will be displayed in depth. + pepper_bin_aug = False # Whether pepper noise will be added to binary textline images. + bin_deg = False # Whether a combination of degrading and binarization will be applied to the image. + degrading = False # Whether images will be artificially degraded. Requires the "degrade_scales" setting. + if degrading or binarization and task == 'cnn-rnn-ocr' and bin_deg: + degrade_scales = None # List of quality factors for degradation. + channels_shuffling = False # Re-arrange color channels. + if channels_shuffling: + shuffle_indexes = None # List of channels to switch between. + rotation = False # Whether images will be rotated by 90 degrees. + rotation_not_90 = False # Whether images will be rotated arbitrarily (skewed). Requires "thetha" setting. + if rotation_not_90: + thetha = None # List of rotation angles in degrees. + if task == 'cnn-rnn-ocr': + white_noise_strap = False # Whether white noise will be applied on some straps on the textline image. + textline_skewing = False # Whether textline images will be skewed for augmentation. + if textline_skewing or binarization and textline_skewing_bin: + skewing_amplitudes = None # List of skewing angles in degrees like [5, 8] + textline_left_in_depth = False # If true, left side of textline image will be displayed in depth. + textline_right_in_depth = False # If true, right side of textline image will be displayed in depth. + textline_up_in_depth = False # If true, upper side of textline image will be displayed in depth. + textline_down_in_depth = False # If true, lower side of textline image will be displayed in depth. + pepper_aug = False # Whether pepper noise will be added to textline images. + if pepper_aug or binarization and pepper_bin_aug: + pepper_indexes = None # List of pepper noise factors, e.g. [0.01, 0.005]. + color_padding_rotation = False # Whether images will be rotated with color padding. Requires "thetha_padd" setting. + if color_padding_rotation: + thetha_padd = None # List of angles (in degrees) used for rotation alongside padding. dir_train = None # Directory of training dataset with subdirectories having the names "images" and "labels". dir_eval = None # Directory of validation dataset with subdirectories having the names "images" and "labels". - dir_output = None # Directory where the output model will be saved. - pretraining = False # Set to true to load pretrained weights of ResNet50 encoder. - scaling_bluring = False # If true, a combination of scaling and blurring will be applied to the image. - scaling_binarization = False # If true, a combination of scaling and binarization will be applied to the image. - bin_deg = False # If true, a combination of degrading and binarization will be applied to the image. - rotation = False # If true, a 90 degree rotation will be implemeneted. - color_padding_rotation = False # If true, rotation and padding will be implemeneted. - rotation_not_90 = False # If true rotation based on provided angles with thetha will be implemeneted. - scaling_brightness = False # If true, a combination of scaling and brightening will be applied to the image. - scaling_flip = False # If true, a combination of scaling and flipping will be applied to the image. - thetha = None # Rotate image by these angles for augmentation. - thetha_padd = None # List of angles used for rotation alongside padding - shuffle_indexes = None # List of shuffling indexes like [[0,2,1], [1,2,0], [1,0,2]] - pepper_indexes = None # List of pepper noise indexes like [0.01, 0.005] - white_padds = None # List of padding size in the case of white padding - skewing_amplitudes = None # List of skewing augmentation amplitudes like [5, 8] - blur_k = None # Blur image for augmentation. - scales = None # Scale patches for augmentation. - padd_colors = None # padding colors. A list elements can be only white and black. like ["white", "black"] or only one of them ["white"] - degrade_scales = None # Degrade image for augmentation. - brightness = None # Brighten image for augmentation. - flip_index = None # Flip image for augmentation. - continue_training = False # Set to true if you would like to continue training an already trained a model. - transformer_patchsize_x = None # Patch size of vision transformer patches in x direction. - transformer_patchsize_y = None # Patch size of vision transformer patches in y direction. - transformer_num_patches_xy = None # Number of patches for vision transformer in x and y direction respectively. - transformer_projection_dim = 64 # Transformer projection dimension. Default value is 64. - transformer_mlp_head_units = [128, 64] # Transformer Multilayer Perceptron (MLP) head units. Default value is [128, 64] - transformer_layers = 8 # transformer layers. Default value is 8. - transformer_num_heads = 4 # Transformer number of heads. Default value is 4. - transformer_cnn_first = True # We have two types of vision transformers. In one type, a CNN is applied first, followed by a transformer. In the other type, this order is reversed. If transformer_cnn_first is true, it means the CNN will be applied before the transformer. Default value is true. - index_start = 0 # Index of model to continue training from. E.g. if you trained for 3 epochs and last index is 2, to continue from model_1.h5, set "index_start" to 3 to start naming model with index 3. - dir_of_start_model = '' # Directory containing pretrained encoder to continue training the model. - is_loss_soft_dice = False # Use soft dice as loss function. When set to true, "weighted_loss" must be false. - weighted_loss = False # Use weighted categorical cross entropy as loss fucntion. When set to true, "is_loss_soft_dice" must be false. - data_is_provided = False # Only set this to true when you have already provided the input data and the train and eval data are in "dir_output". - task = "segmentation" # This parameter defines task of model which can be segmentation, enhancement or classification. - f1_threshold_classification = None # This threshold is used to consider models with an evaluation f1 scores bigger than it. The selected model weights undergo a weights ensembling. And avreage ensembled model will be written to output. - classification_classes_name = None # Dictionary of classification classes names. - backbone_type = None # As backbone we have 2 types of backbones. A vision transformer alongside a CNN and we call it "transformer" and only CNN called "nontransformer" - save_interval = None - dir_img_bin = None - number_of_backgrounds_per_image = 1 - dir_rgb_backgrounds = None - dir_rgb_foregrounds = None - characters_txt_file = None # Directory of characters text file needed for cnn_rnn_ocr model training. The file ends with .txt + dir_output = None # Directory where the augmented training data and the model checkpoints will be saved. + pretraining = False # Set to true to (down)load pretrained weights of ResNet50 encoder. + save_interval = None # frequency for writing model checkpoints (positive integer for number of batches saved under "model_step_{batch:04d}", otherwise epoch saved under "model_{epoch:02d}") + continue_training = False # Whether to continue training an existing model. + if continue_training: + dir_of_start_model = '' # Directory of model checkpoint to load to continue training. (E.g. if you already trained for 3 epochs, set "dir_of_start_model=dir_output/model_03".) + index_start = 0 # Epoch counter initial value to continue training. (E.g. if you already trained for 3 epochs, set "index_start=3" to continue naming checkpoints model_04, model_05 etc.) + data_is_provided = False # Whether the preprocessed input data (subdirectories "images" and "labels" in both subdirectories "train" and "eval" of "dir_output") has already been generated (in the first epoch of a previous run). -@ex.automain -def run( - _config, - n_classes, - n_epochs, - input_height, - input_width, - weight_decay, - weighted_loss, - index_start, - dir_of_start_model, - is_loss_soft_dice, - n_batch, - patches, - augmentation, - flip_aug, - blur_aug, - padding_white, - padding_black, - scaling, - shifting, - degrading, - channels_shuffling, - brightening, - binarization, - adding_rgb_background, - adding_rgb_foreground, - add_red_textlines, - blur_k, - scales, - degrade_scales, - shuffle_indexes, - brightness, - dir_train, - data_is_provided, - scaling_bluring, - scaling_brightness, - scaling_binarization, - rotation, - rotation_not_90, - thetha, - thetha_padd, - scaling_flip, - continue_training, - transformer_projection_dim, - transformer_mlp_head_units, - transformer_layers, - transformer_num_heads, - transformer_cnn_first, - transformer_patchsize_x, - transformer_patchsize_y, - transformer_num_patches_xy, - backbone_type, - save_interval, - flip_index, - dir_eval, - dir_output, - pretraining, - learning_rate, - task, - f1_threshold_classification, - classification_classes_name, - dir_img_bin, - number_of_backgrounds_per_image, - dir_rgb_backgrounds, - dir_rgb_foregrounds, - characters_txt_file, - color_padding_rotation, - bin_deg, - image_inversion, - white_noise_strap, - textline_skewing, - textline_skewing_bin, - textline_left_in_depth, - textline_left_in_depth_bin, - textline_right_in_depth, - textline_right_in_depth_bin, - textline_up_in_depth, - textline_up_in_depth_bin, - textline_down_in_depth, - textline_down_in_depth_bin, - pepper_bin_aug, - pepper_aug, - padd_colors, - pepper_indexes, - white_padds, - skewing_amplitudes, - max_len, +@ex.main +def run(_config, + _log, + task, + pretraining, + data_is_provided, + dir_train, + dir_eval, + dir_output, + n_classes, + n_epochs, + n_batch, + input_height, + input_width, + weight_decay, + learning_rate, + continue_training, + save_interval, + augmentation, + # dependent config keys need a default, + # otherwise yields sacred.utils.ConfigAddedError + ## if rotation_not_90 + thetha=None, + is_loss_soft_dice=False, + weighted_loss=False, + ## if continue_training + index_start=0, + dir_of_start_model=None, + backbone_type=None, + ## if backbone_type=transformer + transformer_projection_dim=None, + transformer_mlp_head_units=None, + transformer_layers=None, + transformer_num_heads=None, + transformer_cnn_first=None, + transformer_patchsize_x=None, + transformer_patchsize_y=None, + transformer_num_patches_xy=None, + ## if task=classification + f1_threshold_classification=None, + classification_classes_name=None, + ## if task=cnn-rnn-ocr + characters_txt_file=None, + color_padding_rotation=False, + thetha_padd=None, + bin_deg=False, + image_inversion=False, + white_noise_strap=False, + textline_skewing=False, + textline_skewing_bin=False, + textline_left_in_depth=False, + textline_left_in_depth_bin=False, + textline_right_in_depth=False, + textline_right_in_depth_bin=False, + textline_up_in_depth=False, + textline_up_in_depth_bin=False, + textline_down_in_depth=False, + textline_down_in_depth_bin=False, + pepper_aug=False, + pepper_bin_aug=False, + pepper_indexes=None, + padd_colors=None, + white_padds=None, + skewing_amplitudes=None, + max_len=None, ): - - if dir_rgb_backgrounds: - list_all_possible_background_images = os.listdir(dir_rgb_backgrounds) - else: - list_all_possible_background_images = None - - if dir_rgb_foregrounds: - list_all_possible_foreground_rgbs = os.listdir(dir_rgb_foregrounds) - else: - list_all_possible_foreground_rgbs = None + """ + run configured experiment via sacred + """ - dir_seg = None - weights = None - model = None - - if task == "segmentation" or task == "enhancement" or task == "binarization": - if data_is_provided: - dir_train_flowing = os.path.join(dir_output, 'train') - dir_eval_flowing = os.path.join(dir_output, 'eval') - + if continue_training: + assert n_epochs > index_start, "with continue_training, n_epochs must be greater than index_start" - dir_flow_train_imgs = os.path.join(dir_train_flowing, 'images') - dir_flow_train_labels = os.path.join(dir_train_flowing, 'labels') + if pretraining and not os.path.isfile(RESNET50_WEIGHTS_PATH): + _log.info("downloading RESNET50 pretrained weights to %s", RESNET50_WEIGHTS_PATH) + download_file(RESNET50_WEIGHTS_URL, RESNET50_WEIGHTS_PATH) - dir_flow_eval_imgs = os.path.join(dir_eval_flowing, 'images') - dir_flow_eval_labels = os.path.join(dir_eval_flowing, 'labels') + # set the gpu configuration + configuration() - configuration() + if task in ["segmentation", "enhancement", "binarization"]: + dir_train_flowing = os.path.join(dir_output, 'train') + dir_eval_flowing = os.path.join(dir_output, 'eval') - else: - dir_img, dir_seg = get_dirs_or_files(dir_train) - dir_img_val, dir_seg_val = get_dirs_or_files(dir_eval) + dir_flow_train_imgs = os.path.join(dir_train_flowing, 'images') + dir_flow_train_labels = os.path.join(dir_train_flowing, 'labels') - # make first a directory in output for both training and evaluations in order to flow data from these directories. - dir_train_flowing = os.path.join(dir_output, 'train') - dir_eval_flowing = os.path.join(dir_output, 'eval') - - dir_flow_train_imgs = os.path.join(dir_train_flowing, 'images/') - dir_flow_train_labels = os.path.join(dir_train_flowing, 'labels/') - - dir_flow_eval_imgs = os.path.join(dir_eval_flowing, 'images/') - dir_flow_eval_labels = os.path.join(dir_eval_flowing, 'labels/') + dir_flow_eval_imgs = os.path.join(dir_eval_flowing, 'images') + dir_flow_eval_labels = os.path.join(dir_eval_flowing, 'labels') + if not data_is_provided: + # first create a directory in output for both training and evaluations + # in order to flow data from these directories. if os.path.isdir(dir_train_flowing): os.system('rm -rf ' + dir_train_flowing) - os.makedirs(dir_train_flowing) - else: - os.makedirs(dir_train_flowing) + os.makedirs(dir_train_flowing) if os.path.isdir(dir_eval_flowing): os.system('rm -rf ' + dir_eval_flowing) - os.makedirs(dir_eval_flowing) - else: - os.makedirs(dir_eval_flowing) + os.makedirs(dir_eval_flowing) os.mkdir(dir_flow_train_imgs) os.mkdir(dir_flow_train_labels) @@ -324,175 +312,182 @@ def run( os.mkdir(dir_flow_eval_imgs) os.mkdir(dir_flow_eval_labels) - # set the gpu configuration - configuration() - - imgs_list=np.array(os.listdir(dir_img)) - segs_list=np.array(os.listdir(dir_seg)) - - imgs_list_test=np.array(os.listdir(dir_img_val)) - segs_list_test=np.array(os.listdir(dir_seg_val)) + dir_img, dir_seg = get_dirs_or_files(dir_train) + dir_img_val, dir_seg_val = get_dirs_or_files(dir_eval) + + imgs_list = list(os.listdir(dir_img)) + segs_list = list(os.listdir(dir_seg)) + + imgs_list_test = list(os.listdir(dir_img_val)) + segs_list_test = list(os.listdir(dir_seg_val)) # writing patches into a sub-folder in order to be flowed from directory. - provide_patches(imgs_list, segs_list, dir_img, dir_seg, dir_flow_train_imgs, - dir_flow_train_labels, input_height, input_width, blur_k, - blur_aug, padding_white, padding_black, flip_aug, binarization, adding_rgb_background,adding_rgb_foreground, add_red_textlines, channels_shuffling, - scaling, shifting, degrading, brightening, scales, degrade_scales, brightness, - flip_index,shuffle_indexes, scaling_bluring, scaling_brightness, scaling_binarization, - rotation, rotation_not_90, thetha, scaling_flip, task, augmentation=augmentation, - patches=patches, dir_img_bin=dir_img_bin,number_of_backgrounds_per_image=number_of_backgrounds_per_image,list_all_possible_background_images=list_all_possible_background_images, dir_rgb_backgrounds=dir_rgb_backgrounds, dir_rgb_foregrounds=dir_rgb_foregrounds,list_all_possible_foreground_rgbs=list_all_possible_foreground_rgbs) - - provide_patches(imgs_list_test, segs_list_test, dir_img_val, dir_seg_val, - dir_flow_eval_imgs, dir_flow_eval_labels, input_height, input_width, - blur_k, blur_aug, padding_white, padding_black, flip_aug, binarization, adding_rgb_background, adding_rgb_foreground, add_red_textlines, channels_shuffling, - scaling, shifting, degrading, brightening, scales, degrade_scales, brightness, - flip_index, shuffle_indexes, scaling_bluring, scaling_brightness, scaling_binarization, - rotation, rotation_not_90, thetha, scaling_flip, task, augmentation=False, patches=patches,dir_img_bin=dir_img_bin,number_of_backgrounds_per_image=number_of_backgrounds_per_image,list_all_possible_background_images=list_all_possible_background_images, dir_rgb_backgrounds=dir_rgb_backgrounds,dir_rgb_foregrounds=dir_rgb_foregrounds,list_all_possible_foreground_rgbs=list_all_possible_foreground_rgbs ) + preprocess_imgs(_config, + imgs_list, + segs_list, + dir_img, + dir_seg, + dir_flow_train_imgs, + dir_flow_train_labels) + preprocess_imgs(_config, + imgs_list_test, + segs_list_test, + dir_img_val, + dir_seg_val, + dir_flow_eval_imgs, + dir_flow_eval_labels, + augmentation=False) if weighted_loss: weights = np.zeros(n_classes) if data_is_provided: - for obj in os.listdir(dir_flow_train_labels): - try: - label_obj = cv2.imread(dir_flow_train_labels + '/' + obj) - label_obj_one_hot = get_one_hot(label_obj, label_obj.shape[0], label_obj.shape[1], n_classes) - weights += (label_obj_one_hot.sum(axis=0)).sum(axis=0) - except: - pass + dirs = dir_flow_train_labels else: - - assert dir_seg is not None - for obj in os.listdir(dir_seg): - try: - label_obj = cv2.imread(dir_seg + '/' + obj) - label_obj_one_hot = get_one_hot(label_obj, label_obj.shape[0], label_obj.shape[1], n_classes) - weights += (label_obj_one_hot.sum(axis=0)).sum(axis=0) - except: - pass + dirs = dir_seg + for obj in os.listdir(dirs): + label_file = os.path.join(dirs, + obj) + try: + label_obj = cv2.imread(label_file) + label_obj_one_hot = get_one_hot(label_obj, label_obj.shape[0], label_obj.shape[1], n_classes) + weights += (label_obj_one_hot.sum(axis=0)).sum(axis=0) + except Exception: + _log.exception("error reading data file '%s'", label_file) weights = 1.00 / weights - weights = weights / float(np.sum(weights)) weights = weights / float(np.min(weights)) weights = weights / float(np.sum(weights)) + if task == "enhancement": + assert not is_loss_soft_dice, "for enhancement, soft_dice loss does not apply" + assert not weighted_loss, "for enhancement, weighted loss does not apply" if continue_training: - if backbone_type=='nontransformer': - if is_loss_soft_dice and (task == "segmentation" or task == "binarization"): - model = load_model(dir_of_start_model, compile=True, custom_objects={'soft_dice_loss': soft_dice_loss}) - if weighted_loss and (task == "segmentation" or task == "binarization"): - model = load_model(dir_of_start_model, compile=True, custom_objects={'loss': weighted_categorical_crossentropy(weights)}) - if not is_loss_soft_dice and not weighted_loss: - model = load_model(dir_of_start_model , compile=True) - elif backbone_type=='transformer': - if is_loss_soft_dice and (task == "segmentation" or task == "binarization"): - model = load_model(dir_of_start_model, compile=True, custom_objects={"PatchEncoder": PatchEncoder, "Patches": Patches,'soft_dice_loss': soft_dice_loss}) - if weighted_loss and (task == "segmentation" or task == "binarization"): - model = load_model(dir_of_start_model, compile=True, custom_objects={'loss': weighted_categorical_crossentropy(weights)}) - if not is_loss_soft_dice and not weighted_loss: - model = load_model(dir_of_start_model , compile=True,custom_objects = {"PatchEncoder": PatchEncoder, "Patches": Patches}) - else: - raise ValueError("backbone_type must be 'nontransformer' or 'transformer'") + custom_objects = dict() + if is_loss_soft_dice: + custom_objects.update(soft_dice_loss=soft_dice_loss) + elif weighted_loss: + custom_objects.update(loss=weighted_categorical_crossentropy(weights)) + if backbone_type == 'transformer': + custom_objects.update(PatchEncoder=PatchEncoder, + Patches=Patches) + model = load_model(dir_of_start_model, compile=False, + custom_objects=custom_objects) else: index_start = 0 - if backbone_type=='nontransformer': - model = resnet50_unet(n_classes, input_height, input_width, task, weight_decay, pretraining) - elif backbone_type=='transformer': + if backbone_type == 'nontransformer': + model = resnet50_unet(n_classes, + input_height, + input_width, + task, + weight_decay, + pretraining) + else: num_patches_x = transformer_num_patches_xy[0] num_patches_y = transformer_num_patches_xy[1] num_patches = num_patches_x * num_patches_y - + if transformer_cnn_first: - if input_height != (num_patches_y * transformer_patchsize_y * 32): - print("Error: transformer_patchsize_y or transformer_num_patches_xy height value error . input_height should be equal to ( transformer_num_patches_xy height value * transformer_patchsize_y * 32)") - sys.exit(1) - if input_width != (num_patches_x * transformer_patchsize_x * 32): - print("Error: transformer_patchsize_x or transformer_num_patches_xy width value error . input_width should be equal to ( transformer_num_patches_xy width value * transformer_patchsize_x * 32)") - sys.exit(1) - if (transformer_projection_dim % (transformer_patchsize_y * transformer_patchsize_x)) != 0: - print("Error: transformer_projection_dim error. The remainder when parameter transformer_projection_dim is divided by (transformer_patchsize_y*transformer_patchsize_x) should be zero") - sys.exit(1) - - - model = vit_resnet50_unet(n_classes, transformer_patchsize_x, transformer_patchsize_y, num_patches, transformer_mlp_head_units, transformer_layers, transformer_num_heads, transformer_projection_dim, input_height, input_width, task, weight_decay, pretraining) + model_builder = vit_resnet50_unet + multiple = 32 else: - if input_height != (num_patches_y * transformer_patchsize_y): - print("Error: transformer_patchsize_y or transformer_num_patches_xy height value error . input_height should be equal to ( transformer_num_patches_xy height value * transformer_patchsize_y)") - sys.exit(1) - if input_width != (num_patches_x * transformer_patchsize_x): - print("Error: transformer_patchsize_x or transformer_num_patches_xy width value error . input_width should be equal to ( transformer_num_patches_xy width value * transformer_patchsize_x)") - sys.exit(1) - if (transformer_projection_dim % (transformer_patchsize_y * transformer_patchsize_x)) != 0: - print("Error: transformer_projection_dim error. The remainder when parameter transformer_projection_dim is divided by (transformer_patchsize_y*transformer_patchsize_x) should be zero") - sys.exit(1) - model = vit_resnet50_unet_transformer_before_cnn(n_classes, transformer_patchsize_x, transformer_patchsize_y, num_patches, transformer_mlp_head_units, transformer_layers, transformer_num_heads, transformer_projection_dim, input_height, input_width, task, weight_decay, pretraining) - + model_builder = vit_resnet50_unet_transformer_before_cnn + multiple = 1 + + assert input_height == ( + num_patches_y * transformer_patchsize_y * multiple), ( + "transformer_patchsize_y or transformer_num_patches_xy height value error: " + "input_height should be equal to " + "(transformer_num_patches_xy height value * transformer_patchsize_y * %d)" % multiple) + assert input_width == ( + num_patches_x * transformer_patchsize_x * multiple), ( + "transformer_patchsize_x or transformer_num_patches_xy width value error: " + "input_width should be equal to " + "(transformer_num_patches_xy width value * transformer_patchsize_x * %d)" % multiple) + assert 0 == (transformer_projection_dim % + (transformer_patchsize_y * transformer_patchsize_x)), ( + "transformer_projection_dim error: " + "The remainder when parameter transformer_projection_dim is divided by " + "(transformer_patchsize_y*transformer_patchsize_x) should be zero") + + model_builder = create_captured_function(model_builder) + model_builder.config = _config + model_builder.logger = _log + model = model_builder(num_patches) + assert model is not None #if you want to see the model structure just uncomment model summary. - model.summary() + #model.summary() - - if task == "segmentation" or task == "binarization": - if not is_loss_soft_dice and not weighted_loss: - model.compile(loss='categorical_crossentropy', - optimizer=Adam(learning_rate=learning_rate), metrics=['accuracy']) - if is_loss_soft_dice: - model.compile(loss=soft_dice_loss, - optimizer=Adam(learning_rate=learning_rate), metrics=['accuracy']) - if weighted_loss: - model.compile(loss=weighted_categorical_crossentropy(weights), - optimizer=Adam(learning_rate=learning_rate), metrics=['accuracy']) - elif task == "enhancement": - model.compile(loss='mean_squared_error', - optimizer=Adam(learning_rate=learning_rate), metrics=['accuracy']) - - - # generating train and evaluation data - train_gen = data_gen(dir_flow_train_imgs, dir_flow_train_labels, batch_size=n_batch, - input_height=input_height, input_width=input_width, n_classes=n_classes, task=task) - val_gen = data_gen(dir_flow_eval_imgs, dir_flow_eval_labels, batch_size=n_batch, - input_height=input_height, input_width=input_width, n_classes=n_classes, task=task) - - ##img_validation_patches = os.listdir(dir_flow_eval_imgs) - ##score_best=[] - ##score_best.append(0) - - save_weights_callback = SaveWeightsAfterSteps(save_interval, dir_output, _config) if save_interval else None - - for i in tqdm(range(index_start, n_epochs + index_start)): - if save_interval: - model.fit( - train_gen, - steps_per_epoch=int(len(os.listdir(dir_flow_train_imgs)) / n_batch) - 1, - validation_data=val_gen, - validation_steps=1, - epochs=1, callbacks=[save_weights_callback]) + if task in ["segmentation", "binarization"]: + if is_loss_soft_dice: + loss = soft_dice_loss + elif weighted_loss: + loss = weighted_categorical_crossentropy(weights) else: - model.fit( - train_gen, - steps_per_epoch=int(len(os.listdir(dir_flow_train_imgs)) / n_batch) - 1, - validation_data=val_gen, - validation_steps=1, - epochs=1) - - model.save(os.path.join(dir_output,'model_'+str(i))) - - with open(os.path.join(os.path.join(dir_output,'model_'+str(i)),"config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON + loss = 'categorical_crossentropy' + else: # task == "enhancement" + loss = 'mean_squared_error' + model.compile(loss=loss, + optimizer=Adam(learning_rate=learning_rate), + metrics=['accuracy', MeanIoU(n_classes, + name='iou', + ignore_class=0, + sparse_y_true=False, + sparse_y_pred=False)]) - #os.system('rm -rf '+dir_train_flowing) - #os.system('rm -rf '+dir_eval_flowing) + def get_dataset(dir_imgs, dir_labs, shuffle=None): + gen_kwargs = dict(labels=None, + label_mode=None, + batch_size=1, # batch after zip below + image_size=(input_height, input_width), + color_mode='rgb', + shuffle=shuffle is not None, + seed=shuffle, + interpolation='nearest', + crop_to_aspect_ratio=False, + # Keras 3 only... + #pad_to_aspect_ratio=False, + #data_format='channel_last', + #verbose=False, + ) + img_gen = image_dataset_from_directory(dir_imgs, **gen_kwargs) + lab_gen = image_dataset_from_directory(dir_labs, **gen_kwargs) + if task in ["segmentation", "binarization"]: + @tf.function + def to_categorical(seg): + seg = tf.image.rgb_to_grayscale(seg) + seg = tf.cast(seg, tf.int8) + seg = tf.squeeze(seg, axis=-1) + return one_hot(seg, n_classes) + lab_gen = lab_gen.map(to_categorical) + return tf.data.Dataset.zip(img_gen, lab_gen).rebatch(n_batch, drop_remainder=True) + train_gen = get_dataset(dir_flow_train_imgs, dir_flow_train_labels, shuffle=np.random.randint(1e6)) + val_gen = get_dataset(dir_flow_eval_imgs, dir_flow_eval_labels) + + callbacks = [TensorBoard(os.path.join(dir_output, 'logs'), write_graph=False), + SaveWeightsAfterSteps(0, dir_output, _config)] + if save_interval: + callbacks.append(SaveWeightsAfterSteps(save_interval, dir_output, _config)) + model.fit( + train_gen.prefetch(tf.data.AUTOTUNE), # .repeat()?? + validation_data=val_gen.prefetch(tf.data.AUTOTUNE), + verbose=1, + epochs=n_epochs, + callbacks=callbacks, + initial_epoch=index_start) - #model.save(dir_output+'/'+'model'+'.h5') - elif task=="cnn-rnn-ocr": + dir_img, dir_lab = get_dirs_or_files(dir_train) - + dir_img_val, dir_lab_val = get_dirs_or_files(dir_eval) + imgs_list = list(os.listdir(dir_img)) + labs_list = list(os.listdir(dir_lab)) + imgs_list_val = list(os.listdir(dir_img_val)) + labs_list_val = list(os.listdir(dir_lab_val)) + with open(characters_txt_file, 'r') as char_txt_f: characters = json.load(char_txt_f) - - AUTOTUNE = tf.data.AUTOTUNE - + padding_token = len(characters) + 5 # Mapping characters to integers. char_to_num = StringLookup(vocabulary=list(characters), mask_token=None) @@ -500,128 +495,122 @@ def run( ##num_to_char = StringLookup( ##vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True ##) - - padding_token = len(characters) + 5 - ls_files_images = os.listdir(dir_img) - n_classes = len(char_to_num.get_vocabulary()) + 2 - + if continue_training: model = load_model(dir_of_start_model) else: index_start = 0 - model = cnn_rnn_ocr_model(image_height=input_height, image_width=input_width, n_classes=n_classes, max_seq=max_len) - - print(model.summary()) - - aug_multip = return_multiplier_based_on_augmnentations(augmentation, color_padding_rotation, rotation_not_90, blur_aug, degrading, bin_deg, - brightening, padding_white, adding_rgb_foreground, adding_rgb_background, binarization, - image_inversion, channels_shuffling, add_red_textlines, white_noise_strap, textline_skewing, textline_skewing_bin, textline_left_in_depth, textline_left_in_depth_bin, textline_right_in_depth, textline_right_in_depth_bin, textline_up_in_depth, textline_up_in_depth_bin, textline_down_in_depth, textline_down_in_depth_bin, pepper_bin_aug, pepper_aug, degrade_scales, number_of_backgrounds_per_image, thetha, thetha_padd, brightness, padd_colors, shuffle_indexes, pepper_indexes, skewing_amplitudes, blur_k, white_padds) - - len_dataset = aug_multip*len(ls_files_images) - - train_ds = data_gen_ocr(padding_token, n_batch, input_height, input_width, max_len, dir_train, ls_files_images, - augmentation, color_padding_rotation, rotation_not_90, blur_aug, degrading, bin_deg, brightening, padding_white, - adding_rgb_foreground, adding_rgb_background, binarization, image_inversion, channels_shuffling, add_red_textlines, white_noise_strap, - textline_skewing, textline_skewing_bin, textline_left_in_depth, textline_left_in_depth_bin, textline_right_in_depth, - textline_right_in_depth_bin, textline_up_in_depth, textline_up_in_depth_bin, textline_down_in_depth, textline_down_in_depth_bin, - pepper_bin_aug, pepper_aug, degrade_scales, number_of_backgrounds_per_image, thetha, thetha_padd, brightness, padd_colors, - shuffle_indexes, pepper_indexes, skewing_amplitudes, blur_k, char_to_num, list_all_possible_background_images, list_all_possible_foreground_rgbs, - dir_rgb_backgrounds, dir_rgb_foregrounds, white_padds, dir_img_bin) - - initial_learning_rate = 1e-4 - decay_steps = int (n_epochs * ( len_dataset / n_batch )) - alpha = 0.01 - lr_schedule = 1e-4#tf.keras.optimizers.schedules.CosineDecay(initial_learning_rate, decay_steps, alpha) + model = cnn_rnn_ocr_model(image_height=input_height, + image_width=input_width, + n_classes=n_classes, + max_seq=max_len) + #print(model.summary()) + + # todo: use Dataset.map() on Dataset.list_files() + # todo: test_ds + def gen(): + return preprocess_imgs(_config, + imgs_list, + labs_list, + dir_img, + dir_lab, + None, # no file I/O, but in-memory + None, # no file I/O, but in-memory + # extra+overrides + char_to_num=char_to_num, + padding_token=padding_token + ) + train_ds = tf.data.Dataset.from_generator(gen) + train_ds = train_ds.padded_batch(n_batch, + padded_shapes=([input_height, input_width, 3], [None]), + padding_values=(0, padding_token), + drop_remainder=True, + #num_parallel_calls=tf.data.AUTOTUNE, + ) + train_ds = train_ds.prefetch(tf.data.AUTOTUNE) + + #initial_learning_rate = 1e-4 + #decay_steps = int (n_epochs * ( len_dataset / n_batch )) + #alpha = 0.01 + #lr_schedule = 1e-4 + #tf.keras.optimizers.schedules.CosineDecay(initial_learning_rate, decay_steps, alpha) + opt = tf.keras.optimizers.Adam(learning_rate=learning_rate) + model.compile(optimizer=opt) # rs: loss seems to be (ctc_batch_cost) in last layer + + callbacks = [TensorBoard(os.path.join(dir_output, 'logs'), write_graph=False), + SaveWeightsAfterSteps(0, dir_output, _config)] + if save_interval: + callbacks.append(SaveWeightsAfterSteps(save_interval, dir_output, _config)) + model.fit( + train_ds, + #validation_data=test_ds, + verbose=1, + epochs=n_epochs, + callbacks=callbacks, + initial_epoch=index_start) - opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)#1e-4)#(lr_schedule) - model.compile(optimizer=opt) - - save_weights_callback = SaveWeightsAfterSteps(save_interval, dir_output, _config) if save_interval else None - - for i in tqdm(range(index_start, n_epochs + index_start)): - if save_interval: - model.fit( - train_ds, - steps_per_epoch=len_dataset / n_batch, - epochs=1, - callbacks=[save_weights_callback] - ) - else: - model.fit( - train_ds, - steps_per_epoch=len_dataset / n_batch, - epochs=1 - ) - - if i >=0: - model.save( os.path.join(dir_output,'model_'+str(i) )) - with open(os.path.join(os.path.join(dir_output,'model_'+str(i)),"config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - elif task=='classification': - configuration() - model = resnet50_classifier(n_classes, input_height, input_width, weight_decay, pretraining) + if continue_training: + model = load_model(dir_of_start_model, compile=False) + else: + index_start = 0 + model = resnet50_classifier(n_classes, + input_height, + input_width, + weight_decay, + pretraining) - opt_adam = Adam(learning_rate=0.001) model.compile(loss='categorical_crossentropy', - optimizer = opt_adam,metrics=['accuracy']) + optimizer=Adam(learning_rate=0.001), # rs: why not learning_rate? + metrics=['accuracy', F1Score(average='macro', name='f1')]) - list_classes = list(classification_classes_name.values()) - testX, testY = generate_data_from_folder_evaluation(dir_eval, input_height, input_width, n_classes, list_classes) + data_args = dict(label_mode="categorical", + class_names=list_classes, + batch_size=n_batch, + image_size=(input_height, input_width), + interpolation="nearest") + trainXY = image_dataset_from_directory(dir_train, shuffle=True, **data_args) + testXY = image_dataset_from_directory(dir_eval, shuffle=False, **data_args) + callbacks = [TensorBoard(os.path.join(dir_output, 'logs'), write_graph=False), + SaveWeightsAfterSteps(0, dir_output, _config, + monitor='val_f1', + #save_best_only=True, # we need all for ensembling + mode='max')] - y_tot=np.zeros((testX.shape[0],n_classes)) + history = model.fit(trainXY, + #class_weight=weights) + validation_data=testXY, + verbose=1, + epochs=n_epochs, + callbacks=callbacks, + initial_epoch=index_start) - score_best= [0] + usable_checkpoints = np.flatnonzero(np.array(history.history['val_f1']) > + f1_threshold_classification) + if len(usable_checkpoints) >= 1: + _log.info("averaging over usable checkpoints: %s", str(usable_checkpoints)) + usable_checkpoints = [os.path.join(dir_output, 'model_{epoch:02d}'.format(epoch=epoch + 1)) + for epoch in usable_checkpoints] + ens_path = os.path.join(dir_output, 'model_ens_avg') + run_ensembling(usable_checkpoints, ens_path) + _log.info("ensemble model saved under '%s'", ens_path) - num_rows = return_number_of_total_training_data(dir_train) - weights=[] - - for i in range(n_epochs): - history = model.fit( generate_data_from_folder_training(dir_train, n_batch , input_height, input_width, n_classes, list_classes), steps_per_epoch=num_rows / n_batch, verbose=1)#,class_weight=weights) - - y_pr_class = [] - for jj in range(testY.shape[0]): - y_pr=model.predict(testX[jj,:,:,:].reshape(1,input_height,input_width,3), verbose=0) - y_pr_ind= np.argmax(y_pr,axis=1) - y_pr_class.append(y_pr_ind) - - y_pr_class = np.array(y_pr_class) - f1score=f1_score(np.argmax(testY,axis=1), y_pr_class, average='macro') - print(i,f1score) - - if f1score>score_best[0]: - score_best[0]=f1score - model.save(os.path.join(dir_output,'model_best')) - - if f1score > f1_threshold_classification: - weights.append(model.get_weights() ) - - - if len(weights) >= 1: - new_weights=list() - for weights_list_tuple in zip(*weights): - new_weights.append( [np.array(weights_).mean(axis=0) for weights_ in zip(*weights_list_tuple)] ) - - new_weights = [np.array(x) for x in new_weights] - model_weight_averaged=tf.keras.models.clone_model(model) - model_weight_averaged.set_weights(new_weights) - - model_weight_averaged.save(os.path.join(dir_output,'model_ens_avg')) - with open(os.path.join( os.path.join(dir_output,'model_ens_avg'), "config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - - with open(os.path.join( os.path.join(dir_output,'model_best'), "config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - elif task=='reading_order': - configuration() - model = machine_based_reading_order_model(n_classes,input_height,input_width,weight_decay,pretraining) - + if continue_training: + model = load_model(dir_of_start_model, compile=False) + else: + index_start = 0 + model = machine_based_reading_order_model(n_classes, + input_height, + input_width, + weight_decay, + pretraining) + dir_flow_train_imgs = os.path.join(dir_train, 'images') dir_flow_train_labels = os.path.join(dir_train, 'labels') - + classes = os.listdir(dir_flow_train_labels) if augmentation: num_rows = len(classes)*(len(thetha) + 1) @@ -630,28 +619,30 @@ def run( #ls_test = os.listdir(dir_flow_train_labels) #f1score_tot = [0] - indexer_start = 0 - # opt = SGD(learning_rate=0.01, momentum=0.9) - opt_adam = tf.keras.optimizers.Adam(learning_rate=0.0001) model.compile(loss="binary_crossentropy", - optimizer = opt_adam,metrics=['accuracy']) - - save_weights_callback = SaveWeightsAfterSteps(save_interval, dir_output, _config) if save_interval else None - - for i in range(n_epochs): - if save_interval: - history = model.fit(generate_arrays_from_folder_reading_order(dir_flow_train_labels, dir_flow_train_imgs, n_batch, input_height, input_width, n_classes, thetha, augmentation), steps_per_epoch=num_rows / n_batch, verbose=1, callbacks=[save_weights_callback]) - else: - history = model.fit(generate_arrays_from_folder_reading_order(dir_flow_train_labels, dir_flow_train_imgs, n_batch, input_height, input_width, n_classes, thetha, augmentation), steps_per_epoch=num_rows / n_batch, verbose=1) - model.save( os.path.join(dir_output,'model_'+str(i+indexer_start) )) - - with open(os.path.join(os.path.join(dir_output,'model_'+str(i)),"config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - ''' - if f1score>f1score_tot[0]: - f1score_tot[0] = f1score - model_dir = os.path.join(dir_out,'model_best') - model.save(model_dir) - ''' + #optimizer=SGD(learning_rate=0.01, momentum=0.9), + optimizer=Adam(learning_rate=0.0001), # rs: why not learning_rate? + metrics=['accuracy']) - + callbacks = [TensorBoard(os.path.join(dir_output, 'logs'), write_graph=False), + SaveWeightsAfterSteps(0, dir_output, _config)] + if save_interval: + callbacks.append(SaveWeightsAfterSteps(save_interval, dir_output, _config)) + + trainXY = generate_arrays_from_folder_reading_order( + dir_flow_train_labels, dir_flow_train_imgs, + n_batch, input_height, input_width, n_classes, + thetha, augmentation) + + history = model.fit(trainXY, + steps_per_epoch=num_rows / n_batch, + verbose=1, + epochs=n_epochs, + callbacks=callbacks, + initial_epoch=index_start) + ''' + if f1score>f1score_tot[0]: + f1score_tot[0] = f1score + model_dir = os.path.join(dir_out,'model_best') + model.save(model_dir) + ''' diff --git a/src/eynollah/training/utils.py b/src/eynollah/training/utils.py index 005810f..4b6033e 100644 --- a/src/eynollah/training/utils.py +++ b/src/eynollah/training/utils.py @@ -1,7 +1,9 @@ import os import math import random +from logging import getLogger from pathlib import Path + import cv2 import numpy as np import seaborn as sns @@ -10,11 +12,12 @@ from scipy.ndimage.filters import gaussian_filter from tqdm import tqdm import imutils import tensorflow as tf -from tensorflow.keras.utils import to_categorical + from PIL import Image, ImageFile, ImageEnhance ImageFile.LOAD_TRUNCATED_IMAGES = True + def vectorize_label(label, char_to_num, padding_token, max_len): label = char_to_num(tf.strings.unicode_split(label, input_encoding="UTF-8")) length = tf.shape(label)[0] @@ -375,111 +378,10 @@ def return_number_of_total_training_data(path_classes): n_tot = n_tot + len(sub_files) return n_tot - - -def generate_data_from_folder_evaluation(path_classes, height, width, n_classes, list_classes): - #sub_classes = os.listdir(path_classes) - #n_classes = len(sub_classes) - all_imgs = [] - labels = [] - #dicts =dict() - #indexer= 0 - for indexer, sub_c in enumerate(list_classes): - sub_files = os.listdir(os.path.join(path_classes,sub_c )) - sub_files = [os.path.join(path_classes,sub_c )+'/' + x for x in sub_files] - #print( os.listdir(os.path.join(path_classes,sub_c )) ) - all_imgs = all_imgs + sub_files - sub_labels = list( np.zeros( len(sub_files) ) +indexer ) - #print( len(sub_labels) ) - labels = labels + sub_labels - #dicts[sub_c] = indexer - #indexer +=1 - - - categories = to_categorical(range(n_classes)).astype(np.int16)#[ [1 , 0, 0 , 0 , 0 , 0] , [0 , 1, 0 , 0 , 0 , 0] , [0 , 0, 1 , 0 , 0 , 0] , [0 , 0, 0 , 1 , 0 , 0] , [0 , 0, 0 , 0 , 1 , 0] , [0 , 0, 0 , 0 , 0 , 1] ] - ret_x= np.zeros((len(labels), height,width, 3)).astype(np.int16) - ret_y= np.zeros((len(labels), n_classes)).astype(np.int16) - - #print(all_imgs) - for i in range(len(all_imgs)): - row = all_imgs[i] - #####img = cv2.imread(row, 0) - #####img= resize_image (img, height, width) - #####img = img.astype(np.uint16) - #####ret_x[i, :,:,0] = img[:,:] - #####ret_x[i, :,:,1] = img[:,:] - #####ret_x[i, :,:,2] = img[:,:] - - img = cv2.imread(row) - img= resize_image (img, height, width) - img = img.astype(np.uint16) - ret_x[i, :,:] = img[:,:,:] - - ret_y[i, :] = categories[ int( labels[i] ) ][:] - - return ret_x/255., ret_y - -def generate_data_from_folder_training(path_classes, n_batch, height, width, n_classes, list_classes): - #sub_classes = os.listdir(path_classes) - #n_classes = len(sub_classes) - - all_imgs = [] - labels = [] - #dicts =dict() - #indexer= 0 - for indexer, sub_c in enumerate(list_classes): - sub_files = os.listdir(os.path.join(path_classes,sub_c )) - sub_files = [os.path.join(path_classes,sub_c )+'/' + x for x in sub_files] - #print( os.listdir(os.path.join(path_classes,sub_c )) ) - all_imgs = all_imgs + sub_files - sub_labels = list( np.zeros( len(sub_files) ) +indexer ) - - #print( len(sub_labels) ) - labels = labels + sub_labels - #dicts[sub_c] = indexer - #indexer +=1 - - ids = np.array(range(len(labels))) - random.shuffle(ids) - - shuffled_labels = np.array(labels)[ids] - shuffled_files = np.array(all_imgs)[ids] - categories = to_categorical(range(n_classes)).astype(np.int16)#[ [1 , 0, 0 , 0 , 0 , 0] , [0 , 1, 0 , 0 , 0 , 0] , [0 , 0, 1 , 0 , 0 , 0] , [0 , 0, 0 , 1 , 0 , 0] , [0 , 0, 0 , 0 , 1 , 0] , [0 , 0, 0 , 0 , 0 , 1] ] - ret_x= np.zeros((n_batch, height,width, 3)).astype(np.int16) - ret_y= np.zeros((n_batch, n_classes)).astype(np.int16) - batchcount = 0 - while True: - for i in range(len(shuffled_files)): - row = shuffled_files[i] - #print(row) - ###img = cv2.imread(row, 0) - ###img= resize_image (img, height, width) - ###img = img.astype(np.uint16) - ###ret_x[batchcount, :,:,0] = img[:,:] - ###ret_x[batchcount, :,:,1] = img[:,:] - ###ret_x[batchcount, :,:,2] = img[:,:] - - img = cv2.imread(row) - img= resize_image (img, height, width) - img = img.astype(np.uint16) - ret_x[batchcount, :,:,:] = img[:,:,:] - - #print(int(shuffled_labels[i]) ) - #print( categories[int(shuffled_labels[i])] ) - ret_y[batchcount, :] = categories[ int( shuffled_labels[i] ) ][:] - - batchcount+=1 - - if batchcount>=n_batch: - ret_x = ret_x/255. - yield ret_x, ret_y - ret_x= np.zeros((n_batch, height,width, 3)).astype(np.int16) - ret_y= np.zeros((n_batch, n_classes)).astype(np.int16) - batchcount = 0 - -def do_brightening(img_in_dir, factor): - im = Image.open(img_in_dir) +def do_brightening(img, factor): + img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + im = Image.fromarray(img_rgb) enhancer = ImageEnhance.Brightness(im) out_img = enhancer.enhance(factor) out_img = out_img.convert('RGB') @@ -694,44 +596,6 @@ def generate_arrays_from_folder_reading_order(classes_file_dir, modal_dir, n_bat ret_y= np.zeros((n_batch, n_classes)).astype(np.int16) batchcount = 0 -def data_gen(img_folder, mask_folder, batch_size, input_height, input_width, n_classes, task='segmentation'): - c = 0 - n = [f for f in os.listdir(img_folder) if not f.startswith('.')] # os.listdir(img_folder) #List of training images - random.shuffle(n) - while True: - img = np.zeros((batch_size, input_height, input_width, 3)).astype('float') - mask = np.zeros((batch_size, input_height, input_width, n_classes)).astype('float') - - for i in range(c, c + batch_size): # initially from 0 to 16, c = 0. - try: - filename = os.path.splitext(n[i])[0] - - train_img = cv2.imread(img_folder + '/' + n[i]) / 255. - train_img = cv2.resize(train_img, (input_width, input_height), - interpolation=cv2.INTER_NEAREST) # Read an image from folder and resize - - img[i - c] = train_img # add to array - img[0], img[1], and so on. - if task == "segmentation" or task=="binarization": - train_mask = cv2.imread(mask_folder + '/' + filename + '.png') - train_mask = get_one_hot(resize_image(train_mask, input_height, input_width), input_height, input_width, - n_classes) - elif task == "enhancement": - train_mask = cv2.imread(mask_folder + '/' + filename + '.png')/255. - train_mask = resize_image(train_mask, input_height, input_width) - - # train_mask = train_mask.reshape(224, 224, 1) # Add extra dimension for parity with train_img size [512 * 512 * 3] - - mask[i - c] = train_mask - except: - img[i - c] = np.ones((input_height, input_width, 3)).astype('float') - mask[i - c] = np.zeros((input_height, input_width, n_classes)).astype('float') - - c += batch_size - if c + batch_size >= len(os.listdir(img_folder)): - c = 0 - random.shuffle(n) - yield img, mask - # TODO: Use otsu_copy from utils def otsu_copy(img): @@ -748,7 +612,7 @@ def otsu_copy(img): return img_r -def get_patches(dir_img_f, dir_seg_f, img, label, height, width, indexer): +def get_patches(img, label, height, width): if img.shape[0] < height or img.shape[1] < width: img, label = do_padding(img, label, height, width) @@ -784,11 +648,7 @@ def get_patches(dir_img_f, dir_seg_f, img, label, height, width, indexer): img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] label_patch = label[index_y_d:index_y_u, index_x_d:index_x_u, :] - cv2.imwrite(dir_img_f + '/img_' + str(indexer) + '.png', img_patch) - cv2.imwrite(dir_seg_f + '/img_' + str(indexer) + '.png', label_patch) - indexer += 1 - - return indexer + yield img_patch, label_patch def do_padding_with_color(img, padding_color='black'): @@ -848,7 +708,7 @@ def do_padding(img, label, height, width): return img_new,label_new -def get_patches_num_scale_new(dir_img_f, dir_seg_f, img, label, height, width, indexer, scaler): +def get_patches_num_scale_new(img, label, height, width, scaler=1.0): img = resize_image(img, int(img.shape[0] * scaler), int(img.shape[1] * scaler)) label = resize_image(label, int(label.shape[0] * scaler), int(label.shape[1] * scaler)) @@ -890,781 +750,502 @@ def get_patches_num_scale_new(dir_img_f, dir_seg_f, img, label, height, width, i img_patch = img[index_y_d:index_y_u, index_x_d:index_x_u, :] label_patch = label[index_y_d:index_y_u, index_x_d:index_x_u, :] - cv2.imwrite(dir_img_f + '/img_' + str(indexer) + '.png', img_patch) - cv2.imwrite(dir_seg_f + '/img_' + str(indexer) + '.png', label_patch) - indexer += 1 - - return indexer + yield img_patch, label_patch -# TODO: (far) too many args # TODO: refactor to combine with data_gen_ocr -def provide_patches( - imgs_list_train, - segs_list_train, - dir_img, - dir_seg, - dir_flow_train_imgs, - dir_flow_train_labels, - input_height, - input_width, - blur_k, - blur_aug, - padding_white, - padding_black, - flip_aug, - binarization, - adding_rgb_background, - adding_rgb_foreground, - add_red_textlines, - channels_shuffling, - scaling, - shifting, - degrading, - brightening, - scales, - degrade_scales, - brightness, - flip_index, - shuffle_indexes, - scaling_bluring, - scaling_brightness, - scaling_binarization, - rotation, - rotation_not_90, - thetha, - scaling_flip, - task, +def preprocess_imgs(config, + imgs_list, + labs_list, + dir_img, + dir_lab, + dir_flow_imgs, + dir_flow_lbls, + logger=None, + **kwargs, +): + if logger is None: + logger = getLogger('') + + # make a copy for this run + config = dict(config) + # add derived keys not part of config + if config.get('dir_rgb_backgrounds', None): + config['list_all_possible_background_images'] = \ + os.listdir(config['dir_rgb_backgrounds']) + if config.get('dir_rgb_foregrounds', None): + config['list_all_possible_foreground_rgbs'] = \ + os.listdir(config['dir_rgb_foregrounds']) + # override keys from call + config.update(kwargs) + + seed = random.random() + random.shuffle(imgs_list, random=lambda: seed) + random.shuffle(labs_list, random=lambda: seed) + + # labs_list not used because stem matching more robust + indexer = 0 + for img, lab in tqdm(zip(imgs_list, labs_list)): + img = cv2.imread(os.path.join(dir_img, img)) + img_name = os.path.splitext(img)[0] + if config['task'] in ["segmentation", "binarization"]: + # assert lab == img_name + '.png' + lab = cv2.imread(os.path.join(dir_lab, img_name + '.png')) + elif config['task'] == "enhancement": + lab = cv2.imread(os.path.join(dir_lab, img)) + elif config['task'] == "cnn-rnn-ocr": + # assert lab == 'img_name + '.txt' + with open(os.path.join(dir_lab, img_name + '.txt'), 'r') as f: + lab = f.read().split('\n')[0] + else: + lab = None + + try: + if config['task'] == "cnn-rnn-ocr": + yield from preprocess_img_ocr(img, img_name, lab, + **config) + continue + for img, lab in preprocess_img(img, img_name, lab, + **config): + cv2.imwrite(os.path.join(dir_flow_imgs, '/img_%d.png' % indexer), + resize_image(img, + config['input_height'], + config['input_width'])) + cv2.imwrite(os.path.join(dir_flow_lbls, '/img_%d.png' % indexer), + resize_image(lab, + config['input_height'], + config['input_width'])) + indexer += 1 + except: + logger.exception("skipping image %s", img_name) + +def preprocess_img(img, + img_name, + lab, + input_height=None, + input_width=None, + augmentation=False, + flip_aug=False, + flip_index=None, + blur_aug=False, + blur_k=None, + padding_white=False, + padding_black=False, + scaling=False, + scaling_bluring=False, + scaling_brightness=False, + scaling_binarization=False, + scaling_flip=False, + scales=None, + shifting=False, + degrading=False, + degrade_scales=None, + brightening=False, + brightness=None, + binarization=False, + dir_img_bin=None, + add_red_textlines=False, + adding_rgb_background=False, + dir_rgb_backgrounds=None, + adding_rgb_foreground=False, + dir_rgb_foregrounds=None, + number_of_backgrounds_per_image=None, + channels_shuffling=False, + shuffle_indexes=None, + rotation=False, + rotation_not_90=False, + thetha=None, + patches=False, + list_all_possible_background_images=None, + list_all_possible_foreground_rgbs=None, + **kwargs, +): + if not patches: + yield img, lab + if augmentation: + if flip_aug: + for f_i in flip_index: + yield cv2.flip(img, f_i), cv2.flip(lab, f_i) + if blur_aug: + for blur_i in blur_k: + yield bluring(img, blur_i), lab + if brightening: + for factor in brightness: + yield do_brightening(img, factor), lab + if binarization: + if dir_img_bin: + img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') + else: + img_bin_corr = otsu_copy(img) + yield img_bin_corr, lab + if degrading: + for degrade_scale_ind in degrade_scales: + yield do_degrading(img, degrade_scale_ind), lab + if rotation_not_90: + for thetha_i in thetha: + yield rotation_not_90_func(img, lab, thetha_i) + if channels_shuffling: + for shuffle_index in shuffle_indexes: + yield return_shuffled_channels(img, shuffle_index), lab + if scaling: + for sc_ind in scales: + yield scale_image_for_no_patch(img, lab, sc_ind) + if shifting: + shift_types = ['xpos', 'xmin', 'ypos', 'ymin', 'xypos', 'xymin'] + for st_ind in shift_types: + yield shift_image_and_label(img, lab, st_ind) + if adding_rgb_background: + img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') + for i_n in range(number_of_backgrounds_per_image): + background_image_chosen_name = random.choice(list_all_possible_background_images) + img_rgb_background_chosen = \ + cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) + img_with_overlayed_background = \ + return_binary_image_with_given_rgb_background( + img_bin_corr, img_rgb_background_chosen) + yield img_with_overlayed_background, lab + if adding_rgb_foreground: + img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') + for i_n in range(number_of_backgrounds_per_image): + background_image_chosen_name = random.choice(list_all_possible_background_images) + foreground_rgb_chosen_name = random.choice(list_all_possible_foreground_rgbs) + img_rgb_background_chosen = \ + cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) + foreground_rgb_chosen = \ + np.load(dir_rgb_foregrounds + '/' + foreground_rgb_chosen_name) + img_with_overlayed_background = \ + return_binary_image_with_given_rgb_background_and_given_foreground_rgb( + img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) + yield img_with_overlayed_background, lab + if add_red_textlines: + img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') + yield return_image_with_red_elements(img, img_bin_corr), lab + else: + yield from get_patches(img, + lab, + input_height, + input_width) + if augmentation: + if rotation: + yield from get_patches(rotation_90(img), + rotation_90(lab), + input_height, + input_width) + if rotation_not_90: + for thetha_i in thetha: + img_max_rotated, label_max_rotated = \ + rotation_not_90_func(img, lab, thetha_i) + yield from get_patches(img_max_rotated, + label_max_rotated, + input_height, + input_width) + if channels_shuffling: + for shuffle_index in shuffle_indexes: + img_shuffled = \ + return_shuffled_channels(img, shuffle_index), + yield from get_patches(img_shuffled, + lab, + input_height, + input_width) + if adding_rgb_background: + img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') + for i_n in range(number_of_backgrounds_per_image): + background_image_chosen_name = random.choice(list_all_possible_background_images) + img_rgb_background_chosen = \ + cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) + img_with_overlayed_background = \ + return_binary_image_with_given_rgb_background( + img_bin_corr, img_rgb_background_chosen) + yield from get_patches(img_with_overlayed_background, + lab, + input_height, + input_width) + if adding_rgb_foreground: + img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') + for i_n in range(number_of_backgrounds_per_image): + background_image_chosen_name = random.choice(list_all_possible_background_images) + foreground_rgb_chosen_name = random.choice(list_all_possible_foreground_rgbs) + img_rgb_background_chosen = \ + cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) + foreground_rgb_chosen = \ + np.load(dir_rgb_foregrounds + '/' + foreground_rgb_chosen_name) + img_with_overlayed_background = \ + return_binary_image_with_given_rgb_background_and_given_foreground_rgb( + img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) + yield from get_patches(img_with_overlayed_background, + lab, + input_height, + input_width) + if add_red_textlines: + img_bin_corr = cv2.imread(os.path.join(dir_img_bin, img_name + '.png')) + img_red_context = \ + return_image_with_red_elements(img, img_bin_corr) + yield from get_patches(img_red_context, + lab, + input_height, + input_width) + if flip_aug: + for f_i in flip_index: + yield from get_patches(cv2.flip(img, f_i), + cv2.flip(lab, f_i), + input_height, + input_width) + if blur_aug: + for blur_i in blur_k: + yield from get_patches(bluring(img, blur_i), + lab, + input_height, + input_width) + if padding_black: + yield from get_patches(do_padding_with_color(img, 'black'), + do_padding_label(lab), + input_height, + input_width) + if padding_white: + yield from get_patches(do_padding_with_color(img, 'white'), + do_padding_label(lab), + input_height, + input_width) + if brightening: + for factor in brightness: + yield from get_patches(do_brightening(img, factor), + lab, + input_height, + input_width) + if scaling: + for sc_ind in scales: + yield from get_patches_num_scale_new(img, + lab, + input_height, + input_width, + scaler=sc_ind) + if degrading: + for degrade_scale_ind in degrade_scales: + img_deg = \ + do_degrading(img, degrade_scale_ind), + yield from get_patches(img_deg, + lab, + input_height, + input_width) + if binarization: + if dir_img_bin: + img_bin_corr = cv2.imread(os.path.join(dir_img_bin, img_name + '.png')) + else: + img_bin_corr = otsu_copy(img) + yield from get_patches(img_bin_corr, + lab, + input_height, + input_width) + if scaling_brightness: + for sc_ind in scales: + for factor in brightness: + img_bright = do_brightening(img, factor) + yield from get_patches_num_scale_new(img_bright, + lab, + input_height, + input_width, + scaler=sc_ind) + if scaling_bluring: + for sc_ind in scales: + for blur_i in blur_k: + img_blur = bluring(img, blur_i), + yield from get_patches_num_scale_new(img_blur, + lab, + input_height, + input_width, + scaler=sc_ind) + if scaling_binarization: + for sc_ind in scales: + img_bin = otsu_copy(img), + yield from get_patches_num_scale_new(img_bin, + lab, + input_height, + input_width, + scaler=sc_ind) + if scaling_flip: + for sc_ind in scales: + for f_i in flip_index: + yield from get_patches_num_scale_new(cv2.flip(img, f_i), + cv2.flip(lab, f_i), + input_height, + input_width, + scaler=sc_ind) + +def preprocess_img_ocr( + img, + img_name, + lab, + char_to_num=None, + padding_token=-1, + max_len=500, + n_batch=1, + input_height=None, + input_width=None, augmentation=False, - patches=False, + color_padding_rotation=None, + thetha_padd=None, + padd_colors=None, + rotation_not_90=None, + thetha=None, + padding_white=None, + white_padds=None, + degrading=False, + bin_deg=None, + degrade_scales=None, + blur_aug=False, + blur_k=None, + brightening=False, + brightness=None, + binarization=False, + image_inversion=False, + channels_shuffling=False, + shuffle_indexes=None, + white_noise_strap=False, + textline_skewing=False, + textline_skewing_bin=False, + skewing_amplitudes=None, + textline_left_in_depth=False, + textline_left_in_depth_bin=False, + textline_right_in_depth=False, + textline_right_in_depth_bin=False, + textline_up_in_depth=False, + textline_up_in_depth_bin=False, + textline_down_in_depth=False, + textline_down_in_depth_bin=False, + pepper_aug=False, + pepper_bin_aug=False, + pepper_indexes=None, dir_img_bin=None, + add_red_textlines=False, + adding_rgb_background=False, + dir_rgb_backgrounds=None, + adding_rgb_foreground=False, + dir_rgb_foregrounds=None, number_of_backgrounds_per_image=None, list_all_possible_background_images=None, - dir_rgb_backgrounds=None, - dir_rgb_foregrounds=None, list_all_possible_foreground_rgbs=None, ): - - # TODO: why sepoarate var if you have seg_i? - indexer = 0 - for im, seg_i in tqdm(zip(imgs_list_train, segs_list_train)): - img_name = os.path.splitext(im)[0] - if task == "segmentation" or task == "binarization": - dir_of_label_file = os.path.join(dir_seg, img_name + '.png') - elif task=="enhancement": - dir_of_label_file = os.path.join(dir_seg, im) - - if not patches: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(cv2.imread(dir_img + '/' + im), input_height, input_width)) - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - indexer += 1 - - if augmentation: - if flip_aug: - for f_i in flip_index: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - resize_image(cv2.flip(cv2.imread(dir_img+'/'+im),f_i),input_height,input_width) ) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.flip(cv2.imread(dir_of_label_file), f_i), input_height, input_width)) - indexer += 1 - - if blur_aug: - for blur_i in blur_k: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - (resize_image(bluring(cv2.imread(dir_img + '/' + im), blur_i), input_height, input_width))) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - indexer += 1 - if brightening: - for factor in brightness: - try: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - (resize_image(do_brightening(dir_img + '/' +im, factor), input_height, input_width))) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - indexer += 1 - except: - pass - - if binarization: - - if dir_img_bin: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - resize_image(img_bin_corr, input_height, input_width)) - else: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - resize_image(otsu_copy(cv2.imread(dir_img + '/' + im)), input_height, input_width)) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - indexer += 1 - - if degrading: - for degrade_scale_ind in degrade_scales: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - (resize_image(do_degrading(cv2.imread(dir_img + '/' + im), degrade_scale_ind), input_height, input_width))) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - indexer += 1 - - if rotation_not_90: - for thetha_i in thetha: - img_max_rotated, label_max_rotated = rotation_not_90_func(cv2.imread(dir_img + '/'+im), - cv2.imread(dir_of_label_file), thetha_i) - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(img_max_rotated, input_height, input_width)) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', resize_image(label_max_rotated, input_height, input_width)) - indexer += 1 - - if channels_shuffling: - for shuffle_index in shuffle_indexes: - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', - (resize_image(return_shuffled_channels(cv2.imread(dir_img + '/' + im), shuffle_index), input_height, input_width))) - - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - indexer += 1 - - if scaling: - for sc_ind in scales: - img_scaled, label_scaled = scale_image_for_no_patch(cv2.imread(dir_img + '/'+im), - cv2.imread(dir_of_label_file), sc_ind) - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(img_scaled, input_height, input_width)) - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', resize_image(label_scaled, input_height, input_width)) - indexer += 1 - if shifting: - shift_types = ['xpos', 'xmin', 'ypos', 'ymin', 'xypos', 'xymin'] - for st_ind in shift_types: - img_shifted, label_shifted = shift_image_and_label(cv2.imread(dir_img + '/'+im), - cv2.imread(dir_of_label_file), st_ind) - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(img_shifted, input_height, input_width)) - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', resize_image(label_shifted, input_height, input_width)) - indexer += 1 - - - if adding_rgb_background: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - for i_n in range(number_of_backgrounds_per_image): - background_image_chosen_name = random.choice(list_all_possible_background_images) - img_rgb_background_chosen = cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) - img_with_overlayed_background = return_binary_image_with_given_rgb_background(img_bin_corr, img_rgb_background_chosen) - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(img_with_overlayed_background, input_height, input_width)) - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - - indexer += 1 - - if adding_rgb_foreground: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - for i_n in range(number_of_backgrounds_per_image): - background_image_chosen_name = random.choice(list_all_possible_background_images) - foreground_rgb_chosen_name = random.choice(list_all_possible_foreground_rgbs) - - img_rgb_background_chosen = cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) - foreground_rgb_chosen = np.load(dir_rgb_foregrounds + '/' + foreground_rgb_chosen_name) - - img_with_overlayed_background = return_binary_image_with_given_rgb_background_and_given_foreground_rgb(img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(img_with_overlayed_background, input_height, input_width)) - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - - indexer += 1 - - if add_red_textlines: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - img_red_context = return_image_with_red_elements(cv2.imread(dir_img + '/'+im), img_bin_corr) - - cv2.imwrite(dir_flow_train_imgs + '/img_' + str(indexer) + '.png', resize_image(img_red_context, input_height, input_width)) - cv2.imwrite(dir_flow_train_labels + '/img_' + str(indexer) + '.png', - resize_image(cv2.imread(dir_of_label_file), input_height, input_width)) - - indexer += 1 - - - - - if patches: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - cv2.imread(dir_img + '/' + im), cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - if augmentation: - if rotation: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - rotation_90(cv2.imread(dir_img + '/' + im)), - rotation_90(cv2.imread(dir_of_label_file)), - input_height, input_width, indexer=indexer) - - if rotation_not_90: - for thetha_i in thetha: - img_max_rotated, label_max_rotated = rotation_not_90_func(cv2.imread(dir_img + '/'+im), - cv2.imread(dir_of_label_file), thetha_i) - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - img_max_rotated, - label_max_rotated, - input_height, input_width, indexer=indexer) - - if channels_shuffling: - for shuffle_index in shuffle_indexes: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - return_shuffled_channels(cv2.imread(dir_img + '/' + im), shuffle_index), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - if adding_rgb_background: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - for i_n in range(number_of_backgrounds_per_image): - background_image_chosen_name = random.choice(list_all_possible_background_images) - img_rgb_background_chosen = cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) - img_with_overlayed_background = return_binary_image_with_given_rgb_background(img_bin_corr, img_rgb_background_chosen) - - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - img_with_overlayed_background, - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - - if adding_rgb_foreground: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - for i_n in range(number_of_backgrounds_per_image): - background_image_chosen_name = random.choice(list_all_possible_background_images) - foreground_rgb_chosen_name = random.choice(list_all_possible_foreground_rgbs) - - img_rgb_background_chosen = cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) - foreground_rgb_chosen = np.load(dir_rgb_foregrounds + '/' + foreground_rgb_chosen_name) - - img_with_overlayed_background = return_binary_image_with_given_rgb_background_and_given_foreground_rgb(img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) - - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - img_with_overlayed_background, - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - - if add_red_textlines: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - img_red_context = return_image_with_red_elements(cv2.imread(dir_img + '/'+im), img_bin_corr) - - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - img_red_context, - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - if flip_aug: - for f_i in flip_index: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - cv2.flip(cv2.imread(dir_img + '/' + im), f_i), - cv2.flip(cv2.imread(dir_of_label_file), f_i), - input_height, input_width, indexer=indexer) - if blur_aug: - for blur_i in blur_k: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - bluring(cv2.imread(dir_img + '/' + im), blur_i), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - if padding_black: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - do_padding_black(cv2.imread(dir_img + '/' + im)), - do_padding_label(cv2.imread(dir_of_label_file)), - input_height, input_width, indexer=indexer) - - if padding_white: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - do_padding_white(cv2.imread(dir_img + '/'+im)), - do_padding_label(cv2.imread(dir_of_label_file)), - input_height, input_width, indexer=indexer) - - if brightening: - for factor in brightness: - try: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - do_brightening(dir_img + '/' +im, factor), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - except: - pass - if scaling: - for sc_ind in scales: - indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, - cv2.imread(dir_img + '/' + im) , - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer, scaler=sc_ind) - - if degrading: - for degrade_scale_ind in degrade_scales: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - do_degrading(cv2.imread(dir_img + '/' + im), degrade_scale_ind), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - if binarization: - if dir_img_bin: - img_bin_corr = cv2.imread(dir_img_bin + '/' + img_name+'.png') - - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - img_bin_corr, - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) - - else: - indexer = get_patches(dir_flow_train_imgs, dir_flow_train_labels, - otsu_copy(cv2.imread(dir_img + '/' + im)), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer) + def scale_image(img): + return scale_padd_image_for_ocr(img, input_height, input_width).astype(np.float32) / 255. + #lab = vectorize_label(lab, char_to_num, padding_token, max_len) + # now padded at Dataset.padded_batch + lab = char_to_num(tf.strings.unicode_split(lab, input_encoding="UTF-8")) + yield scale_image(img), lab + #to_yield = {"image": ret_x, "label": ret_y} - if scaling_brightness: - for sc_ind in scales: - for factor in brightness: - try: - indexer = get_patches_num_scale_new(dir_flow_train_imgs, - dir_flow_train_labels, - do_brightening(dir_img + '/' + im, factor) - ,cv2.imread(dir_of_label_file) - ,input_height, input_width, indexer=indexer, scaler=sc_ind) - except: - pass - - if scaling_bluring: - for sc_ind in scales: - for blur_i in blur_k: - indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, - bluring(cv2.imread(dir_img + '/' + im), blur_i), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer, scaler=sc_ind) + if dir_img_bin: + img_bin_corr = cv2.imread(os.path.join(dir_img_bin, img_name + '.png')) + else: + img_bin_corr = None - if scaling_binarization: - for sc_ind in scales: - indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, - otsu_copy(cv2.imread(dir_img + '/' + im)), - cv2.imread(dir_of_label_file), - input_height, input_width, indexer=indexer, scaler=sc_ind) - - if scaling_flip: - for sc_ind in scales: - for f_i in flip_index: - indexer = get_patches_num_scale_new(dir_flow_train_imgs, dir_flow_train_labels, - - cv2.flip( cv2.imread(dir_img + '/' + im), f_i), - cv2.flip(cv2.imread(dir_of_label_file), f_i), - input_height, input_width, indexer=indexer, scaler=sc_ind) - - - -def data_gen_ocr( - padding_token, - n_batch, - input_height, - input_width, - max_len, - dir_train, - ls_files_images, - augmentation, - color_padding_rotation, - rotation_not_90, - blur_aug, - degrading, - bin_deg, - brightening, - padding_white, - adding_rgb_foreground, - adding_rgb_background, - binarization, - image_inversion, - channels_shuffling, - add_red_textlines, - white_noise_strap, - textline_skewing, - textline_skewing_bin, - textline_left_in_depth, - textline_left_in_depth_bin, - textline_right_in_depth, - textline_right_in_depth_bin, - textline_up_in_depth, - textline_up_in_depth_bin, - textline_down_in_depth, - textline_down_in_depth_bin, - pepper_bin_aug, - pepper_aug, - degrade_scales, - number_of_backgrounds_per_image, - thetha, - thetha_padd, - brightness, - padd_colors, - shuffle_indexes, - pepper_indexes, - skewing_amplitudes, - blur_k, - char_to_num, - list_all_possible_background_images, - list_all_possible_foreground_rgbs, - dir_rgb_backgrounds, - dir_rgb_foregrounds, - white_padds, - dir_img_bin=None, -): - - random.shuffle(ls_files_images) - - ret_x= np.zeros((n_batch, input_height, input_width, 3)).astype(np.float32) - ret_y= np.zeros((n_batch, max_len)).astype(np.int16)+padding_token - batchcount = 0 - - def increment_batchcount(img_out, batchcount, ret_x, ret_y): - to_yield = None - img_out = scale_padd_image_for_ocr(img, input_height, input_width) - ret_x[batchcount, :,:,:] = img_out[:,:,:] - ret_y[batchcount, :] = vectorize_label(txt_inp, char_to_num, padding_token, max_len) - batchcount += 1 - if batchcount>=n_batch: - ret_x = ret_x/255. - to_yield = {"image": ret_x, "label": ret_y} - ret_x= np.zeros((n_batch, input_height, input_width, 3)).astype(np.float32) - ret_y= np.zeros((n_batch, max_len)).astype(np.int16)+padding_token - batchcount = 0 - return img_out, batchcount, ret_x, ret_y, to_yield - - # TODO: Why while True + yield, why not return a list? - while True: - for i in ls_files_images: - print(i, 'i') - f_name = Path(i).stem#.split('.')[0] - - txt_inp = open(os.path.join(dir_train, "labels/"+f_name+'.txt'),'r').read().split('\n')[0] - - img = cv2.imread(os.path.join(dir_train, "images/"+i) ) - if dir_img_bin: - img_bin_corr = cv2.imread(os.path.join(dir_img_bin, f_name+'.png') ) - else: - img_bin_corr = None - - - if augmentation: - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if color_padding_rotation: - for thetha_ind in thetha_padd: - for padd_col in padd_colors: - img_out = rotation_not_90_func_single_image(do_padding_for_ocr(img, 1.2, padd_col), thetha_ind) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if rotation_not_90: - for thetha_ind in thetha: - img_out = rotation_not_90_func_single_image(img, thetha_ind) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if blur_aug: - for blur_type in blur_k: - img_out = bluring(img, blur_type) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if degrading: - for deg_scale_ind in degrade_scales: - try: - img_out = do_degrading(img, deg_scale_ind) - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if bin_deg: - for deg_scale_ind in degrade_scales: - try: - img_out = do_degrading(img_bin_corr, deg_scale_ind) - # TODO: qualify except - except: - img_out = np.copy(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if brightening: - for bright_scale_ind in brightness: - try: - # FIXME: dir_img is not defined in this scope, will always fail - img_out = do_brightening(dir_img, bright_scale_ind) - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if padding_white: - for padding_size in white_padds: - for padd_col in padd_colors: - img_out = do_padding_for_ocr(img, padding_size, padd_col) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if adding_rgb_foreground: - for i_n in range(number_of_backgrounds_per_image): - background_image_chosen_name = random.choice(list_all_possible_background_images) - foreground_rgb_chosen_name = random.choice(list_all_possible_foreground_rgbs) - - img_rgb_background_chosen = cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) - foreground_rgb_chosen = np.load(dir_rgb_foregrounds + '/' + foreground_rgb_chosen_name) - - img_out = return_binary_image_with_given_rgb_background_and_given_foreground_rgb(img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) - - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - - if adding_rgb_background: - for i_n in range(number_of_backgrounds_per_image): - background_image_chosen_name = random.choice(list_all_possible_background_images) - img_rgb_background_chosen = cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) - img_out = return_binary_image_with_given_rgb_background(img_bin_corr, img_rgb_background_chosen) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if binarization: - img_out = scale_padd_image_for_ocr(img_bin_corr, input_height, input_width) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if image_inversion: - img_out = invert_image(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if channels_shuffling: - for shuffle_index in shuffle_indexes: - img_out = return_shuffled_channels(img, shuffle_index) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if add_red_textlines: - img_out = return_image_with_red_elements(img, img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if white_noise_strap: - img_out = return_image_with_strapped_white_noises(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_skewing: - for des_scale_ind in skewing_amplitudes: - try: - img_out = do_deskewing(img, des_scale_ind) - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_skewing_bin: - for des_scale_ind in skewing_amplitudes: - try: - img_out = do_deskewing(img_bin_corr, des_scale_ind) - # TODO: qualify except - except: - img_out = np.copy(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_left_in_depth: - try: - img_out = do_direction_in_depth(img, 'left') - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_left_in_depth_bin: - try: - img_out = do_direction_in_depth(img_bin_corr, 'left') - # TODO: qualify except - except: - img_out = np.copy(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_right_in_depth: - try: - img_out = do_direction_in_depth(img_bin_corr, 'right') - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - - if textline_right_in_depth_bin: - try: - img_out = do_direction_in_depth(img_bin_corr, 'right') - # TODO: qualify except - except: - img_out = np.copy(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_up_in_depth: - try: - img_out = do_direction_in_depth(img, 'up') - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_up_in_depth_bin: - try: - img_out = do_direction_in_depth(img_bin_corr, 'up') - # TODO: qualify except - except: - img_out = np.copy(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_down_in_depth: - try: - img_out = do_direction_in_depth(img, 'down') - # TODO: qualify except - except: - img_out = np.copy(img) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if textline_down_in_depth_bin: - try: - img_out = do_direction_in_depth(img_bin_corr, 'down') - # TODO: qualify except - except: - img_out = np.copy(img_bin_corr) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if pepper_bin_aug: - for pepper_ind in pepper_indexes: - img_out = add_salt_and_pepper_noise(img_bin_corr, pepper_ind, pepper_ind) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - if pepper_aug: - for pepper_ind in pepper_indexes: - img_out = add_salt_and_pepper_noise(img, pepper_ind, pepper_ind) - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - else: - img_out, batchcount, ret_x, ret_y, to_yield = increment_batchcount(img_out, batchcount, ret_x, ret_y) - if to_yield: yield to_yield - - -# TODO: what is aug_multip and why calculate it in this way -def return_multiplier_based_on_augmnentations( - augmentation, - color_padding_rotation, - rotation_not_90, - blur_aug, - degrading, - bin_deg, - brightening, - padding_white, - adding_rgb_foreground, - adding_rgb_background, - binarization, - image_inversion, - channels_shuffling, - add_red_textlines, - white_noise_strap, - textline_skewing, - textline_skewing_bin, - textline_left_in_depth, - textline_left_in_depth_bin, - textline_right_in_depth, - textline_right_in_depth_bin, - textline_up_in_depth, - textline_up_in_depth_bin, - textline_down_in_depth, - textline_down_in_depth_bin, - pepper_bin_aug, - pepper_aug, - degrade_scales, - number_of_backgrounds_per_image, - thetha, - thetha_padd, - brightness, - padd_colors, - shuffle_indexes, - pepper_indexes, - skewing_amplitudes, - blur_k, - white_padds, -): - aug_multip = 1 if not augmentation: - return 1 + return - if binarization: - aug_multip += 1 - if image_inversion: - aug_multip += 1 - if add_red_textlines: - aug_multip += 1 - if white_noise_strap: - aug_multip += 1 - if textline_right_in_depth: - aug_multip += 1 - if textline_left_in_depth: - aug_multip += 1 - if textline_up_in_depth: - aug_multip += 1 - if textline_down_in_depth: - aug_multip += 1 - if textline_right_in_depth_bin: - aug_multip += 1 - if textline_left_in_depth_bin: - aug_multip += 1 - if textline_up_in_depth_bin: - aug_multip += 1 - if textline_down_in_depth_bin: - aug_multip += 1 - if adding_rgb_foreground: - aug_multip += number_of_backgrounds_per_image - if adding_rgb_background: - aug_multip += number_of_backgrounds_per_image - if bin_deg: - aug_multip += len(degrade_scales) - if degrading: - aug_multip += len(degrade_scales) - if rotation_not_90: - aug_multip += len(thetha) - if textline_skewing: - aug_multip += len(skewing_amplitudes) - if textline_skewing_bin: - aug_multip += len(skewing_amplitudes) if color_padding_rotation: - aug_multip += len(thetha_padd)*len(padd_colors) - if channels_shuffling: - aug_multip += len(shuffle_indexes) + for thetha_ind in thetha_padd: + for padd_col in padd_colors: + img_pad = do_padding_for_ocr(img, 1.2, padd_col) + img_rot = rotation_not_90_func_single_image(img_pad, thetha_ind) + yield scale_image(img_rot), lab + if rotation_not_90: + for thetha_ind in thetha: + img_rot = rotation_not_90_func_single_image(img, thetha_ind) + yield scale_image(img_rot), lab if blur_aug: - aug_multip += len(blur_k) + for blur_type in blur_k: + img_blur = bluring(img, blur_type) + yield scale_image(img_blur), lab + if degrading: + for deg_scale_ind in degrade_scales: + img_deg = do_degrading(img, deg_scale_ind) + yield scale_image(img_deg), lab + if bin_deg: + for deg_scale_ind in degrade_scales: + img_deg = do_degrading(img_bin_corr, deg_scale_ind) + yield scale_image(img_deg), lab if brightening: - aug_multip += len(brightness) + for bright_scale_ind in brightness: + img_bright = do_brightening(img, bright_scale_ind) + yield scale_image(img_bright), lab if padding_white: - aug_multip += len(white_padds)*len(padd_colors) + for padding_size in white_padds: + for padd_col in padd_colors: + img_pad = do_padding_for_ocr(img, padding_size, padd_col) + yield scale_image(img_pad), lab + if adding_rgb_foreground: + for i_n in range(number_of_backgrounds_per_image): + background_image_chosen_name = random.choice(list_all_possible_background_images) + foreground_rgb_chosen_name = random.choice(list_all_possible_foreground_rgbs) + + img_rgb_background_chosen = \ + cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) + foreground_rgb_chosen = \ + np.load(dir_rgb_foregrounds + '/' + foreground_rgb_chosen_name) + + img_fg = \ + return_binary_image_with_given_rgb_background_and_given_foreground_rgb( + img_bin_corr, img_rgb_background_chosen, foreground_rgb_chosen) + yield scale_image(img_fg), lab + if adding_rgb_background: + for i_n in range(number_of_backgrounds_per_image): + background_image_chosen_name = random.choice(list_all_possible_background_images) + img_rgb_background_chosen = \ + cv2.imread(dir_rgb_backgrounds + '/' + background_image_chosen_name) + img_bg = \ + return_binary_image_with_given_rgb_background(img_bin_corr, img_rgb_background_chosen) + yield scale_image(img_bg), lab + if binarization: + yield scale_image(img_bin_corr), lab + if image_inversion: + img_inv = invert_image(img_bin_corr) + yield scale_image(img_inv), lab + if channels_shuffling: + for shuffle_index in shuffle_indexes: + img_shuf = return_shuffled_channels(img, shuffle_index) + yield scale_image(img_shuf), lab + if add_red_textlines: + img_red = return_image_with_red_elements(img, img_bin_corr) + yield scale_image(img_red), lab + if white_noise_strap: + img_noisy = return_image_with_strapped_white_noises(img) + yield scale_image(img_noisy), lab + if textline_skewing: + for des_scale_ind in skewing_amplitudes: + img_rot = do_deskewing(img, des_scale_ind) + yield scale_image(img_rot), lab + if textline_skewing_bin: + for des_scale_ind in skewing_amplitudes: + img_rot = do_deskewing(img_bin_corr, des_scale_ind) + yield scale_image(img_rot), lab + if textline_left_in_depth: + img_warp = do_direction_in_depth(img, 'left') + yield scale_image(img_warp), lab + if textline_left_in_depth_bin: + img_warp = do_direction_in_depth(img_bin_corr, 'left') + yield scale_image(img_warp), lab + if textline_right_in_depth: + img_warp = do_direction_in_depth(img, 'right') + yield scale_image(img_warp), lab + if textline_right_in_depth_bin: + img_warp = do_direction_in_depth(img_bin_corr, 'right') + yield scale_image(img_warp), lab + if textline_up_in_depth: + img_warp = do_direction_in_depth(img, 'up') + yield scale_image(img_warp), lab + if textline_up_in_depth_bin: + img_warp = do_direction_in_depth(img_bin_corr, 'up') + yield scale_image(img_warp), lab + if textline_down_in_depth: + img_warp = do_direction_in_depth(img, 'down') + yield scale_image(img_warp), lab + if textline_down_in_depth_bin: + img_warp = do_direction_in_depth(img_bin_corr, 'down') + yield scale_image(img_warp), lab if pepper_aug: - aug_multip += len(pepper_indexes) + for pepper_ind in pepper_indexes: + img_noisy = add_salt_and_pepper_noise(img, pepper_ind, pepper_ind) + yield scale_image(img_noisy), lab if pepper_bin_aug: - aug_multip += len(pepper_indexes) - - return aug_multip + for pepper_ind in pepper_indexes: + img_noisy = add_salt_and_pepper_noise(img_bin_corr, pepper_ind, pepper_ind) + yield scale_image(img_noisy), lab diff --git a/src/eynollah/training/weights_ensembling.py b/src/eynollah/training/weights_ensembling.py index 6dce7fd..e3ede24 100644 --- a/src/eynollah/training/weights_ensembling.py +++ b/src/eynollah/training/weights_ensembling.py @@ -1,136 +1,66 @@ -import sys -from glob import glob -from os import environ, devnull -from os.path import join -from warnings import catch_warnings, simplefilter import os +from warnings import catch_warnings, simplefilter +import click import numpy as np -from PIL import Image -import cv2 -environ['TF_CPP_MIN_LOG_LEVEL'] = '3' -stderr = sys.stderr -sys.stderr = open(devnull, 'w') + +os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' + +from ocrd_utils import tf_disable_interactive_logs +tf_disable_interactive_logs() import tensorflow as tf from tensorflow.keras.models import load_model -from tensorflow.python.keras import backend as tensorflow_backend -sys.stderr = stderr -from tensorflow.keras import layers -import tensorflow.keras.losses -from tensorflow.keras.layers import * -import click -import logging - -class Patches(layers.Layer): - def __init__(self, patch_size_x, patch_size_y): - super(Patches, self).__init__() - self.patch_size_x = patch_size_x - self.patch_size_y = patch_size_y - - def call(self, images): - #print(tf.shape(images)[1],'images') - #print(self.patch_size,'self.patch_size') - batch_size = tf.shape(images)[0] - patches = tf.image.extract_patches( - images=images, - sizes=[1, self.patch_size_y, self.patch_size_x, 1], - strides=[1, self.patch_size_y, self.patch_size_x, 1], - rates=[1, 1, 1, 1], - padding="VALID", - ) - #patch_dims = patches.shape[-1] - patch_dims = tf.shape(patches)[-1] - patches = tf.reshape(patches, [batch_size, -1, patch_dims]) - return patches - def get_config(self): - - config = super().get_config().copy() - config.update({ - 'patch_size_x': self.patch_size_x, - 'patch_size_y': self.patch_size_y, - }) - return config - - - -class PatchEncoder(layers.Layer): - def __init__(self, **kwargs): - super(PatchEncoder, self).__init__() - self.num_patches = num_patches - self.projection = layers.Dense(units=projection_dim) - self.position_embedding = layers.Embedding( - input_dim=num_patches, output_dim=projection_dim - ) - - def call(self, patch): - positions = tf.range(start=0, limit=self.num_patches, delta=1) - encoded = self.projection(patch) + self.position_embedding(positions) - return encoded - def get_config(self): - - config = super().get_config().copy() - config.update({ - 'num_patches': self.num_patches, - 'projection': self.projection, - 'position_embedding': self.position_embedding, - }) - return config +from ..patch_encoder import ( + PatchEncoder, + Patches, +) - -def start_new_session(): - ###config = tf.compat.v1.ConfigProto() - ###config.gpu_options.allow_growth = True +def run_ensembling(model_dirs, out_dir): + all_weights = [] - ###self.session = tf.compat.v1.Session(config=config) # tf.InteractiveSession() - ###tensorflow_backend.set_session(self.session) - - config = tf.compat.v1.ConfigProto() - config.gpu_options.allow_growth = True - - session = tf.compat.v1.Session(config=config) # tf.InteractiveSession() - tensorflow_backend.set_session(session) - return session - -def run_ensembling(dir_models, out): - ls_models = os.listdir(dir_models) - - - weights=[] - - for model_name in ls_models: - model = load_model(os.path.join(dir_models,model_name) , compile=False, custom_objects={'PatchEncoder':PatchEncoder, 'Patches': Patches}) - weights.append(model.get_weights()) + for model_dir in model_dirs: + assert os.path.isdir(model_dir), model_dir + model = load_model(model_dir, compile=False, + custom_objects=dict(PatchEncoder=PatchEncoder, + Patches=Patches)) + all_weights.append(model.get_weights()) - new_weights = list() + new_weights = [] + for layer_weights in zip(*all_weights): + layer_weights = np.array([np.array(weights).mean(axis=0) + for weights in zip(*layer_weights)]) + new_weights.append(layer_weights) - for weights_list_tuple in zip(*weights): - new_weights.append( - [np.array(weights_).mean(axis=0)\ - for weights_ in zip(*weights_list_tuple)]) - - - - new_weights = [np.array(x) for x in new_weights] - + #model = tf.keras.models.clone_model(model) model.set_weights(new_weights) - model.save(out) - os.system('cp '+os.path.join(os.path.join(dir_models,model_name) , "config.json ")+out) + + model.save(out_dir) + os.system('cp ' + os.path.join(model_dirs[0], "config.json ") + out_dir + "/") @click.command() @click.option( - "--dir_models", - "-dm", - help="directory of models", + "--in", + "-i", + help="input directory of checkpoint models to be read", + multiple=True, + required=True, type=click.Path(exists=True, file_okay=False), ) @click.option( "--out", "-o", help="output directory where ensembled model will be written.", + required=True, type=click.Path(exists=False, file_okay=False), ) +def ensemble_cli(in_, out): + """ + mix multiple model weights + + Load a sequence of models and mix them into a single ensemble model + by averaging their weights. Write the resulting model. + """ + run_ensembling(in_, out) -def main(dir_models, out): - run_ensembling(dir_models, out) - diff --git a/src/eynollah/utils/__init__.py b/src/eynollah/utils/__init__.py index 29359eb..38a50be 100644 --- a/src/eynollah/utils/__init__.py +++ b/src/eynollah/utils/__init__.py @@ -1,10 +1,11 @@ -from typing import Tuple +from typing import Iterable, List, Tuple from logging import getLogger import time import math try: import matplotlib.pyplot as plt + import matplotlib.patches as patches except ImportError: plt = None import numpy as np @@ -19,6 +20,8 @@ from .contour import (contours_in_same_horizon, find_new_features_of_contours, return_contours_of_image, return_parent_contours) + + def pairwise(iterable): # pairwise('ABCDEFG') → AB BC CD DE EF FG @@ -29,286 +32,132 @@ def pairwise(iterable): yield a, b a = b -def return_x_start_end_mothers_childs_and_type_of_reading_order( - x_min_hor_some, x_max_hor_some, cy_hor_some, peak_points, cy_hor_diff): +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): + """ + Analyse which separators overlap multiple column candidates, + and how they overlap each other. - x_start=[] - x_end=[] - kind=[]#if covers 2 and more than 2 columns set it to 1 otherwise 0 - len_sep=[] - y_sep=[] - y_diff=[] - new_main_sep_y=[] + Ignore separators not spanning multiple columns. - indexer=0 + For the separators to be returned, try to remove or unify them when there + is no region between them (vertically) and their neighbours. + + Arguments: + * the text mask (with all separators suppressed) + * the x column coordinates + * the y start coordinate to consider in total + * the y end coordinate to consider in total + * the x start coordinate of the horizontal separators + * the x end coordinate of the horizontal separators + * the y start coordinate of the horizontal separators + * the y center coordinate of the horizontal separators + * the y end coordinate of the horizontal separators + + Returns: + a tuple of: + * the x start column index of the resulting multi-span separators + * the x end column index of the resulting multi-span separators + * the y start coordinate of the resulting multi-span separators + * the y center coordinate of the resulting multi-span separators + * the y end coordinate of the resulting multi-span separators + """ + + x_start = [0] + x_end = [len(peak_points) - 1] + y_min = [top] + y_mid = [top] + y_max = [top + 2] + indexer = 1 for i in range(len(x_min_hor_some)): - starting=x_min_hor_some[i]-peak_points - starting=starting[starting>=0] - min_start=np.argmin(starting) - ending=peak_points-x_max_hor_some[i] - len_ending_neg=len(ending[ending<=0]) - - ending=ending[ending>0] - max_end=np.argmin(ending)+len_ending_neg + #print(indexer, "%d:%d" % (x_min_hor_some[i], x_max_hor_some[i]), cy_hor_some[i]) + starting = x_min_hor_some[i] - peak_points + min_start = np.flatnonzero(starting >= 0)[-1] # last left-of + ending = x_max_hor_some[i] - peak_points + max_end = np.flatnonzero(ending <= 0)[0] # first right-of + #print(indexer, "%d:%d" % (min_start, max_end)) if (max_end-min_start)>=2: - if (max_end-min_start)==(len(peak_points)-1): - new_main_sep_y.append(indexer) - + # column range of separator spans more than one column candidate #print((max_end-min_start),len(peak_points),'(max_end-min_start)') - y_sep.append(cy_hor_some[i]) - y_diff.append(cy_hor_diff[i]) + y_min.append(y_min_hor_some[i]) + y_mid.append(cy_hor_some[i]) + y_max.append(y_max_hor_some[i]) x_end.append(max_end) - - x_start.append( min_start) - - len_sep.append(max_end-min_start) - if max_end==min_start+1: - kind.append(0) - else: - kind.append(1) - + x_start.append(min_start) indexer+=1 - - x_start_returned = np.array(x_start, dtype=int) - x_end_returned = np.array(x_end, dtype=int) - y_sep_returned = np.array(y_sep, dtype=int) - y_diff_returned = np.array(y_diff, dtype=int) - - all_args_uniq = contours_in_same_horizon(y_sep_returned) - args_to_be_unified=[] - y_unified=[] - y_diff_unified=[] - x_s_unified=[] - x_e_unified=[] - if len(all_args_uniq)>0: - #print('burda') - if type(all_args_uniq[0]) is list: - for dd in range(len(all_args_uniq)): - if len(all_args_uniq[dd])==2: - x_s_same_hor=np.array(x_start_returned)[all_args_uniq[dd]] - x_e_same_hor=np.array(x_end_returned)[all_args_uniq[dd]] - y_sep_same_hor=np.array(y_sep_returned)[all_args_uniq[dd]] - y_diff_same_hor=np.array(y_diff_returned)[all_args_uniq[dd]] - #print('burda2') - if (x_s_same_hor[0]==x_e_same_hor[1]-1 or - x_s_same_hor[1]==x_e_same_hor[0]-1 and - x_s_same_hor[0]!=x_s_same_hor[1] and - x_e_same_hor[0]!=x_e_same_hor[1]): - #print('burda3') - for arg_in in all_args_uniq[dd]: - #print(arg_in,'arg_in') - args_to_be_unified.append(arg_in) - y_selected=np.min(y_sep_same_hor) - y_diff_selected=np.max(y_diff_same_hor) - x_s_selected=np.min(x_s_same_hor) - x_e_selected=np.max(x_e_same_hor) - - x_s_unified.append(x_s_selected) - x_e_unified.append(x_e_selected) - y_unified.append(y_selected) - y_diff_unified.append(y_diff_selected) - #print(x_s_same_hor,'x_s_same_hor') - #print(x_e_same_hor[:]-1,'x_e_same_hor') - #print('#############################') - #print(x_s_unified,'y_selected') - #print(x_e_unified,'x_s_selected') - #print(y_unified,'x_e_same_hor') - - args_lines_not_unified=list( set(range(len(y_sep_returned)))-set(args_to_be_unified) ) - #print(args_lines_not_unified,'args_lines_not_unified') - - x_start_returned_not_unified=list( np.array(x_start_returned)[args_lines_not_unified] ) - x_end_returned_not_unified=list( np.array(x_end_returned)[args_lines_not_unified] ) - y_sep_returned_not_unified=list (np.array(y_sep_returned)[args_lines_not_unified] ) - y_diff_returned_not_unified=list (np.array(y_diff_returned)[args_lines_not_unified] ) - - for dv in range(len(y_unified)): - y_sep_returned_not_unified.append(y_unified[dv]) - y_diff_returned_not_unified.append(y_diff_unified[dv]) - x_start_returned_not_unified.append(x_s_unified[dv]) - x_end_returned_not_unified.append(x_e_unified[dv]) - - #print(y_sep_returned,'y_sep_returned') - #print(x_start_returned,'x_start_returned') - #print(x_end_returned,'x_end_returned') - - x_start_returned = np.array(x_start_returned_not_unified, dtype=int) - x_end_returned = np.array(x_end_returned_not_unified, dtype=int) - y_sep_returned = np.array(y_sep_returned_not_unified, dtype=int) - y_diff_returned = np.array(y_diff_returned_not_unified, dtype=int) - - #print(y_sep_returned,'y_sep_returned2') - #print(x_start_returned,'x_start_returned2') - #print(x_end_returned,'x_end_returned2') - #print(new_main_sep_y,'new_main_sep_y') - #print(x_start,'x_start') #print(x_end,'x_end') - if len(new_main_sep_y)>0: - min_ys=np.min(y_sep) - max_ys=np.max(y_sep) - - y_mains=[] - y_mains.append(min_ys) - y_mains_sep_ohne_grenzen=[] - - for ii in range(len(new_main_sep_y)): - y_mains.append(y_sep[new_main_sep_y[ii]]) - y_mains_sep_ohne_grenzen.append(y_sep[new_main_sep_y[ii]]) - - y_mains.append(max_ys) - - y_mains_sorted=np.sort(y_mains) - diff=np.diff(y_mains_sorted) - argm=np.argmax(diff) - - y_min_new=y_mains_sorted[argm] - y_max_new=y_mains_sorted[argm+1] - - #print(y_min_new,'y_min_new') - #print(y_max_new,'y_max_new') - #print(y_sep[new_main_sep_y[0]],y_sep,'yseps') - x_start=np.array(x_start) - x_end=np.array(x_end) - kind=np.array(kind) - y_sep=np.array(y_sep) - if (y_min_new in y_mains_sep_ohne_grenzen and - y_max_new in y_mains_sep_ohne_grenzen): - x_start=x_start[(y_sep>y_min_new) & (y_sepy_min_new) & (y_sepy_min_new) & (y_sepy_min_new) & (y_sepy_min_new) & (y_sep<=y_max_new)] - #print('burda1') - x_end=x_end[(y_sep>y_min_new) & (y_sep<=y_max_new)] - #print('burda2') - kind=kind[(y_sep>y_min_new) & (y_sep<=y_max_new)] - y_sep=y_sep[(y_sep>y_min_new) & (y_sep<=y_max_new)] - elif (y_min_new not in y_mains_sep_ohne_grenzen and - y_max_new in y_mains_sep_ohne_grenzen): - x_start=x_start[(y_sep>=y_min_new) & (y_sep=y_min_new) & (y_sep=y_min_new) & (y_sep=y_min_new) & (y_sep=y_min_new) & (y_sep<=y_max_new)] - x_end=x_end[(y_sep>=y_min_new) & (y_sep<=y_max_new)] - kind=kind[(y_sep>=y_min_new) & (y_sep<=y_max_new)] - y_sep=y_sep[(y_sep>=y_min_new) & (y_sep<=y_max_new)] + x_start = np.array(x_start, dtype=int) + x_end = np.array(x_end, dtype=int) + y_min = np.array(y_min, dtype=int) + y_mid = np.array(y_mid, dtype=int) + y_max = np.array(y_max, dtype=int) + #print(y_mid,'y_mid') #print(x_start,'x_start') #print(x_end,'x_end') - #print(len_sep) - deleted=[] - for i in range(len(x_start)-1): - nodes_i=set(range(x_start[i],x_end[i]+1)) - for j in range(i+1,len(x_start)): - if nodes_i==set(range(x_start[j],x_end[j]+1)): - deleted.append(j) - #print(np.unique(deleted)) + # remove redundant separators (with nothing in between) + args_emptysep = set() + args_ysorted = np.argsort(y_mid) + for i in range(len(y_mid)): + # find nearest neighbours above with nothing in between + prev = (~np.eye(len(y_mid), dtype=bool)[i] & + (y_mid[i] >= y_mid) & + # complete subsumption: + # (x_start[i] >= x_start) & + # (x_end[i] <= x_end) + # partial overlap + (x_start[i] < x_end) & + (x_end[i] > x_start) + ) + prev[list(args_emptysep)] = False # but no pair we already saw + if not prev.any(): + continue + prev = np.flatnonzero(prev[args_ysorted]) + j = args_ysorted[prev[-1]] + if not np.any(regions_without_separators[y_max[j]: y_min[i], + peak_points[min(x_start[i], x_start[j])]: + peak_points[max(x_end[i], x_end[j])]]): + args_emptysep.add(i) + if x_start[j] > x_start[i]: + # print(j, "now starts at", x_start[i]) + x_start[j] = x_start[i] + if x_end[j] < x_end[i]: + x_end[j] = x_end[i] + # print(j, "now ends at", x_end[i]) + # print(j, i, "%d:%d" % (y_mid[j], y_mid[i]), "%d:%d" % (x_start[i], x_end[i]), "empty prev sep") + continue + # find nearest neighbours below with nothing in between + nExt = (~np.eye(len(y_mid), dtype=bool)[i] & + (y_mid[i] <= y_mid) & + (x_start[i] >= x_start) & + (x_end[i] <= x_end)) + nExt[list(args_emptysep)] = False # but no pair we already saw + if not nExt.any(): + continue + nExt = np.flatnonzero(nExt[args_ysorted]) + j = args_ysorted[nExt[0]] + if not np.any(regions_without_separators[y_max[i]: y_min[j], + peak_points[x_start[i]]: + peak_points[x_end[i]]]): + args_emptysep.add(i) + # print(j, i, "%d:%d" % (y_mid[j], y_mid[i]), "%d:%d" % (x_start[i], x_end[i]), "empty next sep") + args_to_be_kept = [arg for arg in args_ysorted + if arg not in args_emptysep] + x_start = x_start[args_to_be_kept] + x_end = x_end[args_to_be_kept] + y_min = y_min[args_to_be_kept] + y_mid = y_mid[args_to_be_kept] + y_max = y_max[args_to_be_kept] - remained_sep_indexes=set(range(len(x_start)))-set(np.unique(deleted) ) - #print(remained_sep_indexes,'remained_sep_indexes') - mother=[]#if it has mother - child=[] - for index_i in remained_sep_indexes: - have_mother=0 - have_child=0 - nodes_ind=set(range(x_start[index_i],x_end[index_i]+1)) - for index_j in remained_sep_indexes: - nodes_ind_j=set(range(x_start[index_j],x_end[index_j]+1)) - if nodes_indnodes_ind_j: - have_child=1 - mother.append(have_mother) - child.append(have_child) - - #print(mother,'mother') - #print(len(remained_sep_indexes)) - #print(len(remained_sep_indexes),len(x_start),len(x_end),len(y_sep),'lens') - y_lines_without_mother=[] - x_start_without_mother=[] - x_end_without_mother=[] - - y_lines_with_child_without_mother=[] - x_start_with_child_without_mother=[] - x_end_with_child_without_mother=[] - - mother = np.array(mother) - child = np.array(child) - #print(mother,'mother') - #print(child,'child') - remained_sep_indexes = np.array(list(remained_sep_indexes)) - x_start = np.array(x_start) - x_end = np.array(x_end) - y_sep = np.array(y_sep) - - if len(remained_sep_indexes)>1: - #print(np.array(remained_sep_indexes),'np.array(remained_sep_indexes)') - #print(np.array(mother),'mother') - remained_sep_indexes_without_mother = remained_sep_indexes[mother==0] - remained_sep_indexes_with_child_without_mother = remained_sep_indexes[(mother==0) & (child==1)] - #print(remained_sep_indexes_without_mother,'remained_sep_indexes_without_mother') - #print(remained_sep_indexes_without_mother,'remained_sep_indexes_without_mother') - - x_end_with_child_without_mother = x_end[remained_sep_indexes_with_child_without_mother] - x_start_with_child_without_mother = x_start[remained_sep_indexes_with_child_without_mother] - y_lines_with_child_without_mother = y_sep[remained_sep_indexes_with_child_without_mother] - - reading_orther_type=0 - x_end_without_mother = x_end[remained_sep_indexes_without_mother] - x_start_without_mother = x_start[remained_sep_indexes_without_mother] - y_lines_without_mother = y_sep[remained_sep_indexes_without_mother] - - if len(remained_sep_indexes_without_mother)>=2: - for i in range(len(remained_sep_indexes_without_mother)-1): - nodes_i=set(range(x_start[remained_sep_indexes_without_mother[i]], - x_end[remained_sep_indexes_without_mother[i]] - # + 1 - )) - for j in range(i+1,len(remained_sep_indexes_without_mother)): - nodes_j=set(range(x_start[remained_sep_indexes_without_mother[j]], - x_end[remained_sep_indexes_without_mother[j]] - # + 1 - )) - set_diff = nodes_i - nodes_j - if set_diff != nodes_i: - reading_orther_type = 1 - else: - reading_orther_type = 0 - #print(reading_orther_type,'javab') - #print(y_lines_with_child_without_mother,'y_lines_with_child_without_mother') - #print(x_start_with_child_without_mother,'x_start_with_child_without_mother') - #print(x_end_with_child_without_mother,'x_end_with_hild_without_mother') - - len_sep_with_child = len(child[child==1]) - - #print(len_sep_with_child,'len_sep_with_child') - there_is_sep_with_child = 0 - if len_sep_with_child >= 1: - there_is_sep_with_child = 1 - #print(all_args_uniq,'all_args_uniq') - #print(args_to_be_unified,'args_to_be_unified') - - return (reading_orther_type, - x_start_returned, - x_end_returned, - y_sep_returned, - y_diff_returned, - y_lines_without_mother, - x_start_without_mother, - x_end_without_mother, - there_is_sep_with_child, - y_lines_with_child_without_mother, - x_start_with_child_without_mother, - x_end_with_child_without_mother, - new_main_sep_y) + return (x_start, + x_end, + y_min, + y_mid, + y_max) def box2rect(box: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]: return (box[1], box[1] + box[3], @@ -393,139 +242,183 @@ def find_num_col_deskew(regions_without_separators, sigma_, multiplier=3.8): return np.std(z) def find_num_col( - regions_without_separators, - num_col_classifier, - tables, - multiplier=3.8, + regions_without_separators, + num_col_classifier, + tables, + multiplier=3.8, + unbalanced=False, + vertical_separators=None ): if not regions_without_separators.any(): return 0, [] - #plt.imshow(regions_without_separators) - #plt.show() + if vertical_separators is None: + vertical_separators = np.zeros_like(regions_without_separators) regions_without_separators_0 = regions_without_separators.sum(axis=0) - ##plt.plot(regions_without_separators_0) - ##plt.show() - sigma_ = 35 # 70#35 - meda_n_updown = regions_without_separators_0[len(regions_without_separators_0) :: -1] + vertical_separators_0 = vertical_separators.sum(axis=0) + # fig, (ax1, ax2) = plt.subplots(2, sharex=True) + # ax1.imshow(regions_without_separators, aspect="auto") + # ax2.plot(regions_without_separators_0) + # plt.show() + sigma_ = 25 # 70#35 + meda_n_updown = regions_without_separators_0[::-1] first_nonzero = next((i for i, x in enumerate(regions_without_separators_0) if x), 0) last_nonzero = next((i for i, x in enumerate(meda_n_updown) if x), 0) last_nonzero = len(regions_without_separators_0) - last_nonzero + last_nonzero = last_nonzero - 50 #- 100 + first_nonzero = first_nonzero + 50 #+ 200 + last_offmargin = len(regions_without_separators_0) - 170 #370 + first_offmargin = 170 #370 + x = vertical_separators_0 y = regions_without_separators_0 # [first_nonzero:last_nonzero] - y_help = np.zeros(len(y) + 20) - y_help[10 : len(y) + 10] = y - x = np.arange(len(y)) - zneg_rev = -y_help + np.max(y_help) - zneg = np.zeros(len(zneg_rev) + 20) - zneg[10 : len(zneg_rev) + 10] = zneg_rev + y_help = np.pad(y, (10, 10), constant_values=(0, 0)) + zneg_rev = y.max() - y_help + zneg = np.pad(zneg_rev, (10, 10), constant_values=(0, 0)) + x = gaussian_filter1d(x, sigma_) z = gaussian_filter1d(y, sigma_) zneg = gaussian_filter1d(zneg, sigma_) - peaks_neg, _ = find_peaks(zneg, height=0) - #plt.plot(zneg) - #plt.plot(peaks_neg, zneg[peaks_neg], 'rx') - #plt.show() peaks, _ = find_peaks(z, height=0) + peaks_neg, _ = find_peaks(zneg, height=0) + # _, (ax1, ax2) = plt.subplots(2, sharex=True) + # ax1.set_title("z") + # ax1.plot(z) + # ax1.scatter(peaks, z[peaks]) + # ax1.axvline(0.06 * len(y), label="first") + # ax1.axvline(0.94 * len(y), label="last") + # ax1.text(0.06 * len(y), 0, "first", rotation=90) + # ax1.text(0.94 * len(y), 0, "last", rotation=90) + # ax1.axhline(10, label="minimum") + # ax1.text(0, 10, "minimum") + # ax2.set_title("zneg") + # ax2.plot(zneg) + # ax2.scatter(peaks_neg, zneg[peaks_neg]) + # ax2.axvline(first_nonzero, label="first nonzero") + # ax2.axvline(last_nonzero, label="last nonzero") + # ax2.text(first_nonzero, 0, "first nonzero", rotation=90) + # ax2.text(last_nonzero, 0, "last nonzero", rotation=90) + # ax2.axvline(first_offmargin, label="first offmargin") + # ax2.axvline(last_offmargin, label="last offmargin") + # ax2.text(first_offmargin, 0, "first offmargin", rotation=90) + # ax2.text(last_offmargin, 0, "last offmargin", rotation=90) + # plt.show() peaks_neg = peaks_neg - 10 - 10 - last_nonzero = last_nonzero - 100 - first_nonzero = first_nonzero + 200 - - peaks_neg = peaks_neg[(peaks_neg > first_nonzero) & - (peaks_neg < last_nonzero)] - peaks = peaks[(peaks > 0.06 * regions_without_separators.shape[1]) & - (peaks < 0.94 * regions_without_separators.shape[1])] - peaks_neg = peaks_neg[(peaks_neg > 370) & - (peaks_neg < (regions_without_separators.shape[1] - 370))] + # print("raw peaks", peaks) + peaks = peaks[(peaks > 0.06 * len(y)) & + (peaks < 0.94 * len(y))] + # print("non-marginal peaks", peaks) interest_pos = z[peaks] + # print("interest_pos", interest_pos) interest_pos = interest_pos[interest_pos > 10] if not interest_pos.any(): return 0, [] + # plt.plot(z) # plt.show() + #print("raw peaks_neg", peaks_neg) + peaks_neg = peaks_neg[(peaks_neg > first_nonzero) & + (peaks_neg < last_nonzero)] + #print("non-zero peaks_neg", peaks_neg) + peaks_neg = peaks_neg[(peaks_neg > first_offmargin) & + (peaks_neg < last_offmargin)] + #print("non-marginal peaks_neg", peaks_neg) interest_neg = z[peaks_neg] + #print("interest_neg", interest_neg) if not interest_neg.any(): return 0, [] min_peaks_pos = np.min(interest_pos) max_peaks_pos = np.max(interest_pos) - if max_peaks_pos / min_peaks_pos >= 35: + #print(min_peaks_pos, max_peaks_pos, max_peaks_pos / min_peaks_pos, 'minmax') + if max_peaks_pos / (min_peaks_pos or 1e-9) >= 35: min_peaks_pos = np.mean(interest_pos) min_peaks_neg = 0 # np.min(interest_neg) - # print(np.min(interest_pos),np.max(interest_pos),np.max(interest_pos)/np.min(interest_pos),'minmax') + # cutoff criterion: fixed fraction of lowest column height dis_talaei = (min_peaks_pos - min_peaks_neg) / multiplier grenze = min_peaks_pos - dis_talaei - # np.mean(y[peaks_neg[0]:peaks_neg[len(peaks_neg)-1]])-np.std(y[peaks_neg[0]:peaks_neg[len(peaks_neg)-1]])/2.0 + #np.mean(y[peaks_neg[0]:peaks_neg[-1]])-np.std(y[peaks_neg[0]:peaks_neg[-1]])/2.0 + + # extra criterion: fixed multiple of lowest gap height + # print("grenze", grenze, multiplier * (5 + np.min(interest_neg))) + grenze = min(grenze, multiplier * (5 + np.min(interest_neg))) # print(interest_neg,'interest_neg') # print(grenze,'grenze') # print(min_peaks_pos,'min_peaks_pos') # print(dis_talaei,'dis_talaei') # print(peaks_neg,'peaks_neg') + # fig, (ax1, ax2) = plt.subplots(2, sharex=True) + # ax1.imshow(regions_without_separators + 5 * vertical_separators, aspect="auto") + # ax2.plot(z, color='red', label='z') + # ax2.plot(zneg[20:], color='blue', label='zneg') + # ax2.plot(x, color='green', label='vsep') + # ax2.scatter(peaks_neg, z[peaks_neg], color='red') + # ax2.scatter(peaks_neg, zneg[20:][peaks_neg], color='blue') + # ax2.axhline(min_peaks_pos, color='red') + # ax2.axhline(grenze, color='blue') + # ax2.annotate("min_peaks_pos", xy=(0, min_peaks_pos), color='red') + # ax2.annotate("grenze", xy=(0, grenze), color='blue') + # ax2.text(0, grenze, "grenze") + # ax2.legend() + # plt.show() + # print("vsep", x[peaks_neg]) + interest_neg = interest_neg - x[peaks_neg] interest_neg_fin = interest_neg[(interest_neg < grenze)] peaks_neg_fin = peaks_neg[(interest_neg < grenze)] - # interest_neg_fin=interest_neg[(interest_neg= 3: - index_sort_interest_neg_fin= np.argsort(interest_neg_fin) - peaks_neg_sorted = np.array(peaks_neg)[index_sort_interest_neg_fin] - interest_neg_fin_sorted = np.array(interest_neg_fin)[index_sort_interest_neg_fin] + # found too few columns here: ignore 'grenze' and take the deepest N peaks + sort_by_height = np.argsort(interest_neg)[:num_col_classifier] + peaks_neg_fin = peaks_neg[sort_by_height] + interest_neg_fin = interest_neg[sort_by_height] + # print(peaks_neg_fin, "peaks_neg[sorted_by_height]") + sort_by_pos = np.argsort(peaks_neg_fin) + peaks_neg_fin = peaks_neg_fin[sort_by_pos] + interest_neg_fin = interest_neg_fin[sort_by_pos] - if len(index_sort_interest_neg_fin)>=num_col_classifier: - peaks_neg_fin = list( peaks_neg_sorted[:num_col_classifier] ) - interest_neg_fin = list( interest_neg_fin_sorted[:num_col_classifier] ) - else: - peaks_neg_fin = peaks_neg[:] - interest_neg_fin = interest_neg[:] - - num_col = (len(interest_neg_fin)) + 1 + num_col = len(interest_neg_fin) + 1 # print(peaks_neg_fin,'peaks_neg_fin') # print(num_col,'diz') - p_l = 0 - p_u = len(y) - 1 - p_m = int(len(y) / 2.0) - p_g_l = int(len(y) / 4.0) - p_g_u = len(y) - int(len(y) / 4.0) - - if num_col == 3: - if ((peaks_neg_fin[0] > p_g_u and - peaks_neg_fin[1] > p_g_u) or - (peaks_neg_fin[0] < p_g_l and - peaks_neg_fin[1] < p_g_l) or - (peaks_neg_fin[0] + 200 < p_m and - peaks_neg_fin[1] < p_m) or - (peaks_neg_fin[0] - 200 > p_m and - peaks_neg_fin[1] > p_m)): - num_col = 1 - peaks_neg_fin = [] - - if num_col == 2: - if (peaks_neg_fin[0] > p_g_u or - peaks_neg_fin[0] < p_g_l): - num_col = 1 - peaks_neg_fin = [] + # cancel if resulting split is highly unbalanced across available width + if unbalanced: + pass + elif ((num_col == 3 and + ((peaks_neg_fin[0] > 0.75 * len(y) and + peaks_neg_fin[1] > 0.75 * len(y)) or + (peaks_neg_fin[0] < 0.25 * len(y) and + peaks_neg_fin[1] < 0.25 * len(y)) or + (peaks_neg_fin[0] < 0.5 * len(y) - 200 and + peaks_neg_fin[1] < 0.5 * len(y)) or + (peaks_neg_fin[0] > 0.5 * len(y) + 200 and + peaks_neg_fin[1] > 0.5 * len(y)))) or + (num_col == 2 and + (peaks_neg_fin[0] > 0.75 * len(y) or + peaks_neg_fin[0] < 0.25 * len(y)))): + num_col = 1 + peaks_neg_fin = [] ##print(len(peaks_neg_fin)) + # filter out peaks that are too close (<400px) to each other: + # among each group, pick the position with smallest amount of text diff_peaks = np.abs(np.diff(peaks_neg_fin)) - cut_off = 400 + cut_off = 300 #400 peaks_neg_true = [] forest = [] - # print(len(peaks_neg_fin),'len_') - for i in range(len(peaks_neg_fin)): if i == 0: forest.append(peaks_neg_fin[i]) if i < len(peaks_neg_fin) - 1: if diff_peaks[i] <= cut_off: forest.append(peaks_neg_fin[i + 1]) - if diff_peaks[i] > cut_off: + else: # print(forest[np.argmin(z[forest]) ] ) if not isNaN(forest[np.argmin(z[forest])]): peaks_neg_true.append(forest[np.argmin(z[forest])]) @@ -537,68 +430,61 @@ def find_num_col( peaks_neg_true.append(forest[np.argmin(z[forest])]) num_col = len(peaks_neg_true) + 1 - p_l = 0 - p_u = len(y) - 1 - p_m = int(len(y) / 2.0) - p_quarter = int(len(y) / 5.0) - p_g_l = int(len(y) / 4.0) - p_g_u = len(y) - int(len(y) / 4.0) - - p_u_quarter = len(y) - p_quarter - + #print(peaks_neg_true, "peaks_neg_true") ##print(num_col,'early') - if num_col == 3: - if ((peaks_neg_true[0] > p_g_u and - peaks_neg_true[1] > p_g_u) or - (peaks_neg_true[0] < p_g_l and - peaks_neg_true[1] < p_g_l) or - (peaks_neg_true[0] < p_m and - peaks_neg_true[1] + 200 < p_m) or - (peaks_neg_true[0] - 200 > p_m and - peaks_neg_true[1] > p_m)): - num_col = 1 - peaks_neg_true = [] - elif (peaks_neg_true[0] < p_g_u and - peaks_neg_true[0] > p_g_l and - peaks_neg_true[1] > p_u_quarter): - peaks_neg_true = [peaks_neg_true[0]] - elif (peaks_neg_true[1] < p_g_u and - peaks_neg_true[1] > p_g_l and - peaks_neg_true[0] < p_quarter): - peaks_neg_true = [peaks_neg_true[1]] + # cancel if resulting split is highly unbalanced across available width + if unbalanced: + pass + elif ((num_col == 3 and + ((peaks_neg_true[0] > 0.75 * len(y) and + peaks_neg_true[1] > 0.75 * len(y)) or + (peaks_neg_true[0] < 0.25 * len(y) and + peaks_neg_true[1] < 0.25 * len(y)) or + (peaks_neg_true[0] < 0.5 * len(y) - 200 and + peaks_neg_true[1] < 0.5 * len(y)) or + (peaks_neg_true[0] > 0.5 * len(y) + 200 and + peaks_neg_true[1] > 0.5 * len(y)))) or + (num_col == 2 and + (peaks_neg_true[0] > 0.75 * len(y) or + peaks_neg_true[0] < 0.25 * len(y)))): + num_col = 1 + peaks_neg_true = [] + elif (num_col == 3 and + (peaks_neg_true[0] < 0.75 * len(y) and + peaks_neg_true[0] > 0.25 * len(y) and + peaks_neg_true[1] > 0.80 * len(y))): + num_col = 2 + peaks_neg_true = [peaks_neg_true[0]] + elif (num_col == 3 and + (peaks_neg_true[1] < 0.75 * len(y) and + peaks_neg_true[1] > 0.25 * len(y) and + peaks_neg_true[0] < 0.20 * len(y))): + num_col = 2 + peaks_neg_true = [peaks_neg_true[1]] - if num_col == 2: - if (peaks_neg_true[0] > p_g_u or - peaks_neg_true[0] < p_g_l): - num_col = 1 - peaks_neg_true = [] + # get rid of too narrow columns (not used) + # if np.count_nonzero(diff_peaks < 360): + # arg_help = np.arange(len(diff_peaks)) + # arg_help_ann = arg_help[diff_peaks < 360] + # peaks_neg_fin_new = [] + # for ii in range(len(peaks_neg_fin)): + # if ii in arg_help_ann: + # if interest_neg_fin[ii] < interest_neg_fin[ii + 1]: + # peaks_neg_fin_new.append(peaks_neg_fin[ii]) + # else: + # peaks_neg_fin_new.append(peaks_neg_fin[ii + 1]) - diff_peaks_abnormal = diff_peaks[diff_peaks < 360] - - if len(diff_peaks_abnormal) > 0: - arg_help = np.arange(len(diff_peaks)) - arg_help_ann = arg_help[diff_peaks < 360] - - peaks_neg_fin_new = [] - - for ii in range(len(peaks_neg_fin)): - if ii in arg_help_ann: - arg_min = np.argmin([interest_neg_fin[ii], interest_neg_fin[ii + 1]]) - if arg_min == 0: - peaks_neg_fin_new.append(peaks_neg_fin[ii]) - else: - peaks_neg_fin_new.append(peaks_neg_fin[ii + 1]) - - elif (ii - 1) not in arg_help_ann: - peaks_neg_fin_new.append(peaks_neg_fin[ii]) - else: - peaks_neg_fin_new = peaks_neg_fin + # elif (ii - 1) not in arg_help_ann: + # peaks_neg_fin_new.append(peaks_neg_fin[ii]) + # else: + # peaks_neg_fin_new = peaks_neg_fin # plt.plot(gaussian_filter1d(y, sigma_)) # plt.plot(peaks_neg_true,z[peaks_neg_true],'*') # plt.plot([0,len(y)], [grenze,grenze]) # plt.show() ##print(len(peaks_neg_true)) + #print(peaks_neg_true, "peaks_neg_true") return len(peaks_neg_true), peaks_neg_true def find_num_col_only_image(regions_without_separators, multiplier=3.8): @@ -810,6 +696,12 @@ def find_num_col_by_vertical_lines(regions_without_separators, multiplier=3.8): peaks, _ = find_peaks(z, height=0) # print(peaks,'peaksnew') + # fig, (ax1, ax2) = plt.subplots(2, sharex=True, suptitle='find_num_col_by_vertical_lines') + # ax1.imshow(regions_without_separators, aspect="auto") + # ax2.plot(z) + # ax2.scatter(peaks, z[peaks]) + # ax2.set_title('find_peaks(regions_without_separators.sum(axis=0), height=0)') + # plt.show() return peaks def return_regions_without_separators(regions_pre): @@ -1196,7 +1088,26 @@ def small_textlines_to_parent_adherence2(textlines_con, textline_iamge, num_col) textlines_con_changed.append(textlines_big_org_form) return textlines_con_changed -def order_of_regions(textline_mask, contours_main, contours_head, y_ref): +def order_of_regions(textline_mask, contours_main, contours_head, y_ref, x_ref): + """ + Order text region contours within a single column bbox in a top-down-left-right way. + + First, determine the vertical gaps. Then iterate over each vertical segment, + identifying the contours centered in that segment. Order them by their + horizontal center, and add them to the overall order. + + Arguments: + * textline_mask: the mask of the textline segmentation, cropped for that box + * contours_main: the paragraph text region contours expected to be here + * contours_head: the heading text region contours expected to be here + * y_ref: the vertical offset of that box within the page + * x_ref: the horizontal offset of that box within the page + + Returns: a tuple of + * the array of contour indexes overall within this box (i.e. into main+head) + * the array of types (1 for paragraph, 2 for heading) + * the array of contour indexes for the respective type (i.e. into contours_main or contours_head) + """ ##plt.imshow(textline_mask) ##plt.show() y = textline_mask.sum(axis=1) # horizontal projection profile @@ -1207,6 +1118,8 @@ def order_of_regions(textline_mask, contours_main, contours_head, y_ref): #z = gaussian_filter1d(y_padded, sigma_gaus) #peaks, _ = find_peaks(z, height=0) #peaks = peaks - 20 + ##plt.plot(z) + ##plt.show() zneg_rev = np.max(y_padded) - y_padded zneg = np.zeros(len(zneg_rev) + 40) zneg[20 : len(zneg_rev) + 20] = zneg_rev @@ -1215,15 +1128,16 @@ def order_of_regions(textline_mask, contours_main, contours_head, y_ref): peaks_neg, _ = find_peaks(zneg, height=0) peaks_neg = peaks_neg - 20 - 20 - ##plt.plot(z) - ##plt.show() - cx_main, cy_main = find_center_of_contours(contours_main) - cx_head, cy_head = find_center_of_contours(contours_head) - - peaks_neg_new = np.append(np.insert(peaks_neg, 0, 0), textline_mask.shape[0]) + peaks_neg_new = np.array([0] + + # peaks can be beyond box due to padding and smoothing + [peak for peak in peaks_neg + if 0 < peak and peak < textline_mask.shape[0]] + + [textline_mask.shape[0]]) # offset from bbox of mask peaks_neg_new += y_ref + cx_main, cy_main = find_center_of_contours(contours_main) + cx_head, cy_head = find_center_of_contours(contours_head) # assert not len(cy_main) or np.min(peaks_neg_new) <= np.min(cy_main) and np.max(cy_main) <= np.max(peaks_neg_new) # assert not len(cy_head) or np.min(peaks_neg_new) <= np.min(cy_head) and np.max(cy_head) <= np.max(peaks_neg_new) @@ -1248,6 +1162,23 @@ def order_of_regions(textline_mask, contours_main, contours_head, y_ref): indexes_in, types_in, cxs_in, cys_in, typed_indexes_in = \ matrix_of_orders[(matrix_of_orders[:, 3] >= top) & (matrix_of_orders[:, 3] < bot)].T + # if indexes_in.size: + # img = textline_mask.copy() + # plt.imshow(img) + # plt.gca().add_patch(patches.Rectangle((0, top-y_ref), img.shape[1], bot-top, alpha=0.5, color='gray')) + # xrange = np.arange(0, img.shape[1], 50) + # yrange = np.arange(0, img.shape[0], 50) + # plt.gca().set_xticks(xrange, xrange + x_ref) + # plt.gca().set_yticks(yrange, yrange + y_ref) + # for idx, type_, cx, cy in zip(typed_indexes_in, types_in, cxs_in, cys_in): + # cnt = (contours_main if type_ == 1 else contours_head)[idx] + # col = 'red' if type_ == 1 else 'blue' + # plt.scatter(cx - x_ref, cy - y_ref, 20, c=col, marker='o') + # plt.text(cx - x_ref, cy - y_ref, str(idx), c=col) + # plt.gca().add_patch(patches.Polygon(cnt[:, 0] - [[x_ref, y_ref]], closed=False, fill=False, color=col)) + # plt.title("box contours centered in %d:%d (red=main / blue=heading)" % (top, bot)) + # plt.show() + sorted_inside = np.argsort(cxs_in) final_indexers_sorted.extend(indexes_in[sorted_inside]) final_types.extend(types_in[sorted_inside]) @@ -1255,17 +1186,36 @@ def order_of_regions(textline_mask, contours_main, contours_head, y_ref): ##matrix_of_orders[:len_main,4]=final_indexers_sorted[:] - # assert len(final_indexers_sorted) == len(contours_main) + len(contours_head) - # assert not len(final_indexers_sorted) or max(final_index_type) == max(len(contours_main) + assert len(set(final_indexers_sorted)) == len(contours_main) + len(contours_head) + assert set(final_index_type) == set(range(len(contours_main))).union(range(len(contours_head))) return np.array(final_indexers_sorted), np.array(final_types), np.array(final_index_type) def combine_hor_lines_and_delete_cross_points_and_get_lines_features_back_new( - img_p_in_ver, img_in_hor,num_col_classifier): + img_p_in_ver: np.ndarray, + img_p_in_hor: np.ndarray, + num_col_classifier: int, +) -> Tuple[np.ndarray, List[float]]: + """ + Given a horizontal and vertical separator mask, combine horizontal separators + (where possible) and make sure they do not cross each other. + + Arguments: + * img_p_in_ver: mask of vertical separators + * img_p_in_hor: mask of horizontal separators + * num_col_classifier: predicted (expected) number of columns + + Returns: a tuple of + * the final horizontal separators + * the y coordinates with horizontal separators spanning the full width + """ + + # cut horizontal seps by vertical seps + img_p_in_hor[img_p_in_ver > 0] = 0 #img_p_in_ver = cv2.erode(img_p_in_ver, self.kernel, iterations=2) _, thresh = cv2.threshold(img_p_in_ver, 0, 255, 0) - contours_lines_ver, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + contours_lines_ver, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) slope_lines_ver, _, x_min_main_ver, _, _, _, y_min_main_ver, y_max_main_ver, cx_main_ver = \ find_features_of_lines(contours_lines_ver) for i in range(len(x_min_main_ver)): @@ -1274,79 +1224,94 @@ def combine_hor_lines_and_delete_cross_points_and_get_lines_features_back_new( int(cx_main_ver[i])-25: int(cx_main_ver[i])+25] = 0 img_p_in_ver[int(y_max_main_ver[i])-30: - int(y_max_main_ver[i]), + int(y_max_main_ver[i]+1), int(cx_main_ver[i])-25: int(cx_main_ver[i])+25] = 0 + height, width = img_p_in_ver.shape - _, thresh = cv2.threshold(img_in_hor, 0, 255, 0) - contours_lines_hor, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + _, thresh = cv2.threshold(img_p_in_hor, 0, 255, 0) + contours_lines_hor, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - slope_lines_hor, dist_x_hor, x_min_main_hor, x_max_main_hor, cy_main_hor, _, _, _, _ = \ - find_features_of_lines(contours_lines_hor) - x_width_smaller_than_acolumn_width=img_in_hor.shape[1]/float(num_col_classifier+1.) + (slope_lines_hor, + dist_x_hor, + x_min_main_hor, + x_max_main_hor, + cy_main_hor, _, + y_min_main_hor, + y_max_main_hor, + _) = find_features_of_lines(contours_lines_hor) - len_lines_bigger_than_x_width_smaller_than_acolumn_width=len( dist_x_hor[dist_x_hor>=x_width_smaller_than_acolumn_width] ) - len_lines_bigger_than_x_width_smaller_than_acolumn_width_per_column=int(len_lines_bigger_than_x_width_smaller_than_acolumn_width / - float(num_col_classifier)) - if len_lines_bigger_than_x_width_smaller_than_acolumn_width_per_column < 10: + avg_col_width = width / float(num_col_classifier + 1) + nseps_wider_than_than_avg_col_width = np.count_nonzero(dist_x_hor>=avg_col_width) + if nseps_wider_than_than_avg_col_width < 10 * num_col_classifier: args_hor=np.arange(len(slope_lines_hor)) - all_args_uniq=contours_in_same_horizon(cy_main_hor) - #print(all_args_uniq,'all_args_uniq') - if len(all_args_uniq)>0: - if type(all_args_uniq[0]) is list: - special_separators=[] - contours_new=[] - for dd in range(len(all_args_uniq)): - merged_all=None - some_args=args_hor[all_args_uniq[dd]] - some_cy=cy_main_hor[all_args_uniq[dd]] - some_x_min=x_min_main_hor[all_args_uniq[dd]] - some_x_max=x_max_main_hor[all_args_uniq[dd]] - - #img_in=np.zeros(separators_closeup_n[:,:,2].shape) - #print(img_p_in_ver.shape[1],some_x_max-some_x_min,'xdiff') - diff_x_some=some_x_max-some_x_min - for jv in range(len(some_args)): - img_p_in=cv2.fillPoly(img_in_hor, pts=[contours_lines_hor[some_args[jv]]], color=(1,1,1)) - if any(i_diff>(img_p_in_ver.shape[1]/float(3.3)) for i_diff in diff_x_some): - img_p_in[int(np.mean(some_cy))-5: - int(np.mean(some_cy))+5, - int(np.min(some_x_min)): - int(np.max(some_x_max)) ]=1 - sum_dis=dist_x_hor[some_args].sum() - diff_max_min_uniques=np.max(x_max_main_hor[some_args])-np.min(x_min_main_hor[some_args]) - - if (diff_max_min_uniques > sum_dis and - sum_dis / float(diff_max_min_uniques) > 0.85 and - diff_max_min_uniques / float(img_p_in_ver.shape[1]) > 0.85 and - np.std(dist_x_hor[some_args]) < 0.55 * np.mean(dist_x_hor[some_args])): - # print(dist_x_hor[some_args], - # dist_x_hor[some_args].sum(), - # np.min(x_min_main_hor[some_args]), - # np.max(x_max_main_hor[some_args]),'jalibdi') - # print(np.mean( dist_x_hor[some_args] ), - # np.std( dist_x_hor[some_args] ), - # np.var( dist_x_hor[some_args] ),'jalibdiha') - special_separators.append(np.mean(cy_main_hor[some_args])) - else: - img_p_in=img_in_hor - special_separators=[] - else: - img_p_in=img_in_hor + sep_pairs=contours_in_same_horizon(cy_main_hor) + img_p_in = np.copy(img_p_in_hor) + if len(sep_pairs): special_separators=[] + contours_new=[] + for pair in sep_pairs: + merged_all=None + some_args=args_hor[pair] + some_cy=cy_main_hor[pair] + some_x_min=x_min_main_hor[pair] + some_x_max=x_max_main_hor[pair] + some_y_min=y_min_main_hor[pair] + some_y_max=y_max_main_hor[pair] + if np.any(img_p_in_ver[some_y_min.min(): some_y_max.max(), + some_x_max.min(): some_x_min.max()]): + # print("horizontal pair cut by vertical sep", pair, some_args, some_cy, + # "%d:%d" % (some_x_min[0], some_x_max[0]), + # "%d:%d" % (some_x_min[1], some_x_max[1])) + continue - img_p_in_ver[img_p_in_ver == 255] = 1 - sep_ver_hor = img_p_in + img_p_in_ver - sep_ver_hor_cross = (sep_ver_hor == 2) * 1 - _, thresh = cv2.threshold(sep_ver_hor_cross.astype(np.uint8), 0, 255, 0) - contours_cross, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - center_cross = np.array(find_center_of_contours(contours_cross), dtype=int) - for cx, cy in center_cross.T: - img_p_in[cy - 30: cy + 30, cx + 5: cx + 40] = 0 - img_p_in[cy - 30: cy + 30, cx - 40: cx - 4] = 0 + #img_in=np.zeros(separators_closeup_n[:,:,2].shape) + #print(img_p_in_ver.shape[1],some_x_max-some_x_min,'xdiff') + sum_xspan = dist_x_hor[some_args].sum() + tot_xspan = np.max(x_max_main_hor[some_args]) - np.min(x_min_main_hor[some_args]) + dev_xspan = np.std(dist_x_hor[some_args]) / np.mean(dist_x_hor[some_args]) + if (tot_xspan > sum_xspan and # no x overlap + sum_xspan > 0.85 * tot_xspan): # x close to each other + # print("merging horizontal pair", pair, some_args, some_cy, + # "%d:%d" % (some_x_min[0], some_x_max[0]), + # "%d:%d" % (some_x_min[1], some_x_max[1])) + img_p_in[int(np.mean(some_cy)) - 5: + int(np.mean(some_cy)) + 5, + np.min(some_x_min): + np.max(some_x_max)] = 255 + + if (tot_xspan > sum_xspan and # no x overlap + sum_xspan > 0.85 * tot_xspan and # x close to each other + tot_xspan > 0.85 * width and # nearly full width + dev_xspan < 0.55): # similar x span + # print(dist_x_hor[some_args], + # dist_x_hor[some_args].sum(), + # np.min(x_min_main_hor[some_args]), + # np.max(x_max_main_hor[some_args]),'jalibdi') + # print(np.mean( dist_x_hor[some_args] ), + # np.std( dist_x_hor[some_args] ), + # np.var( dist_x_hor[some_args] ),'jalibdiha') + special_separators.append(np.mean(cy_main_hor[some_args])) + # print("special separator for midline", special_separators[-1]) + # plt.subplot(1, 2, 1, title='original horizontal (1) / vertical (2) seps') + # plt.imshow(1 * (img_p_in_hor > 0) + 2 * (img_p_in_ver > 0)) + # plt.subplot(1, 2, 2, title='extended horizontal seps') + # plt.imshow(img_p_in) + # plt.show() + else: + img_p_in = img_p_in_hor + special_separators = [] + + #img_p_in_ver[img_p_in_ver == 255] = 1 + # sep_ver_hor_cross = 255 * ((img_p_in > 0) & (img_p_in_ver > 0)) + # contours_cross, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + # center_cross = np.array(find_center_of_contours(contours_cross), dtype=int) + # for cx, cy in center_cross.T: + # img_p_in[cy - 30: cy + 30, cx + 5: cx + 40] = 0 + # img_p_in[cy - 30: cy + 30, cx - 40: cx - 4] = 0 else: - img_p_in=np.copy(img_in_hor) - special_separators=[] + img_p_in = np.copy(img_p_in_hor) + special_separators = [] return img_p_in, special_separators def return_points_with_boundies(peaks_neg_fin, first_point, last_point): @@ -1357,189 +1322,207 @@ def return_points_with_boundies(peaks_neg_fin, first_point, last_point): peaks_neg_tot.append(last_point) return peaks_neg_tot -def find_number_of_columns_in_document(region_pre_p, num_col_classifier, tables, label_lines, contours_h=None): - t_ins_c0 = time.time() - separators_closeup=( (region_pre_p[:,:]==label_lines))*1 - separators_closeup[0:110,:]=0 - separators_closeup[separators_closeup.shape[0]-150:,:]=0 +def find_number_of_columns_in_document( + region_pre_p: np.ndarray, + num_col_classifier: int, + tables: bool, + label_seps: int, + contours_h: List[np.ndarray] = None, + logger=None +) -> Tuple[int, List[int], np.ndarray, List[int], np.ndarray]: + """ + Extract vertical and horizontal separators, vertical splits and horizontal column boundaries on page. + + Arguments: + * region_pre_p: segmentation map of the page + * num_col_classifier: predicted (expected) number of columns of the page + * tables: whether tables may be present + * label_seps: segmentation map class label for separators + * contours_h: polygons of potential headings (serving as additional horizontal separators) + * logger + + Returns: a tuple of + * the actual number of columns found + * the x coordinates of the column boundaries + * an array of the separators (bounding boxes and types) + * the y coordinates of the page splits + * a mask of the separators + """ + if logger is None: + logger = getLogger(__package__) + + separators_closeup = 1 * (region_pre_p == label_seps) + separators_closeup[0:110] = 0 + separators_closeup[-150:] = 0 kernel = np.ones((5,5),np.uint8) - separators_closeup=separators_closeup.astype(np.uint8) - separators_closeup = cv2.dilate(separators_closeup,kernel,iterations = 1) - separators_closeup = cv2.erode(separators_closeup,kernel,iterations = 1) + separators_closeup = separators_closeup.astype(np.uint8) + separators_closeup = cv2.morphologyEx(separators_closeup, cv2.MORPH_CLOSE, kernel, iterations=1) - separators_closeup_new=np.zeros((separators_closeup.shape[0] ,separators_closeup.shape[1] )) - separators_closeup_n=np.copy(separators_closeup) - separators_closeup_n=separators_closeup_n.astype(np.uint8) + separators_closeup_n = separators_closeup.astype(np.uint8) # to be returned - separators_closeup_n_binary=np.zeros(( separators_closeup_n.shape[0],separators_closeup_n.shape[1]) ) - separators_closeup_n_binary[:,:]=separators_closeup_n[:,:] - separators_closeup_n_binary[:,:][separators_closeup_n_binary[:,:]!=0]=1 + separators_closeup_n_binary = separators_closeup_n.copy() - _, thresh_e = cv2.threshold(separators_closeup_n_binary, 0, 255, 0) - contours_line_e, _ = cv2.findContours(thresh_e.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - _, dist_xe, _, _, _, _, y_min_main, y_max_main, _ = \ - find_features_of_lines(contours_line_e) - dist_ye = y_max_main - y_min_main - args_e=np.arange(len(contours_line_e)) - args_hor_e=args_e[(dist_ye<=50) & - (dist_xe>=3*dist_ye)] - cnts_hor_e=[] - for ce in args_hor_e: - cnts_hor_e.append(contours_line_e[ce]) + # find horizontal lines by contour properties + contours_sep_e, _ = cv2.findContours(separators_closeup_n_binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + cnts_hor_e = [] + for cnt in contours_sep_e: + max_xe = cnt[:, 0, 0].max() + min_xe = cnt[:, 0, 0].min() + max_ye = cnt[:, 0, 1].max() + min_ye = cnt[:, 0, 1].min() + med_ye = int(np.median(cnt[:, 0, 1])) + dist_xe = max_xe - min_xe + dist_ye = max_ye - min_ye + if dist_ye <= 50 and dist_xe >= 3 * dist_ye: + cnts_hor_e.append(cnt) - separators_closeup_n_binary=cv2.fillPoly(separators_closeup_n_binary, pts=cnts_hor_e, color=0) - gray = cv2.bitwise_not(separators_closeup_n_binary) - gray=gray.astype(np.uint8) + # delete horizontal contours (leaving only the edges) + separators_closeup_n_binary = cv2.fillPoly(separators_closeup_n_binary, pts=cnts_hor_e, color=0) + edges = cv2.adaptiveThreshold(separators_closeup_n_binary * 255, 255, + cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 15, -2) + horizontal = np.copy(edges) + vertical = np.copy(edges) - bw = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, \ - cv2.THRESH_BINARY, 15, -2) - horizontal = np.copy(bw) - vertical = np.copy(bw) - - cols = horizontal.shape[1] - horizontal_size = cols // 30 - # Create structure element for extracting horizontal lines through morphology operations + horizontal_size = horizontal.shape[1] // 30 + # find horizontal lines by morphology horizontalStructure = cv2.getStructuringElement(cv2.MORPH_RECT, (horizontal_size, 1)) - # Apply morphology operations - horizontal = cv2.erode(horizontal, horizontalStructure) - horizontal = cv2.dilate(horizontal, horizontalStructure) - - kernel = np.ones((5,5),np.uint8) - horizontal = cv2.dilate(horizontal,kernel,iterations = 2) - horizontal = cv2.erode(horizontal,kernel,iterations = 2) + horizontal = cv2.morphologyEx(horizontal, cv2.MORPH_OPEN, horizontalStructure) + horizontal = cv2.morphologyEx(horizontal, cv2.MORPH_CLOSE, kernel, iterations=2) + # re-insert deleted horizontal contours horizontal = cv2.fillPoly(horizontal, pts=cnts_hor_e, color=255) - rows = vertical.shape[0] - verticalsize = rows // 30 - # Create structure element for extracting vertical lines through morphology operations - verticalStructure = cv2.getStructuringElement(cv2.MORPH_RECT, (1, verticalsize)) - # Apply morphology operations - vertical = cv2.erode(vertical, verticalStructure) - vertical = cv2.dilate(vertical, verticalStructure) - vertical = cv2.dilate(vertical,kernel,iterations = 1) + vertical_size = vertical.shape[0] // 30 + # find vertical lines by morphology + verticalStructure = cv2.getStructuringElement(cv2.MORPH_RECT, (1, vertical_size)) + vertical = cv2.morphologyEx(vertical, cv2.MORPH_OPEN, verticalStructure) + vertical = cv2.dilate(vertical, kernel, iterations=1) horizontal, special_separators = \ combine_hor_lines_and_delete_cross_points_and_get_lines_features_back_new( vertical, horizontal, num_col_classifier) - separators_closeup_new[:,:][vertical[:,:]!=0]=1 - separators_closeup_new[:,:][horizontal[:,:]!=0]=1 - _, thresh = cv2.threshold(vertical, 0, 255, 0) - contours_line_vers, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - slope_lines, dist_x, x_min_main, x_max_main, cy_main, slope_lines_org, y_min_main, y_max_main, cx_main = \ - find_features_of_lines(contours_line_vers) + contours_sep_vers, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + slope_seps, dist_x, x_min_seps, x_max_seps, cy_seps, slope_seps_org, y_min_seps, y_max_seps, cx_seps = \ + find_features_of_lines(contours_sep_vers) - args=np.arange(len(slope_lines)) - args_ver=args[slope_lines==1] - dist_x_ver=dist_x[slope_lines==1] - y_min_main_ver=y_min_main[slope_lines==1] - y_max_main_ver=y_max_main[slope_lines==1] - x_min_main_ver=x_min_main[slope_lines==1] - x_max_main_ver=x_max_main[slope_lines==1] - cx_main_ver=cx_main[slope_lines==1] - dist_y_ver=y_max_main_ver-y_min_main_ver + args=np.arange(len(slope_seps)) + args_ver=args[slope_seps==1] + dist_x_ver=dist_x[slope_seps==1] + y_min_seps_ver=y_min_seps[slope_seps==1] + y_max_seps_ver=y_max_seps[slope_seps==1] + x_min_seps_ver=x_min_seps[slope_seps==1] + x_max_seps_ver=x_max_seps[slope_seps==1] + cx_seps_ver=cx_seps[slope_seps==1] + dist_y_ver=y_max_seps_ver-y_min_seps_ver len_y=separators_closeup.shape[0]/3.0 _, thresh = cv2.threshold(horizontal, 0, 255, 0) - contours_line_hors, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - slope_lines, dist_x, x_min_main, x_max_main, cy_main, slope_lines_org, y_min_main, y_max_main, cx_main = \ - find_features_of_lines(contours_line_hors) + contours_sep_hors, _ = cv2.findContours(thresh.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + slope_seps, dist_x, x_min_seps, x_max_seps, cy_seps, slope_seps_org, y_min_seps, y_max_seps, cx_seps = \ + find_features_of_lines(contours_sep_hors) - slope_lines_org_hor=slope_lines_org[slope_lines==0] - args=np.arange(len(slope_lines)) + slope_seps_org_hor=slope_seps_org[slope_seps==0] + args=np.arange(len(slope_seps)) len_x=separators_closeup.shape[1]/5.0 - dist_y=np.abs(y_max_main-y_min_main) + dist_y=np.abs(y_max_seps-y_min_seps) - args_hor=args[slope_lines==0] - dist_x_hor=dist_x[slope_lines==0] - y_min_main_hor=y_min_main[slope_lines==0] - y_max_main_hor=y_max_main[slope_lines==0] - x_min_main_hor=x_min_main[slope_lines==0] - x_max_main_hor=x_max_main[slope_lines==0] - dist_y_hor=dist_y[slope_lines==0] - cy_main_hor=cy_main[slope_lines==0] + args_hor=args[slope_seps==0] + dist_x_hor=dist_x[slope_seps==0] + y_min_seps_hor=y_min_seps[slope_seps==0] + y_max_seps_hor=y_max_seps[slope_seps==0] + x_min_seps_hor=x_min_seps[slope_seps==0] + x_max_seps_hor=x_max_seps[slope_seps==0] + dist_y_hor=dist_y[slope_seps==0] + cy_seps_hor=cy_seps[slope_seps==0] args_hor=args_hor[dist_x_hor>=len_x/2.0] - x_max_main_hor=x_max_main_hor[dist_x_hor>=len_x/2.0] - x_min_main_hor=x_min_main_hor[dist_x_hor>=len_x/2.0] - cy_main_hor=cy_main_hor[dist_x_hor>=len_x/2.0] - y_min_main_hor=y_min_main_hor[dist_x_hor>=len_x/2.0] - y_max_main_hor=y_max_main_hor[dist_x_hor>=len_x/2.0] + x_max_seps_hor=x_max_seps_hor[dist_x_hor>=len_x/2.0] + x_min_seps_hor=x_min_seps_hor[dist_x_hor>=len_x/2.0] + cy_seps_hor=cy_seps_hor[dist_x_hor>=len_x/2.0] + y_min_seps_hor=y_min_seps_hor[dist_x_hor>=len_x/2.0] + y_max_seps_hor=y_max_seps_hor[dist_x_hor>=len_x/2.0] dist_y_hor=dist_y_hor[dist_x_hor>=len_x/2.0] - slope_lines_org_hor=slope_lines_org_hor[dist_x_hor>=len_x/2.0] + slope_seps_org_hor=slope_seps_org_hor[dist_x_hor>=len_x/2.0] dist_x_hor=dist_x_hor[dist_x_hor>=len_x/2.0] - matrix_of_lines_ch=np.zeros((len(cy_main_hor)+len(cx_main_ver),10)) - matrix_of_lines_ch[:len(cy_main_hor),0]=args_hor - matrix_of_lines_ch[len(cy_main_hor):,0]=args_ver - matrix_of_lines_ch[len(cy_main_hor):,1]=cx_main_ver - matrix_of_lines_ch[:len(cy_main_hor),2]=x_min_main_hor+50#x_min_main_hor+150 - matrix_of_lines_ch[len(cy_main_hor):,2]=x_min_main_ver - matrix_of_lines_ch[:len(cy_main_hor),3]=x_max_main_hor-50#x_max_main_hor-150 - matrix_of_lines_ch[len(cy_main_hor):,3]=x_max_main_ver - matrix_of_lines_ch[:len(cy_main_hor),4]=dist_x_hor - matrix_of_lines_ch[len(cy_main_hor):,4]=dist_x_ver - matrix_of_lines_ch[:len(cy_main_hor),5]=cy_main_hor - matrix_of_lines_ch[:len(cy_main_hor),6]=y_min_main_hor - matrix_of_lines_ch[len(cy_main_hor):,6]=y_min_main_ver - matrix_of_lines_ch[:len(cy_main_hor),7]=y_max_main_hor - matrix_of_lines_ch[len(cy_main_hor):,7]=y_max_main_ver - matrix_of_lines_ch[:len(cy_main_hor),8]=dist_y_hor - matrix_of_lines_ch[len(cy_main_hor):,8]=dist_y_ver - matrix_of_lines_ch[len(cy_main_hor):,9]=1 + matrix_of_seps_ch = np.zeros((len(cy_seps_hor)+len(cx_seps_ver), 10), dtype=int) + matrix_of_seps_ch[:len(cy_seps_hor),0]=args_hor + matrix_of_seps_ch[len(cy_seps_hor):,0]=args_ver + matrix_of_seps_ch[len(cy_seps_hor):,1]=cx_seps_ver + matrix_of_seps_ch[:len(cy_seps_hor),2]=x_min_seps_hor+50#x_min_seps_hor+150 + matrix_of_seps_ch[len(cy_seps_hor):,2]=x_min_seps_ver + matrix_of_seps_ch[:len(cy_seps_hor),3]=x_max_seps_hor-50#x_max_seps_hor-150 + matrix_of_seps_ch[len(cy_seps_hor):,3]=x_max_seps_ver + matrix_of_seps_ch[:len(cy_seps_hor),4]=dist_x_hor + matrix_of_seps_ch[len(cy_seps_hor):,4]=dist_x_ver + matrix_of_seps_ch[:len(cy_seps_hor),5]=cy_seps_hor + matrix_of_seps_ch[:len(cy_seps_hor),6]=y_min_seps_hor + matrix_of_seps_ch[len(cy_seps_hor):,6]=y_min_seps_ver + matrix_of_seps_ch[:len(cy_seps_hor),7]=y_max_seps_hor + matrix_of_seps_ch[len(cy_seps_hor):,7]=y_max_seps_ver + matrix_of_seps_ch[:len(cy_seps_hor),8]=dist_y_hor + 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: - _, dist_x_head, x_min_main_head, x_max_main_head, cy_main_head, _, y_min_main_head, y_max_main_head, _ = \ + _, 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((matrix_of_lines_ch.shape[0]+len(cy_main_head),matrix_of_lines_ch.shape[1])) - matrix_l_n[:matrix_of_lines_ch.shape[0],:]=np.copy(matrix_of_lines_ch[:,:]) - args_head=np.arange(len(cy_main_head)) + len(cy_main_hor) + matrix_l_n = np.zeros((len(cy_head), matrix_of_seps_ch.shape[1]), dtype=int) + args_head = np.arange(len(cy_head)) + matrix_l_n[:, 0] = args_head + matrix_l_n[:, 2] = x_min_head + matrix_l_n[:, 3] = x_max_head + matrix_l_n[:, 4] = dist_x_head + matrix_l_n[:, 5] = cy_head + matrix_l_n[:, 6] = y_min_head + matrix_l_n[:, 7] = y_max_head + matrix_l_n[:, 8] = y_max_head - y_min_head + matrix_l_n[:, 9] = 2 # mark as heading (so it can be split into 2 horizontal separators as needed) + matrix_of_seps_ch = np.append( + matrix_of_seps_ch, matrix_l_n, axis=0) - matrix_l_n[matrix_of_lines_ch.shape[0]:,0]=args_head - matrix_l_n[matrix_of_lines_ch.shape[0]:,2]=x_min_main_head+30 - matrix_l_n[matrix_of_lines_ch.shape[0]:,3]=x_max_main_head-30 - matrix_l_n[matrix_of_lines_ch.shape[0]:,4]=dist_x_head - matrix_l_n[matrix_of_lines_ch.shape[0]:,5]=y_min_main_head-3-8 - matrix_l_n[matrix_of_lines_ch.shape[0]:,6]=y_min_main_head-5-8 - matrix_l_n[matrix_of_lines_ch.shape[0]:,7]=y_max_main_head#y_min_main_head+1-8 - matrix_l_n[matrix_of_lines_ch.shape[0]:,8]=4 - matrix_of_lines_ch=np.copy(matrix_l_n) + # ensure no seps are out of bounds + matrix_of_seps_ch[:, 1] = np.maximum(np.minimum(matrix_of_seps_ch[:, 1], region_pre_p.shape[1]), 0) + matrix_of_seps_ch[:, 2] = np.maximum(matrix_of_seps_ch[:, 2], 0) + matrix_of_seps_ch[:, 3] = np.minimum(matrix_of_seps_ch[:, 3], region_pre_p.shape[1]) + matrix_of_seps_ch[:, 5] = np.maximum(np.minimum(matrix_of_seps_ch[:, 5], region_pre_p.shape[0]), 0) + matrix_of_seps_ch[:, 6] = np.maximum(matrix_of_seps_ch[:, 6], 0) + matrix_of_seps_ch[:, 7] = np.minimum(matrix_of_seps_ch[:, 7], region_pre_p.shape[0]) + + cy_seps_splitters=cy_seps_hor[(x_min_seps_hor<=.16*region_pre_p.shape[1]) & + (x_max_seps_hor>=.84*region_pre_p.shape[1])] + cy_seps_splitters = np.append(cy_seps_splitters, special_separators) - cy_main_splitters=cy_main_hor[(x_min_main_hor<=.16*region_pre_p.shape[1]) & - (x_max_main_hor>=.84*region_pre_p.shape[1])] - cy_main_splitters=np.array( list(cy_main_splitters)+list(special_separators)) if contours_h is not None: - try: - cy_main_splitters_head=cy_main_head[(x_min_main_head<=.16*region_pre_p.shape[1]) & - (x_max_main_head>=.84*region_pre_p.shape[1])] - cy_main_splitters=np.array( list(cy_main_splitters)+list(cy_main_splitters_head)) - except: - pass - args_cy_splitter=np.argsort(cy_main_splitters) - cy_main_splitters_sort=cy_main_splitters[args_cy_splitter] + y_min_splitters_head = y_min_head[(x_min_head<=.16*region_pre_p.shape[1]) & + (x_max_head>=.84*region_pre_p.shape[1])] + y_max_splitters_head = y_max_head[(x_min_head<=.16*region_pre_p.shape[1]) & + (x_max_head>=.84*region_pre_p.shape[1])] + cy_seps_splitters = np.append(cy_seps_splitters, y_min_splitters_head) + cy_seps_splitters = np.append(cy_seps_splitters, y_max_splitters_head) - splitter_y_new=[] - splitter_y_new.append(0) - for i in range(len(cy_main_splitters_sort)): - splitter_y_new.append( cy_main_splitters_sort[i] ) - splitter_y_new.append(region_pre_p.shape[0]) - splitter_y_new_diff=np.diff(splitter_y_new)/float(region_pre_p.shape[0])*100 - - args_big_parts=np.arange(len(splitter_y_new_diff))[ splitter_y_new_diff>22 ] + cy_seps_splitters = np.sort(cy_seps_splitters).astype(int) + splitter_y_new = [0] + list(cy_seps_splitters) + [region_pre_p.shape[0]] + big_part = 22 * region_pre_p.shape[0] // 100 # percent height regions_without_separators=return_regions_without_separators(region_pre_p) - length_y_threshold=regions_without_separators.shape[0]/4.0 num_col_fin=0 peaks_neg_fin_fin=[] - for itiles in args_big_parts: - regions_without_separators_tile=regions_without_separators[int(splitter_y_new[itiles]): - int(splitter_y_new[itiles+1]),:] + num_big_parts = 0 + for top, bot in pairwise(splitter_y_new): + if bot - top < big_part: + continue + num_big_parts += 1 try: - num_col, peaks_neg_fin = find_num_col(regions_without_separators_tile, - num_col_classifier, tables, multiplier=7.0) + num_col, peaks_neg_fin = find_num_col(regions_without_separators[top: bot], + num_col_classifier, tables, + vertical_separators=1 * (vertical[top: bot] > 0), + multiplier=7.0) + logger.debug("big part %d:%d has %d columns", top, bot, num_col + 1) + # print(peaks_neg_fin) except: num_col = 0 peaks_neg_fin = [] @@ -1547,576 +1530,402 @@ def find_number_of_columns_in_document(region_pre_p, num_col_classifier, tables, num_col_fin=num_col peaks_neg_fin_fin=peaks_neg_fin - if len(args_big_parts)==1 and (len(peaks_neg_fin_fin)+1)=500] peaks_neg_fin=peaks_neg_fin[peaks_neg_fin<=(vertical.shape[1]-500)] peaks_neg_fin_fin=peaks_neg_fin[:] - return num_col_fin, peaks_neg_fin_fin,matrix_of_lines_ch,splitter_y_new,separators_closeup_n + return num_col_fin, peaks_neg_fin_fin, matrix_of_seps_ch, splitter_y_new, separators_closeup_n def return_boxes_of_images_by_order_of_reading_new( - splitter_y_new, regions_without_separators, - matrix_of_lines_ch, + splitter_y_new, + regions_without_separators, + regions_with_separators, + matrix_of_seps_ch, num_col_classifier, erosion_hurts, tables, right2left_readingorder, logger=None): + """ + Iterate through the vertical parts of a page, each with its own set of columns, + and from the matrix of horizontal separators for that part, find an ordered + list of bounding boxes through all columns and regions. + + Arguments: + * splitter_y_new: the y coordinates separating the parts + * regions_without_separators: (text) region mask with separators suppressed; + (needed to find per-part columns and to combine separators if possible) + * regions_with_separators: (full) region map with separators included; + (needed to elongate separators if possible) + * matrix_of_seps: type and coordinates of horizontal and vertical separators, + as well as headings + * num_col_classifier: predicted number of columns for the entire page + * erosion_hurts: whether region masks have already been eroded + (and thus gaps can be expected to be wider) + * tables: bool + * right2left_readingorder: whether to invert the default left-to-right order + + Returns: a tuple of + * the ordered list of bounding boxes + * a list of arrays: the x coordinates delimiting the columns for every page part + (according to splitter) + """ if right2left_readingorder: regions_without_separators = cv2.flip(regions_without_separators,1) + regions_with_separators = cv2.flip(regions_with_separators,1) if logger is None: logger = getLogger(__package__) logger.debug('enter return_boxes_of_images_by_order_of_reading_new') + # def dbg_imshow(box, title): + # xmin, xmax, ymin, ymax = box + # plt.imshow(regions_with_separators) #, extent=[0, width_tot, bot, top]) + # plt.gca().add_patch(patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, + # fill=False, linewidth=1, edgecolor='r')) + # plt.title(title + " at %d:%d, %d:%d" % (ymin, ymax, xmin, xmax)) + # plt.show() + # def dbg_plt(box=None, title=None, rectangles=None, rectangles_showidx=False): + # minx, maxx, miny, maxy = box or (0, None, 0, None) + # img = regions_without_separators[miny:maxy, minx:maxx] + # plt.imshow(img) + # step = max(img.shape) // 10 + # xrange = np.arange(0, img.shape[1], step) + # yrange = np.arange(0, img.shape[0], step) + # ax = plt.gca() + # ax.set_xticks(xrange) + # ax.set_yticks(yrange) + # ax.set_xticklabels(xrange + minx) + # ax.set_yticklabels(yrange + miny) + # def format_coord(x, y): + # return 'x={:g}, y={:g}'.format(x + minx, y + miny) + # ax.format_coord = format_coord + # if title: + # plt.title(title) + # if rectangles: + # for i, (xmin, xmax, ymin, ymax) in enumerate(rectangles): + # ax.add_patch(patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, + # fill=False, linewidth=1, edgecolor='r')) + # if rectangles_showidx: + # ax.text((xmin+xmax)/2, (ymin+ymax)/2, str(i), c='r') + # plt.show() + # dbg_plt(title="return_boxes_of_images_by_order_of_reading_new") + boxes=[] peaks_neg_tot_tables = [] splitter_y_new = np.array(splitter_y_new, dtype=int) - for i in range(len(splitter_y_new)-1): - #print(splitter_y_new[i],splitter_y_new[i+1]) - matrix_new = matrix_of_lines_ch[:,:][(matrix_of_lines_ch[:,6]> splitter_y_new[i] ) & - (matrix_of_lines_ch[:,7]< splitter_y_new[i+1] )] + height_tot, width_tot = regions_without_separators.shape + big_part = 22 * height_tot // 100 # percent height + _, ccomps, cstats, _ = cv2.connectedComponentsWithStats(regions_without_separators.astype(np.uint8)) + args_ver = matrix_of_seps_ch[:, 9] == 1 + mask_ver = np.zeros_like(regions_without_separators, dtype=bool) + for i in np.flatnonzero(args_ver): + mask_ver[matrix_of_seps_ch[i, 6]: matrix_of_seps_ch[i, 7], + matrix_of_seps_ch[i, 2]: matrix_of_seps_ch[i, 3]] = True + vertical_seps = 1 * ((regions_with_separators == 6) & mask_ver) + for top, bot in pairwise(splitter_y_new): + # print("%d:%d" % (top, bot), 'i') + # dbg_plt([0, None, top, bot], "image cut for y split %d:%d" % (top, bot)) + matrix_new = matrix_of_seps_ch[(matrix_of_seps_ch[:,6] >= top) & + (matrix_of_seps_ch[:,7] < bot)] #print(len( matrix_new[:,9][matrix_new[:,9]==1] )) #print(matrix_new[:,8][matrix_new[:,9]==1],'gaddaaa') # check to see is there any vertical separator to find holes. #if (len(matrix_new[:,9][matrix_new[:,9]==1]) > 0 and # np.max(matrix_new[:,8][matrix_new[:,9]==1]) >= - # 0.1 * (np.abs(splitter_y_new[i+1]-splitter_y_new[i]))): - if True: - try: - num_col, peaks_neg_fin = find_num_col( - regions_without_separators[splitter_y_new[i]:splitter_y_new[i+1], :], - num_col_classifier, tables, multiplier=6. if erosion_hurts else 7.) - except: - peaks_neg_fin=[] - num_col = 0 - try: - if (len(peaks_neg_fin)+1)= big_part else 1, + tables, vertical_separators=vertical_seps[top: bot], + multiplier=6. if erosion_hurts else 7., + unbalanced=True) + try: + if ((len(peaks_neg_fin) + 1 < num_col_classifier or + num_col_classifier == 6) and + # we do not expect to get all columns in small parts (headings etc.): + bot - top >= big_part): + # found too few columns here + #print('burda') + logger.debug("searching for more than %d columns in big part %d:%d", + len(peaks_neg_fin) + 1, top, bot) + peaks_neg_fin_org = np.copy(peaks_neg_fin) + #print("peaks_neg_fin_org", peaks_neg_fin_org) + if len(peaks_neg_fin) == 0: + num_col, peaks_neg_fin = find_num_col( + regions_without_separators[top:bot], + num_col_classifier, tables, + vertical_separators=vertical_seps[top: bot], + # try to be less strict (lower threshold than above) + multiplier=7. if erosion_hurts else 8., + unbalanced=True) + #print(peaks_neg_fin,'peaks_neg_fin') + peaks_neg_fin_early = [0] + peaks_neg_fin + [width_tot-1] + + #print(peaks_neg_fin_early,'burda2') + peaks_neg_fin_rev=[] + for left, right in pairwise(peaks_neg_fin_early): + # print("%d:%d" % (left, right), 'i_n') + # dbg_plt([left, right, top, bot], + # "image cut for y split %d:%d / x gap %d:%d" % ( + # top, bot, left, right)) + # plt.plot(regions_without_separators[top:bot, left:right].sum(axis=0)) + # plt.title("vertical projection (sum over y)") + # plt.show() + # try to get more peaks with different multipliers + num_col_expected = round((right - left) / width_tot * num_col_classifier) + args = regions_without_separators[top:bot, left:right], num_col_expected, tables + kwargs = dict(vertical_separators=vertical_seps[top: bot, left:right]) + _, peaks_neg_fin1 = find_num_col(*args, **kwargs, multiplier=7.) + _, peaks_neg_fin2 = find_num_col(*args, **kwargs, multiplier=5.) + if len(peaks_neg_fin1) >= len(peaks_neg_fin2): + peaks_neg_fin = peaks_neg_fin1 + else: + peaks_neg_fin = peaks_neg_fin2 + # print(peaks_neg_fin) + logger.debug("found %d additional column boundaries in %d:%d", + len(peaks_neg_fin), left, right) + # add offset to local result + peaks_neg_fin = list(np.array(peaks_neg_fin) + left) #print(peaks_neg_fin,'peaks_neg_fin') - for p_n in peaks_neg_fin: - peaks_neg_fin_early.append(p_n) - peaks_neg_fin_early.append(regions_without_separators.shape[1]-1) - #print(peaks_neg_fin_early,'burda2') - peaks_neg_fin_rev=[] - for i_n in range(len(peaks_neg_fin_early)-1): - #print(i_n,'i_n') - #plt.plot(regions_without_separators[splitter_y_new[i]: - # splitter_y_new[i+1], - # peaks_neg_fin_early[i_n]: - # peaks_neg_fin_early[i_n+1]].sum(axis=0) ) - #plt.show() - try: - num_col, peaks_neg_fin1 = find_num_col( - regions_without_separators[splitter_y_new[i]:splitter_y_new[i+1], - peaks_neg_fin_early[i_n]:peaks_neg_fin_early[i_n+1]], - num_col_classifier,tables, multiplier=7.) - except: - peaks_neg_fin1=[] - try: - num_col, peaks_neg_fin2 = find_num_col( - regions_without_separators[splitter_y_new[i]:splitter_y_new[i+1], - peaks_neg_fin_early[i_n]:peaks_neg_fin_early[i_n+1]], - num_col_classifier,tables, multiplier=5.) - except: - peaks_neg_fin2=[] + peaks_neg_fin_rev.extend(peaks_neg_fin) + if right < peaks_neg_fin_early[-1]: + # all but the last column: interject the preexisting boundary + peaks_neg_fin_rev.append(right) + #print(peaks_neg_fin_rev,'peaks_neg_fin_rev') - if len(peaks_neg_fin1)>=len(peaks_neg_fin2): - peaks_neg_fin=list(np.copy(peaks_neg_fin1)) - else: - peaks_neg_fin=list(np.copy(peaks_neg_fin2)) - peaks_neg_fin=list(np.array(peaks_neg_fin)+peaks_neg_fin_early[i_n]) - - if i_n!=(len(peaks_neg_fin_early)-2): - peaks_neg_fin_rev.append(peaks_neg_fin_early[i_n+1]) - #print(peaks_neg_fin,'peaks_neg_fin') - peaks_neg_fin_rev=peaks_neg_fin_rev+peaks_neg_fin - - if len(peaks_neg_fin_rev)>=len(peaks_neg_fin_org): - peaks_neg_fin=list(np.sort(peaks_neg_fin_rev)) - num_col=len(peaks_neg_fin) - else: - peaks_neg_fin=list(np.copy(peaks_neg_fin_org)) - num_col=len(peaks_neg_fin) - - #print(peaks_neg_fin,'peaks_neg_fin') - except: - logger.exception("cannot find peaks consistent with columns") - #num_col, peaks_neg_fin = find_num_col( - # regions_without_separators[splitter_y_new[i]:splitter_y_new[i+1],:], - # multiplier=7.0) - x_min_hor_some=matrix_new[:,2][ (matrix_new[:,9]==0) ] - x_max_hor_some=matrix_new[:,3][ (matrix_new[:,9]==0) ] - cy_hor_some=matrix_new[:,5][ (matrix_new[:,9]==0) ] - cy_hor_diff=matrix_new[:,7][ (matrix_new[:,9]==0) ] - arg_org_hor_some=matrix_new[:,0][ (matrix_new[:,9]==0) ] - - if right2left_readingorder: - x_max_hor_some_new = regions_without_separators.shape[1] - x_min_hor_some - x_min_hor_some_new = regions_without_separators.shape[1] - x_max_hor_some - x_min_hor_some =list(np.copy(x_min_hor_some_new)) - x_max_hor_some =list(np.copy(x_max_hor_some_new)) - - peaks_neg_tot=return_points_with_boundies(peaks_neg_fin,0, regions_without_separators[:,:].shape[1]) - peaks_neg_tot_tables.append(peaks_neg_tot) - - reading_order_type, x_starting, x_ending, y_type_2, y_diff_type_2, \ - y_lines_without_mother, x_start_without_mother, x_end_without_mother, there_is_sep_with_child, \ - y_lines_with_child_without_mother, x_start_with_child_without_mother, x_end_with_child_without_mother, \ - new_main_sep_y = return_x_start_end_mothers_childs_and_type_of_reading_order( - x_min_hor_some, x_max_hor_some, cy_hor_some, peaks_neg_tot, cy_hor_diff) - - all_columns = set(range(len(peaks_neg_tot) - 1)) - if ((reading_order_type==1) or - (reading_order_type==0 and - (len(y_lines_without_mother)>=2 or there_is_sep_with_child==1))): - try: - y_grenze = splitter_y_new[i] + 300 - #check if there is a big separator in this y_mains_sep_ohne_grenzen - - args_early_ys=np.arange(len(y_type_2)) - #print(args_early_ys,'args_early_ys') - #print(splitter_y_new[i], splitter_y_new[i+1]) - - x_starting_up = x_starting[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - x_ending_up = x_ending[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - y_type_2_up = y_type_2[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - y_diff_type_2_up = y_diff_type_2[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - args_up = args_early_ys[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - if len(y_type_2_up) > 0: - y_main_separator_up = y_type_2_up [(x_starting_up==0) & - (x_ending_up==(len(peaks_neg_tot)-1) )] - y_diff_main_separator_up = y_diff_type_2_up[(x_starting_up==0) & - (x_ending_up==(len(peaks_neg_tot)-1) )] - args_main_to_deleted = args_up[(x_starting_up==0) & - (x_ending_up==(len(peaks_neg_tot)-1) )] - #print(y_main_separator_up,y_diff_main_separator_up,args_main_to_deleted,'fffffjammmm') - if len(y_diff_main_separator_up) > 0: - args_to_be_kept = np.array(list( set(args_early_ys) - set(args_main_to_deleted) )) - #print(args_to_be_kept,'args_to_be_kept') - boxes.append([0, peaks_neg_tot[len(peaks_neg_tot)-1], - splitter_y_new[i], y_diff_main_separator_up.max()]) - splitter_y_new[i] = y_diff_main_separator_up.max() - - #print(splitter_y_new[i],'splitter_y_new[i]') - y_type_2 = y_type_2[args_to_be_kept] - x_starting = x_starting[args_to_be_kept] - x_ending = x_ending[args_to_be_kept] - y_diff_type_2 = y_diff_type_2[args_to_be_kept] - - #print('galdiha') - y_grenze = splitter_y_new[i] + 200 - args_early_ys2=np.arange(len(y_type_2)) - y_type_2_up=y_type_2[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - x_starting_up=x_starting[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - x_ending_up=x_ending[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - y_diff_type_2_up=y_diff_type_2[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - args_up2=args_early_ys2[(y_type_2 > splitter_y_new[i]) & - (y_type_2 <= y_grenze)] - #print(y_type_2_up,x_starting_up,x_ending_up,'didid') - nodes_in = set() - for ij in range(len(x_starting_up)): - nodes_in.update(range(x_starting_up[ij], - x_ending_up[ij])) - #print(nodes_in,'nodes_in') - - if nodes_in == set(range(len(peaks_neg_tot)-1)): - pass - elif nodes_in == set(range(1, len(peaks_neg_tot)-1)): - pass - else: - #print('burdaydikh') - args_to_be_kept2=np.array(list( set(args_early_ys2)-set(args_up2) )) - - if len(args_to_be_kept2)>0: - y_type_2 = y_type_2[args_to_be_kept2] - x_starting = x_starting[args_to_be_kept2] - x_ending = x_ending[args_to_be_kept2] - y_diff_type_2 = y_diff_type_2[args_to_be_kept2] - else: - pass - #print('burdaydikh2') - elif len(y_diff_main_separator_up)==0: - nodes_in = set() - for ij in range(len(x_starting_up)): - nodes_in.update(range(x_starting_up[ij], - x_ending_up[ij])) - #print(nodes_in,'nodes_in2') - #print(np.array(range(len(peaks_neg_tot)-1)),'np.array(range(len(peaks_neg_tot)-1))') - - if nodes_in == set(range(len(peaks_neg_tot)-1)): - pass - elif nodes_in == set(range(1,len(peaks_neg_tot)-1)): - pass - else: - #print('burdaydikh') - #print(args_early_ys,'args_early_ys') - #print(args_up,'args_up') - args_to_be_kept2=np.array(list( set(args_early_ys) - set(args_up) )) - - #print(args_to_be_kept2,'args_to_be_kept2') - #print(len(y_type_2),len(x_starting),len(x_ending),len(y_diff_type_2)) - if len(args_to_be_kept2)>0: - y_type_2 = y_type_2[args_to_be_kept2] - x_starting = x_starting[args_to_be_kept2] - x_ending = x_ending[args_to_be_kept2] - y_diff_type_2 = y_diff_type_2[args_to_be_kept2] - else: - pass - #print('burdaydikh2') - - #int(splitter_y_new[i]) - y_lines_by_order=[] - x_start_by_order=[] - x_end_by_order=[] - if (len(x_end_with_child_without_mother)==0 and reading_order_type==0) or reading_order_type==1: - if reading_order_type==1: - y_lines_by_order.append(splitter_y_new[i]) - x_start_by_order.append(0) - x_end_by_order.append(len(peaks_neg_tot)-2) - else: - #print(x_start_without_mother,x_end_without_mother,peaks_neg_tot,'dodo') - columns_covered_by_mothers = set() - for dj in range(len(x_start_without_mother)): - columns_covered_by_mothers.update( - range(x_start_without_mother[dj], - x_end_without_mother[dj])) - columns_not_covered = list(all_columns - columns_covered_by_mothers) - y_type_2 = np.append(y_type_2, np.ones(len(columns_not_covered) + - len(x_start_without_mother), - dtype=int) * splitter_y_new[i]) - ##y_lines_by_order = np.append(y_lines_by_order, [splitter_y_new[i]] * len(columns_not_covered)) - ##x_start_by_order = np.append(x_start_by_order, [0] * len(columns_not_covered)) - x_starting = np.append(x_starting, np.array(columns_not_covered, int)) - x_starting = np.append(x_starting, x_start_without_mother) - x_ending = np.append(x_ending, np.array(columns_not_covered, int) + 1) - x_ending = np.append(x_ending, x_end_without_mother) - - ind_args=np.arange(len(y_type_2)) - #ind_args=np.array(ind_args) - #print(ind_args,'ind_args') - for column in range(len(peaks_neg_tot)-1): - #print(column,'column') - ind_args_in_col=ind_args[x_starting==column] - #print('babali2') - #print(ind_args_in_col,'ind_args_in_col') - ind_args_in_col=np.array(ind_args_in_col) - #print(len(y_type_2)) - y_column=y_type_2[ind_args_in_col] - x_start_column=x_starting[ind_args_in_col] - x_end_column=x_ending[ind_args_in_col] - #print('babali3') - ind_args_col_sorted=np.argsort(y_column) - y_col_sort=y_column[ind_args_col_sorted] - x_start_column_sort=x_start_column[ind_args_col_sorted] - x_end_column_sort=x_end_column[ind_args_col_sorted] - #print('babali4') - for ii in range(len(y_col_sort)): - #print('babali5') - y_lines_by_order.append(y_col_sort[ii]) - x_start_by_order.append(x_start_column_sort[ii]) - x_end_by_order.append(x_end_column_sort[ii]-1) - else: - #print(x_start_without_mother,x_end_without_mother,peaks_neg_tot,'dodo') - columns_covered_by_mothers = set() - for dj in range(len(x_start_without_mother)): - columns_covered_by_mothers.update( - range(x_start_without_mother[dj], - x_end_without_mother[dj])) - columns_not_covered = list(all_columns - columns_covered_by_mothers) - y_type_2 = np.append(y_type_2, np.ones(len(columns_not_covered) + len(x_start_without_mother), - dtype=int) * splitter_y_new[i]) - ##y_lines_by_order = np.append(y_lines_by_order, [splitter_y_new[i]] * len(columns_not_covered)) - ##x_start_by_order = np.append(x_start_by_order, [0] * len(columns_not_covered)) - x_starting = np.append(x_starting, np.array(columns_not_covered, int)) - x_starting = np.append(x_starting, x_start_without_mother) - x_ending = np.append(x_ending, np.array(columns_not_covered, int) + 1) - x_ending = np.append(x_ending, x_end_without_mother) - - columns_covered_by_with_child_no_mothers = set() - for dj in range(len(x_end_with_child_without_mother)): - columns_covered_by_with_child_no_mothers.update( - range(x_start_with_child_without_mother[dj], - x_end_with_child_without_mother[dj])) - columns_not_covered_child_no_mother = list( - all_columns - columns_covered_by_with_child_no_mothers) - #indexes_to_be_spanned=[] - for i_s in range(len(x_end_with_child_without_mother)): - columns_not_covered_child_no_mother.append(x_start_with_child_without_mother[i_s]) - columns_not_covered_child_no_mother = np.sort(columns_not_covered_child_no_mother) - ind_args = np.arange(len(y_type_2)) - x_end_with_child_without_mother = np.array(x_end_with_child_without_mother, int) - x_start_with_child_without_mother = np.array(x_start_with_child_without_mother, int) - for i_s_nc in columns_not_covered_child_no_mother: - if i_s_nc in x_start_with_child_without_mother: - x_end_biggest_column = \ - x_end_with_child_without_mother[x_start_with_child_without_mother==i_s_nc][0] - args_all_biggest_lines = ind_args[(x_starting==i_s_nc) & - (x_ending==x_end_biggest_column)] - y_column_nc = y_type_2[args_all_biggest_lines] - x_start_column_nc = x_starting[args_all_biggest_lines] - x_end_column_nc = x_ending[args_all_biggest_lines] - y_column_nc = np.sort(y_column_nc) - for i_c in range(len(y_column_nc)): - if i_c==(len(y_column_nc)-1): - ind_all_lines_between_nm_wc=ind_args[(y_type_2>y_column_nc[i_c]) & - (y_type_2=i_s_nc) & - (x_ending<=x_end_biggest_column)] - else: - ind_all_lines_between_nm_wc=ind_args[(y_type_2>y_column_nc[i_c]) & - (y_type_2=i_s_nc) & - (x_ending<=x_end_biggest_column)] - y_all_between_nm_wc = y_type_2[ind_all_lines_between_nm_wc] - x_starting_all_between_nm_wc = x_starting[ind_all_lines_between_nm_wc] - x_ending_all_between_nm_wc = x_ending[ind_all_lines_between_nm_wc] - - x_diff_all_between_nm_wc = x_ending_all_between_nm_wc - x_starting_all_between_nm_wc - if len(x_diff_all_between_nm_wc)>0: - biggest=np.argmax(x_diff_all_between_nm_wc) - - columns_covered_by_mothers = set() - for dj in range(len(x_starting_all_between_nm_wc)): - columns_covered_by_mothers.update( - range(x_starting_all_between_nm_wc[dj], - x_ending_all_between_nm_wc[dj])) - child_columns = set(range(i_s_nc, x_end_biggest_column)) - columns_not_covered = list(child_columns - columns_covered_by_mothers) - - should_longest_line_be_extended=0 - if (len(x_diff_all_between_nm_wc) > 0 and - set(list(range(x_starting_all_between_nm_wc[biggest], - x_ending_all_between_nm_wc[biggest])) + - list(columns_not_covered)) != child_columns): - should_longest_line_be_extended=1 - index_lines_so_close_to_top_separator = \ - np.arange(len(y_all_between_nm_wc))[(y_all_between_nm_wc>y_column_nc[i_c]) & - (y_all_between_nm_wc<=(y_column_nc[i_c]+500))] - if len(index_lines_so_close_to_top_separator) > 0: - indexes_remained_after_deleting_closed_lines= \ - np.array(list(set(list(range(len(y_all_between_nm_wc)))) - - set(list(index_lines_so_close_to_top_separator)))) - if len(indexes_remained_after_deleting_closed_lines) > 0: - y_all_between_nm_wc = \ - y_all_between_nm_wc[indexes_remained_after_deleting_closed_lines] - x_starting_all_between_nm_wc = \ - x_starting_all_between_nm_wc[indexes_remained_after_deleting_closed_lines] - x_ending_all_between_nm_wc = \ - x_ending_all_between_nm_wc[indexes_remained_after_deleting_closed_lines] - - y_all_between_nm_wc = np.append(y_all_between_nm_wc, y_column_nc[i_c]) - x_starting_all_between_nm_wc = np.append(x_starting_all_between_nm_wc, i_s_nc) - x_ending_all_between_nm_wc = np.append(x_ending_all_between_nm_wc, x_end_biggest_column) - - if len(x_diff_all_between_nm_wc) > 0: - try: - y_all_between_nm_wc = np.append(y_all_between_nm_wc, y_column_nc[i_c]) - x_starting_all_between_nm_wc = np.append(x_starting_all_between_nm_wc, x_starting_all_between_nm_wc[biggest]) - x_ending_all_between_nm_wc = np.append(x_ending_all_between_nm_wc, x_ending_all_between_nm_wc[biggest]) - except: - logger.exception("cannot append") - - y_all_between_nm_wc = np.append(y_all_between_nm_wc, [y_column_nc[i_c]] * len(columns_not_covered)) - x_starting_all_between_nm_wc = np.append(x_starting_all_between_nm_wc, np.array(columns_not_covered, int)) - x_ending_all_between_nm_wc = np.append(x_ending_all_between_nm_wc, np.array(columns_not_covered, int) + 1) - - ind_args_between=np.arange(len(x_ending_all_between_nm_wc)) - for column in range(int(i_s_nc), int(x_end_biggest_column)): - ind_args_in_col=ind_args_between[x_starting_all_between_nm_wc==column] - #print('babali2') - #print(ind_args_in_col,'ind_args_in_col') - ind_args_in_col=np.array(ind_args_in_col) - #print(len(y_type_2)) - y_column=y_all_between_nm_wc[ind_args_in_col] - x_start_column=x_starting_all_between_nm_wc[ind_args_in_col] - x_end_column=x_ending_all_between_nm_wc[ind_args_in_col] - #print('babali3') - ind_args_col_sorted=np.argsort(y_column) - y_col_sort=y_column[ind_args_col_sorted] - x_start_column_sort=x_start_column[ind_args_col_sorted] - x_end_column_sort=x_end_column[ind_args_col_sorted] - #print('babali4') - for ii in range(len(y_col_sort)): - #print('babali5') - y_lines_by_order.append(y_col_sort[ii]) - x_start_by_order.append(x_start_column_sort[ii]) - x_end_by_order.append(x_end_column_sort[ii]-1) - else: - #print(column,'column') - ind_args_in_col=ind_args[x_starting==i_s_nc] - #print('babali2') - #print(ind_args_in_col,'ind_args_in_col') - ind_args_in_col=np.array(ind_args_in_col) - #print(len(y_type_2)) - y_column=y_type_2[ind_args_in_col] - x_start_column=x_starting[ind_args_in_col] - x_end_column=x_ending[ind_args_in_col] - #print('babali3') - ind_args_col_sorted=np.argsort(y_column) - y_col_sort=y_column[ind_args_col_sorted] - x_start_column_sort=x_start_column[ind_args_col_sorted] - x_end_column_sort=x_end_column[ind_args_col_sorted] - #print('babali4') - for ii in range(len(y_col_sort)): - y_lines_by_order.append(y_col_sort[ii]) - x_start_by_order.append(x_start_column_sort[ii]) - x_end_by_order.append(x_end_column_sort[ii]-1) - - for il in range(len(y_lines_by_order)): - y_copy = list(y_lines_by_order) - x_start_copy = list(x_start_by_order) - x_end_copy = list(x_end_by_order) - - #print(y_copy,'y_copy') - y_itself=y_copy.pop(il) - x_start_itself=x_start_copy.pop(il) - x_end_itself=x_end_copy.pop(il) - - #print(y_copy,'y_copy2') - for column in range(int(x_start_itself), int(x_end_itself)+1): - #print(column,'cols') - y_in_cols=[] - for yic in range(len(y_copy)): - #print('burda') - if (y_copy[yic]>y_itself and - column>=x_start_copy[yic] and - column<=x_end_copy[yic]): - y_in_cols.append(y_copy[yic]) - #print('burda2') - #print(y_in_cols,'y_in_cols') - if len(y_in_cols)>0: - y_down=np.min(y_in_cols) - else: - y_down=splitter_y_new[i+1] - #print(y_itself,'y_itself') - boxes.append([peaks_neg_tot[column], - peaks_neg_tot[column+1], - y_itself, - y_down]) - except: - logger.exception("cannot assign boxes") - boxes.append([0, peaks_neg_tot[len(peaks_neg_tot)-1], - splitter_y_new[i], splitter_y_new[i+1]]) - else: - y_lines_by_order=[] - x_start_by_order=[] - x_end_by_order=[] - if len(x_starting)>0: - columns_covered_by_lines_covered_more_than_2col = set() - for dj in range(len(x_starting)): - if set(range(x_starting[dj], x_ending[dj])) != all_columns: - columns_covered_by_lines_covered_more_than_2col.update( - range(x_starting[dj], x_ending[dj])) - columns_not_covered = list(all_columns - columns_covered_by_lines_covered_more_than_2col) - - y_type_2 = np.append(y_type_2, np.ones(len(columns_not_covered) + 1, - dtype=int) * splitter_y_new[i]) - ##y_lines_by_order = np.append(y_lines_by_order, [splitter_y_new[i]] * len(columns_not_covered)) - ##x_start_by_order = np.append(x_start_by_order, [0] * len(columns_not_covered)) - x_starting = np.append(x_starting, np.array(columns_not_covered, x_starting.dtype)) - x_ending = np.append(x_ending, np.array(columns_not_covered, x_ending.dtype) + 1) - if len(new_main_sep_y) > 0: - x_starting = np.append(x_starting, 0) - x_ending = np.append(x_ending, len(peaks_neg_tot) - 1) - else: - x_starting = np.append(x_starting, x_starting[0]) - x_ending = np.append(x_ending, x_ending[0]) + if len(peaks_neg_fin_rev) >= len(peaks_neg_fin_org): + #print("found more peaks than at first glance", peaks_neg_fin_rev, peaks_neg_fin_org) + peaks_neg_fin = peaks_neg_fin_rev else: - columns_not_covered = list(all_columns) - y_type_2 = np.append(y_type_2, np.ones(len(columns_not_covered), - dtype=int) * splitter_y_new[i]) - ##y_lines_by_order = np.append(y_lines_by_order, [splitter_y_new[i]] * len(columns_not_covered)) - ##x_start_by_order = np.append(x_start_by_order, [0] * len(columns_not_covered)) - x_starting = np.append(x_starting, np.array(columns_not_covered, x_starting.dtype)) - x_ending = np.append(x_ending, np.array(columns_not_covered, x_ending.dtype) + 1) + peaks_neg_fin = peaks_neg_fin_org + num_col = len(peaks_neg_fin) + #print(peaks_neg_fin,'peaks_neg_fin') + except: + logger.exception("cannot find peaks consistent with columns") + #num_col, peaks_neg_fin = find_num_col( + # regions_without_separators[top:bot,:], + # multiplier=7.0) + peaks_neg_tot = np.array([0] + peaks_neg_fin + [width_tot]) + #print(peaks_neg_tot,'peaks_neg_tot') + peaks_neg_tot_tables.append(peaks_neg_tot) - ind_args = np.arange(len(y_type_2)) - - for column in range(len(peaks_neg_tot)-1): - #print(column,'column') - ind_args_in_col=ind_args[x_starting==column] - ind_args_in_col=np.array(ind_args_in_col) - #print(len(y_type_2)) - y_column=y_type_2[ind_args_in_col] - x_start_column=x_starting[ind_args_in_col] - x_end_column=x_ending[ind_args_in_col] + all_columns = set(range(len(peaks_neg_tot) - 1)) + #print("all_columns", all_columns) - ind_args_col_sorted=np.argsort(y_column) - y_col_sort=y_column[ind_args_col_sorted] - x_start_column_sort=x_start_column[ind_args_col_sorted] - x_end_column_sort=x_end_column[ind_args_col_sorted] - #print('babali4') - for ii in range(len(y_col_sort)): - #print('babali5') - y_lines_by_order.append(y_col_sort[ii]) - x_start_by_order.append(x_start_column_sort[ii]) - x_end_by_order.append(x_end_column_sort[ii]-1) + # elongate horizontal separators+headings as much as possible without overlap + args_nonver = matrix_new[:, 9] != 1 + for i in np.flatnonzero(args_nonver): + xmin, xmax, ymin, ymax, typ = matrix_new[i, [2, 3, 6, 7, 9]] + cut = regions_with_separators[ymin: ymax] + # dbg_imshow([xmin, xmax, ymin, ymax], "separator %d (%s)" % (i, "heading" if typ else "horizontal")) + starting = xmin - peaks_neg_tot + min_start = np.flatnonzero(starting >= 0)[-1] # last left-of + ending = xmax - peaks_neg_tot + max_end = np.flatnonzero(ending <= 0)[0] # first right-of + # skip elongation unless this is already a multi-column separator/heading: + if not max_end - min_start > 1: + continue + # is there anything left of min_start? + for j in range(min_start): + # dbg_imshow([peaks_neg_tot[j], xmin, ymin, ymax], "start of %d candidate %d" % (i, j)) + if not np.any(cut[:, peaks_neg_tot[j]: xmin]): + # print("elongated sep", i, "typ", typ, "start", xmin, "to", j, peaks_neg_tot[j]) + matrix_new[i, 2] = peaks_neg_tot[j] + 1 # elongate to start of this column + break + # is there anything right of max_end? + for j in range(len(peaks_neg_tot) - 1, max_end, -1): + # dbg_imshow([xmax, peaks_neg_tot[j], ymin, ymax], "end of %d candidate %d" % (i, j)) + if not np.any(cut[:, xmax: peaks_neg_tot[j]]): + # print("elongated sep", i, "typ", typ, "end", xmax, "to", j, peaks_neg_tot[j]) + matrix_new[i, 3] = peaks_neg_tot[j] - 1 # elongate to end of this column + break - for il in range(len(y_lines_by_order)): - y_copy = list(y_lines_by_order) - x_start_copy = list(x_start_by_order) - x_end_copy = list(x_end_by_order) + args_hor = matrix_new[:, 9] == 0 + x_min_hor_some = matrix_new[:, 2][args_hor] + x_max_hor_some = matrix_new[:, 3][args_hor] + y_min_hor_some = matrix_new[:, 6][args_hor] + y_max_hor_some = matrix_new[:, 7][args_hor] + cy_hor_some = matrix_new[:, 5][args_hor] - #print(y_copy,'y_copy') - y_itself=y_copy.pop(il) - x_start_itself=x_start_copy.pop(il) - x_end_itself=x_end_copy.pop(il) + args_head = matrix_new[:, 9] == 2 + x_min_hor_head = matrix_new[:, 2][args_head] + x_max_hor_head = matrix_new[:, 3][args_head] + y_min_hor_head = matrix_new[:, 6][args_head] + y_max_hor_head = matrix_new[:, 7][args_head] + cy_hor_head = matrix_new[:, 5][args_head] - for column in range(x_start_itself, x_end_itself+1): - #print(column,'cols') - y_in_cols=[] - for yic in range(len(y_copy)): - #print('burda') - if (y_copy[yic]>y_itself and - column>=x_start_copy[yic] and - column<=x_end_copy[yic]): - y_in_cols.append(y_copy[yic]) - #print('burda2') - #print(y_in_cols,'y_in_cols') - if len(y_in_cols)>0: - y_down=np.min(y_in_cols) - else: - y_down=splitter_y_new[i+1] - #print(y_itself,'y_itself') + # split headings at toplines (y_min_head) and baselines (y_max_head) + # instead of merely adding their center (cy_head) as horizontal separator + # (x +/- 30px to avoid crossing col peaks by accident) + x_min_hor_some = np.append(x_min_hor_some, np.tile(x_min_hor_head + 30, 2)) + x_max_hor_some = np.append(x_max_hor_some, np.tile(x_max_hor_head - 30, 2)) + y_min_hor_some = np.append(y_min_hor_some, # toplines + np.concatenate((y_min_hor_head - 2, + y_max_hor_head - 0))) + y_max_hor_some = np.append(y_max_hor_some, # baselines + np.concatenate((y_min_hor_head + 0, + y_max_hor_head + 2))) + cy_hor_some = np.append(cy_hor_some, # centerlines + np.concatenate((y_min_hor_head - 1, + y_max_hor_head + 1))) + + # analyse connected components of regions to gain additional separators + # and prepare a map for cross-column boxes + ccounts = np.bincount(ccomps[top: bot].flatten()) + ccounts_median = np.median(ccounts) + col_ccounts = np.stack([np.bincount(ccomps[top: bot, left: right].flatten(), + minlength=ccounts.size) + for left, right in pairwise(peaks_neg_tot)]) + labelcolmap = dict() + for label, label_count in enumerate(ccounts): + if not label: + continue + # ignore small labels for the purpose of finding multicol seps + if label_count < 0.5 * ccounts_median: + continue + label_left, label_top, label_width, label_height, label_area = cstats[label] + # if label_count < 0.9 * label_area: + # # mostly not in this part of the page + # continue + if label_count < 0.01 * (top - bot) * width_tot: + continue + #assert np.sum(col_ccounts[:, label]) == label_count + label_right = label_left + label_width + label_bot = label_top + label_height + label_start = np.flatnonzero(peaks_neg_tot > label_left)[0] - 1 + label_end = np.flatnonzero(peaks_neg_tot >= label_right)[0] + if label_end - label_start < 2: + continue + if np.count_nonzero(col_ccounts[:, label] > 0.1 * label_count) < 2: + continue + # store as dict for multi-column boxes: + for start in range(label_start, label_end): + labelcolmap.setdefault(start, list()).append( + (label_end, label_top, label_bot, sum(col_ccounts[start: label_end, label]))) + # make additional separators: + x_min_hor_some = np.append(x_min_hor_some, [label_left] * 2) + x_max_hor_some = np.append(x_max_hor_some, [label_right] * 2) + y_min_hor_some = np.append(y_min_hor_some, [label_top - 2, label_bot]) + y_max_hor_some = np.append(y_max_hor_some, [label_top, label_bot + 2]) + cy_hor_some = np.append(cy_hor_some, [label_top - 1, label_bot + 1]) + + if right2left_readingorder: + x_max_hor_some = width_tot - x_min_hor_some + x_min_hor_some = width_tot - x_max_hor_some + + x_starting, x_ending, y_min, y_mid, y_max = return_multicol_separators_x_start_end( + regions_without_separators, peaks_neg_tot, top, bot, + x_min_hor_some, x_max_hor_some, cy_hor_some, y_min_hor_some, y_max_hor_some) + # dbg_plt([0, None, top, bot], "non-empty multi-column separators in current split", + # list(zip(peaks_neg_tot[x_starting], peaks_neg_tot[x_ending], + # y_min - top, y_max - top)), True) + + # core algorithm: + # 1. iterate through multi-column separators, pre-ordered by their y coord + # 2. for each separator, iterate from its starting to its ending column + # 3. in each starting column, determine the next downwards separator, + # 4. if there is none, then fill up the column to the bottom; + # otherwise, fill up to that next separator + # 5. moreover, determine the next rightward column that would not cut through + # any regions, advancing to that column, and storing a new in-order bbox + # for that down/right span + # 6. if there was a next separator, and it ends no further than the current one, + # then recurse on that separator from step 1, then continue (with the next + # column for the current separator) at step 2, or (with the next separator + # in order) at step 1 + args = list(range(len(y_mid))) + while len(args): + cur = args[0] + args = args[1:] + # print("iter", cur, y_mid[cur], "%d:%d" % (x_starting[cur], x_ending[cur])) + def get_span(start, y_top, y_bot): + # for last, l_top, l_bot, l_count in labelcolmap.get(start, []): + # if y_top < l_bot and y_bot > l_top and last > start + 1: + # width = (peaks_neg_tot[last] - peaks_neg_tot[start]) + # print("span", start, last, l_top, l_bot, l_count, + # "box area", (y_bot - y_top) * width, + # "label area", (min(y_bot, l_bot) - max(y_top, l_top)) * width, + # "box height", (y_bot - y_top), + # "label height", sum(regions_without_separators[ + # y_top: y_bot, peaks_neg_tot[start + 1]])) + return max((last for last, l_top, l_bot, l_count in labelcolmap.get(start, []) + # yield the right-most column that does not cut through + # any regions in this horizontal span + if y_top < l_bot and y_bot > l_top + # Ignore if it ends here, anyway + and last > start + 1 + # Ensure this is not just a tiny region near larger regions + and l_count > 0.1 * max(l_count2 for _, l_top2, l_bot2, l_count2 in labelcolmap[start] + if y_top < l_bot2 and y_bot > l_top2) + # or just a small cut of the respective region + # (i.e. box should cover at least 10% of the label). + and ((min(y_bot, l_bot) - max(y_top, l_top)) * + (peaks_neg_tot[last] - peaks_neg_tot[start])) > 0.1 * l_count + # But do allow cutting tiny passages with less 10% of height + # (i.e. label is already almost separated by columns) + and sum(regions_without_separators[ + y_top: y_bot, peaks_neg_tot[start + 1]]) > 0.1 * (y_bot - y_top)), + # Otherwise advance only 1 column. + default=start + 1) + def add_sep(cur): + column = x_starting[cur] + while column < x_ending[cur]: + nxt = np.flatnonzero((y_mid[cur] < y_mid) & + (column >= x_starting) & + (column < x_ending)) + if len(nxt): + nxt = nxt[0] + # print("column", column) + last = get_span(column, y_max[cur], y_min[nxt]) + last = min(last, x_ending[nxt], x_ending[cur]) + # print("nxt", nxt, y_mid[nxt], "%d:%d" % (column, last)) boxes.append([peaks_neg_tot[column], - peaks_neg_tot[column+1], - y_itself, - y_down]) - #else: - #boxes.append([ 0, regions_without_separators[:,:].shape[1] ,splitter_y_new[i],splitter_y_new[i+1]]) + peaks_neg_tot[last], + y_mid[cur], + y_mid[nxt]]) + # dbg_plt(boxes[-1], "recursive column %d:%d box [%d]" % (column, last, len(boxes))) + column = last + if (last == x_ending[nxt] and + x_ending[nxt] <= x_ending[cur] and + x_starting[nxt] >= x_starting[cur] and + nxt in args): + # child – recur + # print("recur", nxt, y_mid[nxt], "%d:%d" % (x_starting[nxt], x_ending[nxt])) + args.remove(nxt) + add_sep(nxt) + else: + # print("column", column) + last = get_span(column, y_max[cur], bot) + # print("bot", bot, "%d:%d" % (column, last)) + boxes.append([peaks_neg_tot[column], + peaks_neg_tot[last], + y_mid[cur], + bot]) + # dbg_plt(boxes[-1], "non-recursive column %d box [%d]" % (column, len(boxes))) + column = last + add_sep(cur) if right2left_readingorder: peaks_neg_tot_tables_new = [] if len(peaks_neg_tot_tables)>=1: for peaks_tab_ind in peaks_neg_tot_tables: - peaks_neg_tot_tables_ind = regions_without_separators.shape[1] - np.array(peaks_tab_ind) + peaks_neg_tot_tables_ind = width_tot - np.array(peaks_tab_ind) peaks_neg_tot_tables_ind = list(peaks_neg_tot_tables_ind[::-1]) peaks_neg_tot_tables_new.append(peaks_neg_tot_tables_ind) for i in range(len(boxes)): - x_start_new = regions_without_separators.shape[1] - boxes[i][1] - x_end_new = regions_without_separators.shape[1] - boxes[i][0] + x_start_new = width_tot - boxes[i][1] + x_end_new = width_tot - boxes[i][0] boxes[i][0] = x_start_new boxes[i][1] = x_end_new peaks_neg_tot_tables = peaks_neg_tot_tables_new + # show final xy-cut + # dbg_plt(None, "final XY-Cut", boxes, True) + logger.debug('exit return_boxes_of_images_by_order_of_reading_new') return boxes, peaks_neg_tot_tables @@ -2130,3 +1939,11 @@ def is_image_filename(fname: str) -> bool: def is_xml_filename(fname: str) -> bool: return fname.lower().endswith('.xml') + +def ensure_array(obj: Iterable) -> np.ndarray: + """convert sequence to array of type `object` so items can be of heterogeneous shape + (but ensure not to convert inner arrays to `object` if len=1) + """ + if not isinstance(obj, np.ndarray): + return np.fromiter(obj, object) + return obj diff --git a/src/eynollah/utils/contour.py b/src/eynollah/utils/contour.py index 6550171..3a67c65 100644 --- a/src/eynollah/utils/contour.py +++ b/src/eynollah/utils/contour.py @@ -14,21 +14,16 @@ from shapely.ops import unary_union, nearest_points from .rotate import rotate_image, rotation_image_new def contours_in_same_horizon(cy_main_hor): - X1 = np.zeros((len(cy_main_hor), len(cy_main_hor))) - X2 = np.zeros((len(cy_main_hor), len(cy_main_hor))) - - X1[0::1, :] = cy_main_hor[:] - X2 = X1.T - - X_dif = np.abs(X2 - X1) - args_help = np.array(range(len(cy_main_hor))) - all_args = [] - for i in range(len(cy_main_hor)): - list_h = list(args_help[X_dif[i, :] <= 20]) - list_h.append(i) - if len(list_h) > 1: - all_args.append(list(set(list_h))) - return np.unique(np.array(all_args, dtype=object)) + """ + Takes an array of y coords, identifies all pairs among them + which are close to each other, and returns all such pairs + by index into the array. + """ + sort = np.argsort(cy_main_hor) + same = np.diff(cy_main_hor[sort]) <= 20 + # groups = np.split(sort, np.arange(len(cy_main_hor) - 1)[~same] + 1) + same = np.flatnonzero(same) + return np.stack((sort[:-1][same], sort[1:][same])).T def find_contours_mean_y_diff(contours_main): M_main = [cv2.moments(contours_main[j]) for j in range(len(contours_main))] @@ -253,13 +248,17 @@ def return_contours_of_image(image): return contours, hierarchy def dilate_textline_contours(all_found_textline_polygons): - return [[polygon2contour(contour2polygon(contour, dilate=6)) - for contour in region] + from . import ensure_array + return [ensure_array( + [polygon2contour(contour2polygon(contour, dilate=6)) + for contour in region]) for region in all_found_textline_polygons] -def dilate_textregion_contours(all_found_textline_polygons): - return [polygon2contour(contour2polygon(contour, dilate=6)) - for contour in all_found_textline_polygons] +def dilate_textregion_contours(all_found_textregion_polygons): + from . import ensure_array + return ensure_array( + [polygon2contour(contour2polygon(contour, dilate=6)) + for contour in all_found_textregion_polygons]) def contour2polygon(contour: Union[np.ndarray, Sequence[Sequence[Sequence[Number]]]], dilate=0): polygon = Polygon([point[0] for point in contour]) diff --git a/src/eynollah/utils/separate_lines.py b/src/eynollah/utils/separate_lines.py index c220234..869cd23 100644 --- a/src/eynollah/utils/separate_lines.py +++ b/src/eynollah/utils/separate_lines.py @@ -399,14 +399,14 @@ def separate_lines(img_patch, contour_text_interest, thetha, x_help, y_help): point_down_rot3=point_down_rot3-y_help point_down_rot4=point_down_rot4-y_help - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(point_up)], - [int(x_max), int(point_up)], - [int(x_max), int(point_down)], - [int(x_min), int(point_down)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(point_up)]], + [[int(x_max), int(point_up)]], + [[int(x_max), int(point_down)]], + [[int(x_min), int(point_down)]]])) elif len(peaks) < 1: pass @@ -458,14 +458,14 @@ def separate_lines(img_patch, contour_text_interest, thetha, x_help, y_help): point_down_rot3=point_down_rot3-y_help point_down_rot4=point_down_rot4-y_help - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(y_min)], - [int(x_max), int(y_min)], - [int(x_max), int(y_max)], - [int(x_min), int(y_max)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(y_min)]], + [[int(x_max), int(y_min)]], + [[int(x_max), int(y_max)]], + [[int(x_min), int(y_max)]]])) elif len(peaks) == 2: dis_to_next = np.abs(peaks[1] - peaks[0]) for jj in range(len(peaks)): @@ -526,14 +526,14 @@ def separate_lines(img_patch, contour_text_interest, thetha, x_help, y_help): point_down_rot3=point_down_rot3-y_help point_down_rot4=point_down_rot4-y_help - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(point_up)], - [int(x_max), int(point_up)], - [int(x_max), int(point_down)], - [int(x_min), int(point_down)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(point_up)]], + [[int(x_max), int(point_up)]], + [[int(x_max), int(point_down)]], + [[int(x_min), int(point_down)]]])) else: for jj in range(len(peaks)): if jj == 0: @@ -602,14 +602,14 @@ def separate_lines(img_patch, contour_text_interest, thetha, x_help, y_help): point_down_rot3=point_down_rot3-y_help point_down_rot4=point_down_rot4-y_help - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(point_up)], - [int(x_max), int(point_up)], - [int(x_max), int(point_down)], - [int(x_min), int(point_down)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(point_up)]], + [[int(x_max), int(point_up)]], + [[int(x_max), int(point_down)]], + [[int(x_min), int(point_down)]]])) return peaks, textline_boxes_rot def separate_lines_vertical(img_patch, contour_text_interest, thetha): @@ -781,14 +781,14 @@ def separate_lines_vertical(img_patch, contour_text_interest, thetha): if point_up_rot2 < 0: point_up_rot2 = 0 - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(point_up)], - [int(x_max), int(point_up)], - [int(x_max), int(point_down)], - [int(x_min), int(point_down)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(point_up)]], + [[int(x_max), int(point_up)]], + [[int(x_max), int(point_down)]], + [[int(x_min), int(point_down)]]])) elif len(peaks) < 1: pass elif len(peaks) == 1: @@ -817,14 +817,14 @@ def separate_lines_vertical(img_patch, contour_text_interest, thetha): if point_up_rot2 < 0: point_up_rot2 = 0 - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(y_min)], - [int(x_max), int(y_min)], - [int(x_max), int(y_max)], - [int(x_min), int(y_max)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(y_min)]], + [[int(x_max), int(y_min)]], + [[int(x_max), int(y_max)]], + [[int(x_min), int(y_max)]]])) elif len(peaks) == 2: dis_to_next = np.abs(peaks[1] - peaks[0]) for jj in range(len(peaks)): @@ -872,14 +872,14 @@ def separate_lines_vertical(img_patch, contour_text_interest, thetha): if point_up_rot2 < 0: point_up_rot2 = 0 - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(point_up)], - [int(x_max), int(point_up)], - [int(x_max), int(point_down)], - [int(x_min), int(point_down)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(point_up)]], + [[int(x_max), int(point_up)]], + [[int(x_max), int(point_down)]], + [[int(x_min), int(point_down)]]])) else: for jj in range(len(peaks)): if jj == 0: @@ -938,14 +938,14 @@ def separate_lines_vertical(img_patch, contour_text_interest, thetha): if point_up_rot2 < 0: point_up_rot2 = 0 - textline_boxes_rot.append(np.array([[int(x_min_rot1), int(point_up_rot1)], - [int(x_max_rot2), int(point_up_rot2)], - [int(x_max_rot3), int(point_down_rot3)], - [int(x_min_rot4), int(point_down_rot4)]])) - textline_boxes.append(np.array([[int(x_min), int(point_up)], - [int(x_max), int(point_up)], - [int(x_max), int(point_down)], - [int(x_min), int(point_down)]])) + textline_boxes_rot.append(np.array([[[int(x_min_rot1), int(point_up_rot1)]], + [[int(x_max_rot2), int(point_up_rot2)]], + [[int(x_max_rot3), int(point_down_rot3)]], + [[int(x_min_rot4), int(point_down_rot4)]]])) + textline_boxes.append(np.array([[[int(x_min), int(point_up)]], + [[int(x_max), int(point_up)]], + [[int(x_max), int(point_down)]], + [[int(x_min), int(point_down)]]])) return peaks, textline_boxes_rot def separate_lines_new_inside_tiles2(img_patch, thetha): @@ -1560,6 +1560,9 @@ def return_deskew_slop(img_patch_org, sigma_des,n_tot_angles=100, angle2, var2 = get_smallest_skew(img_resized, sigma_des, angles2, map=map, logger=logger, plotter=plotter) if var2 > var: angle = angle2 + # precision stage: + angles = np.linspace(angle - 2.5, angle + 2.5, n_tot_angles // 2) + angle, _ = get_smallest_skew(img_resized, sigma_des, angles, map=map, logger=logger, plotter=plotter) return angle def get_smallest_skew(img, sigma_des, angles, logger=None, plotter=None, map=map): diff --git a/src/eynollah/utils/utils_ocr.py b/src/eynollah/utils/utils_ocr.py index b738e29..928c164 100644 --- a/src/eynollah/utils/utils_ocr.py +++ b/src/eynollah/utils/utils_ocr.py @@ -370,8 +370,8 @@ def break_curved_line_into_small_pieces_and_then_merge(img_curved, mask_curved, return img_curved, img_bin_curved def return_textline_contour_with_added_box_coordinate(textline_contour, box_ind): - textline_contour[:,0] = textline_contour[:,0] + box_ind[2] - textline_contour[:,1] = textline_contour[:,1] + box_ind[0] + textline_contour[:,:,0] += box_ind[2] + textline_contour[:,:,1] += box_ind[0] return textline_contour diff --git a/src/eynollah/writer.py b/src/eynollah/writer.py index 1781230..4f0827f 100644 --- a/src/eynollah/writer.py +++ b/src/eynollah/writer.py @@ -2,11 +2,12 @@ # pylint: disable=import-error from pathlib import Path import os.path -from typing import Optional import logging -from .utils.xml import create_page_xml, xml_reading_order -from .utils.counter import EynollahIdCounter +from typing import Optional +import numpy as np +from shapely import affinity, clip_by_rect +from ocrd_utils import points_from_polygon from ocrd_models.ocrd_page import ( BorderType, CoordsType, @@ -19,6 +20,10 @@ from ocrd_models.ocrd_page import ( to_xml ) +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, curved_line, pcgts=None): @@ -38,20 +43,14 @@ class EynollahXmlWriter: def image_filename_stem(self): return Path(Path(self.image_filename).name).stem - def calculate_page_coords(self, cont_page): - self.logger.debug('enter calculate_page_coords') - points_page_print = "" - for _, contour in enumerate(cont_page[0]): - if len(contour) == 2: - points_page_print += str(int((contour[0]) / self.scale_x)) - points_page_print += ',' - points_page_print += str(int((contour[1]) / self.scale_y)) - else: - points_page_print += str(int((contour[0][0]) / self.scale_x)) - points_page_print += ',' - points_page_print += str(int((contour[0][1] ) / self.scale_y)) - points_page_print = points_page_print + ' ' - return points_page_print[:-1] + def calculate_points(self, contour, offset=None): + self.logger.debug('enter calculate_points') + poly = contour2polygon(contour) + 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.width_org, self.height_org)) + return points_from_polygon(poly.exterior.coords[:-1]) def serialize_lines_in_region(self, text_region, all_found_textline_polygons, region_idx, page_coord, all_box_coord, slopes, counter, ocr_all_textlines_textregion): self.logger.debug('enter serialize_lines_in_region') @@ -64,16 +63,12 @@ class EynollahXmlWriter: text_region.add_TextLine(textline) text_region.set_orientation(-slopes[region_idx]) region_bboxes = all_box_coord[region_idx] - points_co = '' - for point in polygon_textline: - if len(point) != 2: - point = point[0] - point_x = point[0] + page_coord[2] - point_y = point[1] + page_coord[0] - point_x = max(0, int(point_x / self.scale_x)) - point_y = max(0, int(point_y / self.scale_y)) - points_co += f'{point_x},{point_y} ' - coords.set_points(points_co[:-1]) + offset = [page_coord[2], page_coord[0]] + # FIXME: or actually... self.curved_line or np.abs(slopes[region_idx]) > 45? + if self.curved_line and np.abs(slopes[region_idx]) > 45: + offset[0] += region_bboxes[2] + offset[1] += region_bboxes[0] + coords.set_points(self.calculate_points(polygon_textline, offset)) def write_pagexml(self, pcgts): self.logger.info("output filename: '%s'", self.output_filename) @@ -168,9 +163,13 @@ class EynollahXmlWriter: # create the file structure pcgts = self.pcgts if self.pcgts else create_page_xml(self.image_filename, self.height_org, self.width_org) page = pcgts.get_Page() - assert page - page.set_Border(BorderType(Coords=CoordsType(points=self.calculate_page_coords(cont_page)))) + if len(cont_page): + page.set_Border(BorderType(Coords=CoordsType(points=self.calculate_points(cont_page[0])))) + if skip_layout_reading_order: + offset = None + else: + offset = [page_coord[2], page_coord[0]] counter = EynollahIdCounter() if len(order_of_texts): _counter_marginals = EynollahIdCounter(region_idx=len(order_of_texts)) @@ -183,8 +182,7 @@ class EynollahXmlWriter: for mm, region_contour in enumerate(found_polygons_text_region): textregion = TextRegionType( id=counter.next_region_id, type_='paragraph', - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord, - skip_layout_reading_order)) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)) ) assert textregion.Coords if conf_contours_textregions: @@ -201,7 +199,7 @@ class EynollahXmlWriter: for mm, region_contour in enumerate(found_polygons_text_region_h): textregion = TextRegionType( id=counter.next_region_id, type_='heading', - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord)) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)) ) assert textregion.Coords if conf_contours_textregions_h: @@ -217,7 +215,7 @@ class EynollahXmlWriter: for mm, region_contour in enumerate(found_polygons_marginals_left): marginal = TextRegionType( id=counter.next_region_id, type_='marginalia', - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord)) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)) ) page.add_TextRegion(marginal) if ocr_all_textlines_marginals_left: @@ -229,7 +227,7 @@ class EynollahXmlWriter: for mm, region_contour in enumerate(found_polygons_marginals_right): marginal = TextRegionType( id=counter.next_region_id, type_='marginalia', - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord)) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)) ) page.add_TextRegion(marginal) if ocr_all_textlines_marginals_right: @@ -242,7 +240,7 @@ class EynollahXmlWriter: for mm, region_contour in enumerate(found_polygons_drop_capitals): dropcapital = TextRegionType( id=counter.next_region_id, type_='drop-capital', - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord)) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)) ) page.add_TextRegion(dropcapital) all_box_coord_drop = [[0, 0, 0, 0]] @@ -257,33 +255,17 @@ class EynollahXmlWriter: for region_contour in found_polygons_text_region_img: page.add_ImageRegion( ImageRegionType(id=counter.next_region_id, - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord)))) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)))) for region_contour in polygons_seplines: page.add_SeparatorRegion( SeparatorRegionType(id=counter.next_region_id, - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, [0, 0, 0, 0])))) + Coords=CoordsType(points=self.calculate_points(region_contour, None)))) for region_contour in found_polygons_tables: page.add_TableRegion( TableRegionType(id=counter.next_region_id, - Coords=CoordsType(points=self.calculate_polygon_coords(region_contour, page_coord)))) + Coords=CoordsType(points=self.calculate_points(region_contour, offset)))) return pcgts - def calculate_polygon_coords(self, contour, page_coord, skip_layout_reading_order=False): - self.logger.debug('enter calculate_polygon_coords') - coords = '' - for point in contour: - if len(point) != 2: - point = point[0] - point_x = point[0] - point_y = point[1] - if not skip_layout_reading_order: - point_x += page_coord[2] - point_y += page_coord[0] - point_x = int(point_x / self.scale_x) - point_y = int(point_y / self.scale_y) - coords += str(point_x) + ',' + str(point_y) + ' ' - return coords[:-1] - diff --git a/tests/cli_tests/test_binarization.py b/tests/cli_tests/test_binarization.py index aa52957..1287ffa 100644 --- a/tests/cli_tests/test_binarization.py +++ b/tests/cli_tests/test_binarization.py @@ -22,7 +22,7 @@ def test_run_eynollah_binarization_filename( '-o', str(outfile), ] + options, [ - 'Predicting' + 'Loaded model' ] ) assert outfile.exists() @@ -46,8 +46,8 @@ def test_run_eynollah_binarization_directory( '-o', str(outdir), ], [ - f'Predicting {image_resources[0].name}', - f'Predicting {image_resources[1].name}', + f'Binarizing [ 1/2] {image_resources[0].name}', + f'Binarizing [ 2/2] {image_resources[1].name}', ] ) assert len(list(outdir.iterdir())) == 2 diff --git a/train/requirements.txt b/train/requirements.txt index 63f3813..8ad884d 100644 --- a/train/requirements.txt +++ b/train/requirements.txt @@ -1,6 +1,6 @@ sacred seaborn -numpy <1.24.0 +numpy tqdm imutils scipy