diff --git a/.dockerignore b/.dockerignore index 562fb6f..5fafd44 100644 --- a/.dockerignore +++ b/.dockerignore @@ -2,5 +2,9 @@ tests dist build env* +venv* *.egg-info models_eynollah* +reloaded +*.h5 +config_files* diff --git a/.github/workflows/test-eynollah.yml b/.github/workflows/test-eynollah.yml index 82de94d..44910ce 100644 --- a/.github/workflows/test-eynollah.yml +++ b/.github/workflows/test-eynollah.yml @@ -64,6 +64,9 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip + # preempt CUDA dependencies (which need core's recipe) + pip install onnxruntime tensorflow tf-keras "torch<2.11" + sed -i '/onnxruntime-gpu/d;/tensorrt/d;/torch/d;/tensorflow/d' requirements*.txt make install-dev EXTRAS=OCR,plotting make deps-test EXTRAS=OCR,plotting diff --git a/.gitignore b/.gitignore index 49835a7..3b4f340 100644 --- a/.gitignore +++ b/.gitignore @@ -1,14 +1,18 @@ *.egg-info __pycache__ sbb_newspapers_org_image/pylint.log -models_eynollah* -models_ocr* -models_layout* default-2021-03-09 output.html /build -/dist *.tif *.sw? TAGS uv.lock +/dist +/ignore +*.log +.env +models/reloaded +models/packages +models/models_eynollah +models/pretrained_model diff --git a/Dockerfile b/Dockerfile index a15776e..a041807 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,7 +15,7 @@ LABEL \ org.opencontainers.image.documentation="https://github.com/qurator-spk/eynollah/blob/${VCS_REF}/README.md" \ org.opencontainers.image.revision=$VCS_REF \ org.opencontainers.image.created=$BUILD_DATE \ - org.opencontainers.image.base.name=ocrd/core-cuda-tf2 + org.opencontainers.image.base.name=ocrd/core-cuda-onnx ENV DEBIAN_FRONTEND=noninteractive # set proper locales @@ -39,9 +39,10 @@ RUN ocrd ocrd-tool ocrd-tool.json dump-tools > $(dirname $(ocrd bashlib filename # prepackage ocrd-all-module-dir.json RUN ocrd ocrd-tool ocrd-tool.json dump-module-dirs > $(dirname $(ocrd bashlib filename))/ocrd-all-module-dir.json # install everything and reduce image size -RUN make install EXTRAS=OCR && rm -rf /build/eynollah -# fixup for broken cuDNN installation (Torch pulls in 8.5.0, which is incompatible with Tensorflow) -RUN pip install nvidia-cudnn-cu11==8.6.0.163 +# FIXME: EXTRAS=OCR (should become extra Dockerfile based on ocrd/core-cuda-tf2 and ocrd/core-cuda-torch) +RUN make install && rm -rf /build/eynollah +# fixup for broken cuDNN installation (Torch may pull in version which is incompatible with Tensorflow) +RUN pip install "nvidia-cudnn-cu12<9.10.2.21" # smoke test RUN eynollah --help diff --git a/Makefile b/Makefile index f54cf5b..2734672 100644 --- a/Makefile +++ b/Makefile @@ -2,7 +2,7 @@ PYTHON ?= python3 PIP ?= pip3 EXTRAS ?= -DOCKER_BASE_IMAGE ?= docker.io/ocrd/core-cuda-tf2:v3.13.0 +DOCKER_BASE_IMAGE ?= docker.io/ocrd/core-cuda-onnx:v3.13.1 DOCKER_TAG ?= ocrd/eynollah DOCKER ?= docker WGET = wget -O @@ -13,7 +13,8 @@ WGET = wget -O #SEG_MODEL := https://github.com/qurator-spk/eynollah/releases/download/v0.3.0/models_eynollah.tar.gz #SEG_MODEL := https://github.com/qurator-spk/eynollah/releases/download/v0.3.1/models_eynollah.tar.gz #SEG_MODEL := https://zenodo.org/records/17194824/files/models_layout_v0_5_0.tar.gz?download=1 -EYNOLLAH_MODELS_URL := https://zenodo.org/records/17727267/files/models_all_v0_8_0.zip +# EYNOLLAH_MODELS_URL := https://zenodo.org/records/21362927/files/models_inference_all_v0_9_0.zip +EYNOLLAH_MODELS_URL := https://zenodo.org/records/21381102/files/models_inference_all_v0_9_1.zip EYNOLLAH_MODELS_ZIP = $(notdir $(EYNOLLAH_MODELS_URL)) EYNOLLAH_MODELS_DIR = $(EYNOLLAH_MODELS_ZIP:%.zip=%) diff --git a/README.md b/README.md index 7e9a90a..be21a7a 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,12 @@ documents using a combination of multiple deep learning models and heuristics; therefore processing can be slow. ## Installation -Python `3.8-3.11` with Tensorflow `<2.13` on Linux are currently supported. -For (limited) GPU support the CUDA toolkit needs to be installed. -A working config is CUDA `11.8` with cuDNN `8.6`. + +Python `3.8-3.11` with ONNX Runtime on Linux are currently supported. + +For GPU support, NVidia drivers supporting CUDA 12 must be installed. +The runtime dependencies will pull in ONNX, TensorRT and CUDA runtime +libraries (including cuDNN) from PyPI. You can either install from PyPI @@ -52,6 +55,13 @@ pip install "eynollah[OCR]" make install EXTRAS=OCR ``` +> **Note**: Requirements for OCR are more involved, +> as they may need Tensorflow (with tf-keras) and/or +> Torch (with transformers). Those two frameworks may +> also have conflicting CUDA dependencies. An ONNX +> conversion for these models may be achieved soon. +> :construction: + ### Docker Use @@ -64,7 +74,12 @@ When using Eynollah with Docker, see [`docker.md`](https://github.com/qurator-sp ## Models -Pretrained models can be downloaded from [Zenodo](https://zenodo.org/records/17727267) or [Hugging Face](https://huggingface.co/SBB?search_models=eynollah). +Pretrained models can be downloaded from [Zenodo](https://doi.org/10.5281/zenodo.17194823) or [Hugging Face](https://huggingface.co/SBB?search_models=eynollah). + +For fast runtime inference, download the ONNX models distributed as `models_inference_...zip`. + +For finetuning training, download the original (Tensorflow / Torch) models distributed as `models_training...zip` +(and install the `[training]` extra). For model documentation and model cards, see [`models.md`](https://github.com/qurator-spk/eynollah/tree/main/docs/models.md). @@ -83,18 +98,26 @@ Eynollah supports five use cases: Some example outputs can be found in [`examples.md`](https://github.com/qurator-spk/eynollah/tree/main/docs/examples.md). +The **generic options** shared by all subcommands are: +```sh + -m + -mv + -D + -l +``` + ### Layout Analysis -The layout analysis module is responsible for detecting layout elements, identifying text lines, and determining reading -order using heuristic methods or a [pretrained model](https://github.com/qurator-spk/eynollah#machine-based-reading-order). +Detects layout elements, i.e. regions of various types and text lines, +and determines their reading order using either heuristic methods or a +[pretrained model](https://github.com/qurator-spk/eynollah#machine-based-reading-order). The command-line interface for layout analysis can be called like this: ```sh -eynollah layout \ +eynollah [GENERIC_OPTIONS] layout \ -i | -di \ -o \ - -m \ [OPTIONS] ``` @@ -106,7 +129,7 @@ The following options can be used to further configure the processing: | `-tab` | apply table detection | | `-ae` | apply enhancement (the resulting image is saved to the output directory) | | `-as` | apply scaling | -| `-cl` | apply contour detection for curved text lines instead of bounding boxes | +| `-cl` | apply contour detection for curved text lines, deskewing all regions independently | | `-ib` | apply binarization (the resulting image is saved to the output directory) | | `-ep` | enable plotting (MUST always be used with `-sl`, `-sd`, `-sa`, `-si` or `-ae`) | | `-ho` | ignore headers for reading order dectection | @@ -115,79 +138,93 @@ The following options can be used to further configure the processing: | `-sl ` | save layout prediction as plot to this directory | | `-sp ` | save cropped page image to this directory | | `-sa ` | save all (plot, enhanced/binary image, layout) to this directory | -| `-thart` | threshold of artifical class in the case of textline detection. The default value is 0.1 | -| `-tharl` | threshold of artifical class in the case of layout detection. The default value is 0.1 | +| `-thart` | confidence threshold of artifical boundary class during textline detection | +| `-tharl` | confidence threshold of artifical boundary class during region detection | | `-ncu` | upper limit of columns in document image | | `-ncl` | lower limit of columns in document image | | `-slro` | skip layout detection and reading order | | `-romb` | apply machine based reading order detection | | `-ipe` | ignore page extraction | +| `-j` | number of CPU jobs to run parallel (useful with -di) | +| `-H` | when to halt when some jobs fail | +The default is to only perform layout detection of main regions +(background, text, images, separators and marginals). -If no further option is set, the tool performs layout detection of main regions (background, text, images, separators -and marginals). -The best output quality is achieved when RGB images are used as input rather than greyscale or binarized images. +The best output quality is achieved when RGB images are used as input +rather than greyscale or binarized images. -Additional documentation can be found in [`usage.md`](https://github.com/qurator-spk/eynollah/tree/main/docs/usage.md). +Additional documentation can be found in +[`usage.md`](https://github.com/qurator-spk/eynollah/tree/main/docs/usage.md). ### Binarization -The binarization module performs document image binarization using pretrained pixelwise segmentation models. +Performs document image binarization (thresholding) +using pretrained pixelwise segmentation models. The command-line interface for binarization can be called like this: ```sh -eynollah binarization \ +eynollah [GENERIC_OPTIONS] binarization \ -i | -di \ -o \ - -m + [OPTIONS] ``` ### Image Enhancement -TODO + +This enlarges and enhances images. Useful in case the scan quality is low. + +```sh +eynollah [GENERIC_OPTIONS] enhancement \ + -i | -di \ + -o \ + [OPTIONS] +``` + +| option | description | +|-------------------|:--------------------------------------------------------------------------------------------| +| `-sos` | save the enhanced image in original image size | +| `-ncu` | upper limit of columns in document image | +| `-ncl` | lower limit of columns in document image | ### OCR -The OCR module performs text recognition using either a CNN-RNN model or a Transformer model. +Performs text recognition using either a CNN-RNN model or a Transformer model. +Needs a PAGE-XML input file. The command-line interface for OCR can be called like this: ```sh -eynollah ocr \ +eynollah [GENERIC_OPTIONS] ocr \ -i | -di \ -dx \ -o \ - -m | --model_name ``` The following options can be used to further configure the ocr processing: | option | description | |-------------------|:-------------------------------------------------------------------------------------------| +| `-trocr` | use transformer OCR model instead of CNN-RNN model | | `-dib` | directory of binarized images (file type must be '.png'), prediction with both RGB and bin | | `-doit` | directory for output images rendered with the predicted text | -| `--model_name` | file path to use specific model for OCR | -| `-trocr` | use transformer ocr model (otherwise cnn_rnn model is used) | -| `-etit` | export textline images and text in xml to output dir (OCR training data) | | `-nmtc` | cropped textline images will not be masked with textline contour | | `-bs` | ocr inference batch size. Default batch size is 2 for trocr and 8 for cnn_rnn models | -| `-ds_pref` | add an abbrevation of dataset name to generated training data | | `-min_conf` | minimum OCR confidence value. OCR with textline conf lower than this will be ignored | ### Reading Order Detection -Reading order detection can be performed either as part of layout analysis based on image input, or, currently under -development, based on pre-existing layout analysis data in PAGE-XML format as input. -The reading order detection module employs a pretrained model to identify the reading order from layouts represented in PAGE-XML files. +Reading order can be detected either during layout analysis, +or as a separate module, which requires a PAGE-XML input file. The command-line interface for machine based reading order can be called like this: ```sh -eynollah machine-based-reading-order \ +eynollah [GENERIC_OPTIONS] machine-based-reading-order \ -i | -di \ -xml | -dx \ - -m \ -o ``` diff --git a/docs/ocrd.md b/docs/ocrd.md index 9e7e268..fc68a31 100644 --- a/docs/ocrd.md +++ b/docs/ocrd.md @@ -5,7 +5,7 @@ formally described in [`ocrd-tool.json`](https://github.com/qurator-spk/eynollah When using Eynollah in OCR-D, the source image file group with (preferably) RGB images should be used as input like this: - ocrd-eynollah-segment -I OCR-D-IMG -O OCR-D-SEG -P models eynollah_layout_v0_5_0 + ocrd-eynollah-segment -I OCR-D-IMG -O OCR-D-SEG -P models eynollah_layout_v0_9_0 If the input file group is PAGE-XML (from a previous OCR-D workflow step), Eynollah behaves as follows: - existing regions are kept and ignored (i.e. in effect they might overlap segments from Eynollah results) @@ -17,7 +17,7 @@ If the input file group is PAGE-XML (from a previous OCR-D workflow step), Eynol (because some other preprocessing step was in effect like `denoised`), then the output PAGE-XML will be based on that as new top-level (`@imageFilename`) - ocrd-eynollah-segment -I OCR-D-XYZ -O OCR-D-SEG -P models eynollah_layout_v0_5_0 + ocrd-eynollah-segment -I OCR-D-XYZ -O OCR-D-SEG -P models eynollah_layout_v0_9_0 In general, it makes more sense to add other workflow steps **after** Eynollah. diff --git a/docs/usage.md b/docs/usage.md index da164de..2f9a688 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -3,49 +3,69 @@ The command-line interface can be called like this: ```sh eynollah \ + [GENERIC_OPTIONS] \ + layout \ -i | -di \ -o \ - -m \ - [OPTIONS] + [LAYOUT_OPTIONS] ``` -## Processing options -The following options can be used to further configure the processing: +## Generic options +Pass any of the following options: -| option | description | +| **option** | **description** | +|---------------------------------------------------|:----------------------------------------------------------------| +| -m | override default directory `$PWD/models_eynollah` | +| -mv | override specific models, e.g. `region_1_2 '' /path/to/my.onnx` | +| -D | allocate models to GPUs, e.g. `col*:CPU,page:GPU1,*:GPU0` | +| -l | override default `INFO` log level e.g. `DEBUG` | + +## Processing options +The following options can be used to further control layout analysis: + +| **option** | **description** | |-------------------|:-------------------------------------------------------------------------------| -| `-fl` | full layout analysis including all steps and segmentation classes | -| `-light` | lighter and faster but simpler method for main region detection and deskewing | +| `-fl` | full layout analysis including all steps and segmentation classes (recommended)| | `-tab` | apply table detection | | `-ae` | apply enhancement (the resulting image is saved to the output directory) | | `-as` | apply scaling | -| `-cl` | apply contour detection for curved text lines instead of bounding boxes | +| `-cl` | apply contour detection for curved text lines, deskewing regions independently | | `-ib` | apply binarization (the resulting image is saved to the output directory) | | `-ep` | enable plotting (MUST always be used with `-sl`, `-sd`, `-sa`, `-si` or `-ae`) | -| `-eoi` | extract only images to output directory (other processing will not be done) | | `-ho` | ignore headers for reading order dectection | | `-si ` | save image regions detected to this directory | | `-sd ` | save deskewed image to this directory | | `-sl ` | save layout prediction as plot to this directory | | `-sp ` | save cropped page image to this directory | | `-sa ` | save all (plot, enhanced/binary image, layout) to this directory | +| `-thart` | confidence threshold of artifical boundary class during textline detection | +| `-tharl` | confidence threshold of artifical boundary class during region detection | +| `-ncu` | upper limit of columns in document image | +| `-ncl` | lower limit of columns in document image | +| `-slro` | skip layout detection and reading order | +| `-romb` | apply machine based reading order detection | +| `-ipe` | ignore page extraction | +| `-j` | number of CPU jobs to run parallel (useful with -di) | +| `-H` | when to halt if some jobs fail, e.g. `0.1` for 10% or `3` for 3 pages | If no option is set, the tool performs detection of main regions (background, text, images, separators and marginals). -### `--full-layout` vs `--no-full-layout` +### `--full-layout` vs shallow -Here are the difference in elements detected depending on the `--full-layout`/`--no-full-layout` command line flags: +Here are the differences in segment types detected: -| | `--full-layout` | `--no-full-layout` | -|--------------------------|-----------------|--------------------| -| reading order | x | x | -| header regions | x | - | -| text regions | x | x | -| text regions / text line | x | x | -| drop-capitals | x | - | -| marginals | x | x | -| marginals / text line | x | x | -| image region | x | x | +| | `-fl` | without | +|--------------------------|-------|---------| +| reading order | x | x | +| header regions | x | - | +| text regions | x | x | +| text regions / textlines | x | x | +| drop-capitals | x | - | +| marginals | x | x | +| marginals / textlines | x | x | +| image regions | x | x | + +(Note: No marginals are detected for pages with 3 columns or more.) ## Use as OCR-D processor Eynollah ships with a CLI interface to be used as [OCR-D](https://ocr-d.de) processor that is described in @@ -54,13 +74,13 @@ Eynollah ships with a CLI interface to be used as [OCR-D](https://ocr-d.de) proc The source image file group with (preferably) RGB images should be used as input for Eynollah like this: ``` -ocrd-eynollah-segment -I OCR-D-IMG -O SEG-LINE -P models +ocrd-eynollah-segment -I OCR-D-IMG -O SEG-LINE -P full_layout true ``` Any image referenced by `@imageFilename` in PAGE-XML is passed on directly to Eynollah as a processor, so that e.g. ``` -ocrd-eynollah-segment -I OCR-D-IMG-BIN -O SEG-LINE -P models +ocrd-eynollah-segment -I OCR-D-IMG-BIN -O SEG-LINE -P full_layout true ``` uses the original (RGB) image despite any binarization that may have occured in previous OCR-D processing steps. diff --git a/models/Makefile b/models/Makefile new file mode 100644 index 0000000..ed97039 --- /dev/null +++ b/models/Makefile @@ -0,0 +1,142 @@ +SHELL = bash -e + +VERSION = v0_9_1 +MODELS_SRC = models_eynollah +MODELS_DST = reloaded/models_eynollah + + +# eynollah-main-regions-aug-rotation_20210425 +# eynollah-main-regions-aug-scaling_20210425 +# eynollah-main-regions-ensembled_20210425 +# eynollah-main-regions_20220314 +# eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 +# eynollah-tables_20210319 + +CURRENT_MODELS := +CURRENT_MODELS += eynollah-main-regions-aug-rotation_20210425 +CURRENT_MODELS += eynollah-main-regions-aug-scaling_20210425 +CURRENT_MODELS += eynollah-main-regions-ensembled_20210425 +CURRENT_MODELS += eynollah-main-regions_20220314 +CURRENT_MODELS += eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 +CURRENT_MODELS += eynollah-tables_20210319 +CURRENT_MODELS += eynollah-main-regions_20220314 +CURRENT_MODELS += model_eynollah_page_extraction_20250915 +CURRENT_MODELS += model_eynollah_reading_order_20250824 +CURRENT_MODELS += modelens_e_l_all_sp_0_1_2_3_4_171024 +CURRENT_MODELS += modelens_full_lay_1__4_3_091124 +CURRENT_MODELS += modelens_table_0t4_201124 +CURRENT_MODELS += modelens_textline_0_1__2_4_16092024 +CURRENT_MODELS += model_eynollah_ocr_cnnrnn_20250930 +CURRENT_MODELS += eynollah-binarization_20210425 +CURRENT_MODELS += eynollah-column-classifier_20210425 +CURRENT_MODELS += eynollah-enhancement_20210425 + +help: + @echo "Targets:" + @echo "" + @echo "Model conversion:" + @echo " all Convert all current models to TensorFlow SavedModel (default)" + @echo " tf[-serving] Convert all current models to TensorFlow SavedModel" + @echo " keras Convert all current models to Keras format (.keras)" + @echo " hdf5 Convert all current models to HDF5 format (.h5)" + @echo " onnx Convert all eligible current models to ONNX (.onnx)" + @echo " reload Reload selected pre-v0.8 models without Lambda layers" + @echo "" + @echo "Model reloading:" + @echo " compare Compare original and v0.8 reloaded models" + @echo "" + @echo "Packaging:" + @echo " packages Build all Zenodo distribution archives" + @echo " upload Upload all distribution archives to Zenodo" + @echo "" + @echo "Cleanup:" + @echo " clean Remove all generated files" + @echo " clean-reload Remove reloaded model directories" + @echo " clean-packages Remove generated package archives" + @echo "" + @echo "Variables:" + @echo " VERSION=$(VERSION)" + + +# tf (SavedModel format) for training +# onnx conversion for fast inference +all: tf onnx + +tf-serving: $(CURRENT_MODELS:%=$(MODELS_DST)/%) +tf: $(CURRENT_MODELS:%=$(MODELS_DST)/%) +keras: $(CURRENT_MODELS:%=$(MODELS_DST)/%.keras) +hdf5: $(CURRENT_MODELS:%=$(MODELS_DST)/%.h5) +onnx: $(CURRENT_MODELS:%=$(MODELS_DST)/%.onnx) + +# distinguish tf from tf-serving: target pattern is the same, +# so check if either is current goal, otherwise assumg tf +$(MODELS_DST)/%: FORMAT = $(or $(filter tf-serving,$(MAKECMDGOALS)), tf) +$(MODELS_DST)/%: $(MODELS_SRC)/% + $(if $(and $(filter tf-serving,$(FORMAT)),$(findstring _ocr,$@)),$(warning skipping $@: OCR CTC decoder fails in TF-Serving) : )\ + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).$(FORMAT).log 2>&1 || { cat $(notdir $<).$(FORMAT).log; false; } + +$(MODELS_DST)/%.keras: $(MODELS_SRC)/% + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).keras.log 2>&1 || { cat $(notdir $<).keras.log; false; } + +$(MODELS_DST)/%.h5: $(MODELS_SRC)/% + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).hdf5.log 2>&1 || { cat $(notdir $<).hdf5.log; false; } + +$(MODELS_DST)/%.onnx: $(MODELS_SRC)/% + $(if $(findstring _ocr,$@),$(warning skipping $@: OCR CTC decoder is buggy in ONNX) : )\ + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).onnx.log 2>&1 || { cat $(notdir $<).onnx.log; false; } + +compare: + for i in `find $(MODELS_DST) -mindepth 2`;do \ + n=$(MODELS_SRC)$${i#$(MODELS_DST)}; \ + du -bs $$n $$i ; \ + done + +clean: clean-reload clean-packages + +clean-reload: + rm -rf $(RELOADABLE_MODELS) + +clean-packages: + rm -rf $(PACKAGES_DIR) + +# NB: this works using symlinks +PACKAGES_DIR = packages +BUNDLES = inference training +CATEGORIES = layout ocr extra all + +packages: $(foreach B,$(BUNDLES),$(foreach C,$(CATEGORIES),$(PACKAGES_DIR)/models_$(B)_$(C)_$(VERSION).zip)) + +$(PACKAGES_DIR)/models_%.zip: + @mkdir -p $(PACKAGES_DIR) ;\ + bundle=$(word 1,$(subst _, ,$*)); \ + category=$(word 2,$(subst _, ,$*)); \ + echo "Packaging $$bundle/$$category to $(notdir $@)" ;\ + cd "dist/$$bundle/$$category" && \ + zip -vqr "$$OLDPWD/$@" models_eynollah + + +upload: $(foreach B,$(BUNDLES),$(foreach C,$(CATEGORIES),upload/$(B)_$(C))) + +upload/%: + @source .env ;\ + bash scripts/zenodo_upload.sh $$ZENODO_ID $(PACKAGES_DIR)/models_$*_$(VERSION).zip --verbose diff --git a/models/dist/inference/all/models_eynollah/characters_org.txt b/models/dist/inference/all/models_eynollah/characters_org.txt new file mode 120000 index 0000000..9b62c85 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/characters_org.txt @@ -0,0 +1 @@ +../../../../models_eynollah/characters_org.txt \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-binarization_20210309 b/models/dist/inference/all/models_eynollah/eynollah-binarization_20210309 new file mode 120000 index 0000000..a043883 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-binarization_20210309 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization_20210309 \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-binarization_20210425.onnx b/models/dist/inference/all/models_eynollah/eynollah-binarization_20210425.onnx new file mode 120000 index 0000000..eb64495 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-binarization_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-binarization_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-column-classifier_20210425.onnx b/models/dist/inference/all/models_eynollah/eynollah-column-classifier_20210425.onnx new file mode 120000 index 0000000..d7a154f --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-column-classifier_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-column-classifier_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-enhancement_20210425.onnx b/models/dist/inference/all/models_eynollah/eynollah-enhancement_20210425.onnx new file mode 120000 index 0000000..1e400f3 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-enhancement_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-enhancement_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx b/models/dist/inference/all/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx new file mode 120000 index 0000000..e1aea76 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx b/models/dist/inference/all/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx new file mode 120000 index 0000000..921f8f3 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx b/models/dist/inference/all/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx new file mode 120000 index 0000000..6c23442 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-main-regions_20220314.onnx b/models/dist/inference/all/models_eynollah/eynollah-main-regions_20220314.onnx new file mode 120000 index 0000000..e15fbd9 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-main-regions_20220314.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions_20220314.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx b/models/dist/inference/all/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx new file mode 120000 index 0000000..9034c95 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/eynollah-tables_20210319.onnx b/models/dist/inference/all/models_eynollah/eynollah-tables_20210319.onnx new file mode 120000 index 0000000..912a400 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/eynollah-tables_20210319.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-tables_20210319.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 b/models/dist/inference/all/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 new file mode 120000 index 0000000..9875a6a --- /dev/null +++ b/models/dist/inference/all/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 b/models/dist/inference/all/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 new file mode 120000 index 0000000..0982ba7 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/model_eynollah_ocr_trocr_20250919 b/models/dist/inference/all/models_eynollah/model_eynollah_ocr_trocr_20250919 new file mode 120000 index 0000000..d5dc0b1 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/model_eynollah_ocr_trocr_20250919 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_trocr_20250919 \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/model_eynollah_page_extraction_20250915.onnx b/models/dist/inference/all/models_eynollah/model_eynollah_page_extraction_20250915.onnx new file mode 120000 index 0000000..688b314 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/model_eynollah_page_extraction_20250915.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_page_extraction_20250915.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/model_eynollah_reading_order_20250824.onnx b/models/dist/inference/all/models_eynollah/model_eynollah_reading_order_20250824.onnx new file mode 120000 index 0000000..2db28c0 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/model_eynollah_reading_order_20250824.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_reading_order_20250824.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx b/models/dist/inference/all/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx new file mode 120000 index 0000000..c5828f6 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/modelens_full_lay_1__4_3_091124.onnx b/models/dist/inference/all/models_eynollah/modelens_full_lay_1__4_3_091124.onnx new file mode 120000 index 0000000..758c260 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/modelens_full_lay_1__4_3_091124.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_full_lay_1__4_3_091124.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/modelens_table_0t4_201124.onnx b/models/dist/inference/all/models_eynollah/modelens_table_0t4_201124.onnx new file mode 120000 index 0000000..5768a79 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/modelens_table_0t4_201124.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_table_0t4_201124.onnx \ No newline at end of file diff --git a/models/dist/inference/all/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx b/models/dist/inference/all/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx new file mode 120000 index 0000000..2b842b4 --- /dev/null +++ b/models/dist/inference/all/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx \ No newline at end of file diff --git a/models/dist/inference/extra/models_eynollah/eynollah-binarization_20210309 b/models/dist/inference/extra/models_eynollah/eynollah-binarization_20210309 new file mode 120000 index 0000000..a043883 --- /dev/null +++ b/models/dist/inference/extra/models_eynollah/eynollah-binarization_20210309 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization_20210309 \ No newline at end of file diff --git a/models/dist/inference/extra/models_eynollah/eynollah-binarization_20210425.onnx b/models/dist/inference/extra/models_eynollah/eynollah-binarization_20210425.onnx new file mode 120000 index 0000000..eb64495 --- /dev/null +++ b/models/dist/inference/extra/models_eynollah/eynollah-binarization_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-binarization_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/extra/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx b/models/dist/inference/extra/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx new file mode 120000 index 0000000..921f8f3 --- /dev/null +++ b/models/dist/inference/extra/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/extra/models_eynollah/eynollah-tables_20210319.onnx b/models/dist/inference/extra/models_eynollah/eynollah-tables_20210319.onnx new file mode 120000 index 0000000..912a400 --- /dev/null +++ b/models/dist/inference/extra/models_eynollah/eynollah-tables_20210319.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-tables_20210319.onnx \ No newline at end of file diff --git a/models/dist/inference/extra/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx b/models/dist/inference/extra/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx new file mode 120000 index 0000000..2b842b4 --- /dev/null +++ b/models/dist/inference/extra/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-binarization_20210425.onnx b/models/dist/inference/layout/models_eynollah/eynollah-binarization_20210425.onnx new file mode 120000 index 0000000..eb64495 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-binarization_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-binarization_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-column-classifier_20210425.onnx b/models/dist/inference/layout/models_eynollah/eynollah-column-classifier_20210425.onnx new file mode 120000 index 0000000..d7a154f --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-column-classifier_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-column-classifier_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-enhancement_20210425.onnx b/models/dist/inference/layout/models_eynollah/eynollah-enhancement_20210425.onnx new file mode 120000 index 0000000..1e400f3 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-enhancement_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-enhancement_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx b/models/dist/inference/layout/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx new file mode 120000 index 0000000..e1aea76 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-aug-rotation_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx b/models/dist/inference/layout/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx new file mode 120000 index 0000000..921f8f3 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-aug-scaling_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx b/models/dist/inference/layout/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx new file mode 120000 index 0000000..6c23442 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions-ensembled_20210425.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-main-regions_20220314.onnx b/models/dist/inference/layout/models_eynollah/eynollah-main-regions_20220314.onnx new file mode 120000 index 0000000..e15fbd9 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-main-regions_20220314.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions_20220314.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx b/models/dist/inference/layout/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx new file mode 120000 index 0000000..9034c95 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/model_eynollah_page_extraction_20250915.onnx b/models/dist/inference/layout/models_eynollah/model_eynollah_page_extraction_20250915.onnx new file mode 120000 index 0000000..688b314 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/model_eynollah_page_extraction_20250915.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_page_extraction_20250915.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/model_eynollah_reading_order_20250824.onnx b/models/dist/inference/layout/models_eynollah/model_eynollah_reading_order_20250824.onnx new file mode 120000 index 0000000..2db28c0 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/model_eynollah_reading_order_20250824.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_reading_order_20250824.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx b/models/dist/inference/layout/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx new file mode 120000 index 0000000..c5828f6 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/modelens_full_lay_1__4_3_091124.onnx b/models/dist/inference/layout/models_eynollah/modelens_full_lay_1__4_3_091124.onnx new file mode 120000 index 0000000..758c260 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/modelens_full_lay_1__4_3_091124.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_full_lay_1__4_3_091124.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/modelens_table_0t4_201124.onnx b/models/dist/inference/layout/models_eynollah/modelens_table_0t4_201124.onnx new file mode 120000 index 0000000..5768a79 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/modelens_table_0t4_201124.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_table_0t4_201124.onnx \ No newline at end of file diff --git a/models/dist/inference/layout/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx b/models/dist/inference/layout/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx new file mode 120000 index 0000000..2b842b4 --- /dev/null +++ b/models/dist/inference/layout/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/modelens_textline_0_1__2_4_16092024.onnx \ No newline at end of file diff --git a/models/dist/inference/ocr/models_eynollah/characters_org.txt b/models/dist/inference/ocr/models_eynollah/characters_org.txt new file mode 120000 index 0000000..9b62c85 --- /dev/null +++ b/models/dist/inference/ocr/models_eynollah/characters_org.txt @@ -0,0 +1 @@ +../../../../models_eynollah/characters_org.txt \ No newline at end of file diff --git a/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 b/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 new file mode 120000 index 0000000..9875a6a --- /dev/null +++ b/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 \ No newline at end of file diff --git a/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 b/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 new file mode 120000 index 0000000..0982ba7 --- /dev/null +++ b/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 \ No newline at end of file diff --git a/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_trocr_20250919 b/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_trocr_20250919 new file mode 120000 index 0000000..d5dc0b1 --- /dev/null +++ b/models/dist/inference/ocr/models_eynollah/model_eynollah_ocr_trocr_20250919 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_trocr_20250919 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/characters_org.txt b/models/dist/training/all/models_eynollah/characters_org.txt new file mode 120000 index 0000000..9b62c85 --- /dev/null +++ b/models/dist/training/all/models_eynollah/characters_org.txt @@ -0,0 +1 @@ +../../../../models_eynollah/characters_org.txt \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-binarization-hybrid_20230504 b/models/dist/training/all/models_eynollah/eynollah-binarization-hybrid_20230504 new file mode 120000 index 0000000..0d90379 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-binarization-hybrid_20230504 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization-hybrid_20230504 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-binarization_20210309 b/models/dist/training/all/models_eynollah/eynollah-binarization_20210309 new file mode 120000 index 0000000..a043883 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-binarization_20210309 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization_20210309 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-binarization_20210425 b/models/dist/training/all/models_eynollah/eynollah-binarization_20210425 new file mode 120000 index 0000000..dc3ebc7 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-binarization_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization_20210425 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-column-classifier_20210425 b/models/dist/training/all/models_eynollah/eynollah-column-classifier_20210425 new file mode 120000 index 0000000..66c19d1 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-column-classifier_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-column-classifier_20210425 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-enhancement_20210425 b/models/dist/training/all/models_eynollah/eynollah-enhancement_20210425 new file mode 120000 index 0000000..1287961 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-enhancement_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-enhancement_20210425 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-main-regions-aug-rotation_20210425 b/models/dist/training/all/models_eynollah/eynollah-main-regions-aug-rotation_20210425 new file mode 120000 index 0000000..d66338f --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-main-regions-aug-rotation_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-aug-rotation_20210425 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-main-regions-aug-scaling_20210425 b/models/dist/training/all/models_eynollah/eynollah-main-regions-aug-scaling_20210425 new file mode 120000 index 0000000..9e11f44 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-main-regions-aug-scaling_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-aug-scaling_20210425 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-main-regions-ensembled_20210425 b/models/dist/training/all/models_eynollah/eynollah-main-regions-ensembled_20210425 new file mode 120000 index 0000000..19fc248 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-main-regions-ensembled_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-ensembled_20210425 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-main-regions_20220314 b/models/dist/training/all/models_eynollah/eynollah-main-regions_20220314 new file mode 120000 index 0000000..6d3d6ca --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-main-regions_20220314 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions_20220314 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 b/models/dist/training/all/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 new file mode 120000 index 0000000..81d5930 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/eynollah-tables_20210319 b/models/dist/training/all/models_eynollah/eynollah-tables_20210319 new file mode 120000 index 0000000..56d65a6 --- /dev/null +++ b/models/dist/training/all/models_eynollah/eynollah-tables_20210319 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-tables_20210319 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 b/models/dist/training/all/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 new file mode 120000 index 0000000..1e8a14b --- /dev/null +++ b/models/dist/training/all/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_cnnrnn_20250930 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 b/models/dist/training/all/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 new file mode 120000 index 0000000..0982ba7 --- /dev/null +++ b/models/dist/training/all/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/model_eynollah_ocr_trocr_20250919 b/models/dist/training/all/models_eynollah/model_eynollah_ocr_trocr_20250919 new file mode 120000 index 0000000..d5dc0b1 --- /dev/null +++ b/models/dist/training/all/models_eynollah/model_eynollah_ocr_trocr_20250919 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_trocr_20250919 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/model_eynollah_page_extraction_20250915 b/models/dist/training/all/models_eynollah/model_eynollah_page_extraction_20250915 new file mode 120000 index 0000000..21ee2d8 --- /dev/null +++ b/models/dist/training/all/models_eynollah/model_eynollah_page_extraction_20250915 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_page_extraction_20250915 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/model_eynollah_reading_order_20250824 b/models/dist/training/all/models_eynollah/model_eynollah_reading_order_20250824 new file mode 120000 index 0000000..232e767 --- /dev/null +++ b/models/dist/training/all/models_eynollah/model_eynollah_reading_order_20250824 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_reading_order_20250824 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 b/models/dist/training/all/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 new file mode 120000 index 0000000..5f4ff0b --- /dev/null +++ b/models/dist/training/all/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/modelens_full_lay_1__4_3_091124 b/models/dist/training/all/models_eynollah/modelens_full_lay_1__4_3_091124 new file mode 120000 index 0000000..e5f4d77 --- /dev/null +++ b/models/dist/training/all/models_eynollah/modelens_full_lay_1__4_3_091124 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_full_lay_1__4_3_091124 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/modelens_table_0t4_201124 b/models/dist/training/all/models_eynollah/modelens_table_0t4_201124 new file mode 120000 index 0000000..60c21ae --- /dev/null +++ b/models/dist/training/all/models_eynollah/modelens_table_0t4_201124 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_table_0t4_201124 \ No newline at end of file diff --git a/models/dist/training/all/models_eynollah/modelens_textline_0_1__2_4_16092024 b/models/dist/training/all/models_eynollah/modelens_textline_0_1__2_4_16092024 new file mode 120000 index 0000000..ec2e51c --- /dev/null +++ b/models/dist/training/all/models_eynollah/modelens_textline_0_1__2_4_16092024 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_textline_0_1__2_4_16092024 \ No newline at end of file diff --git a/models/dist/training/extra/models_eynollah/eynollah-binarization_20210309 b/models/dist/training/extra/models_eynollah/eynollah-binarization_20210309 new file mode 120000 index 0000000..a043883 --- /dev/null +++ b/models/dist/training/extra/models_eynollah/eynollah-binarization_20210309 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization_20210309 \ No newline at end of file diff --git a/models/dist/training/extra/models_eynollah/eynollah-binarization_20210425 b/models/dist/training/extra/models_eynollah/eynollah-binarization_20210425 new file mode 120000 index 0000000..dc3ebc7 --- /dev/null +++ b/models/dist/training/extra/models_eynollah/eynollah-binarization_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-binarization_20210425 \ No newline at end of file diff --git a/models/dist/training/extra/models_eynollah/eynollah-main-regions-aug-scaling_20210425 b/models/dist/training/extra/models_eynollah/eynollah-main-regions-aug-scaling_20210425 new file mode 120000 index 0000000..9e11f44 --- /dev/null +++ b/models/dist/training/extra/models_eynollah/eynollah-main-regions-aug-scaling_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-aug-scaling_20210425 \ No newline at end of file diff --git a/models/dist/training/extra/models_eynollah/eynollah-tables_20210319 b/models/dist/training/extra/models_eynollah/eynollah-tables_20210319 new file mode 120000 index 0000000..56d65a6 --- /dev/null +++ b/models/dist/training/extra/models_eynollah/eynollah-tables_20210319 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-tables_20210319 \ No newline at end of file diff --git a/models/dist/training/extra/models_eynollah/modelens_textline_0_1__2_4_16092024 b/models/dist/training/extra/models_eynollah/modelens_textline_0_1__2_4_16092024 new file mode 120000 index 0000000..ec2e51c --- /dev/null +++ b/models/dist/training/extra/models_eynollah/modelens_textline_0_1__2_4_16092024 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_textline_0_1__2_4_16092024 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-column-classifier_20210425 b/models/dist/training/layout/models_eynollah/eynollah-column-classifier_20210425 new file mode 120000 index 0000000..66c19d1 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-column-classifier_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-column-classifier_20210425 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-enhancement_20210425 b/models/dist/training/layout/models_eynollah/eynollah-enhancement_20210425 new file mode 120000 index 0000000..1287961 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-enhancement_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-enhancement_20210425 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-main-regions-aug-rotation_20210425 b/models/dist/training/layout/models_eynollah/eynollah-main-regions-aug-rotation_20210425 new file mode 120000 index 0000000..d66338f --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-main-regions-aug-rotation_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-aug-rotation_20210425 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-main-regions-aug-scaling_20210425 b/models/dist/training/layout/models_eynollah/eynollah-main-regions-aug-scaling_20210425 new file mode 120000 index 0000000..9e11f44 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-main-regions-aug-scaling_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-aug-scaling_20210425 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-main-regions-ensembled_20210425 b/models/dist/training/layout/models_eynollah/eynollah-main-regions-ensembled_20210425 new file mode 120000 index 0000000..19fc248 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-main-regions-ensembled_20210425 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions-ensembled_20210425 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-main-regions_20220314 b/models/dist/training/layout/models_eynollah/eynollah-main-regions_20220314 new file mode 120000 index 0000000..6d3d6ca --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-main-regions_20220314 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions_20220314 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 b/models/dist/training/layout/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 new file mode 120000 index 0000000..81d5930 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 @@ -0,0 +1 @@ +../../../../models_eynollah/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/model_eynollah_page_extraction_20250915 b/models/dist/training/layout/models_eynollah/model_eynollah_page_extraction_20250915 new file mode 120000 index 0000000..3e1a782 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/model_eynollah_page_extraction_20250915 @@ -0,0 +1 @@ +../../../../reloaded/models_eynollah/model_eynollah_page_extraction_20250915 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/model_eynollah_reading_order_20250824 b/models/dist/training/layout/models_eynollah/model_eynollah_reading_order_20250824 new file mode 120000 index 0000000..232e767 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/model_eynollah_reading_order_20250824 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_reading_order_20250824 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 b/models/dist/training/layout/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 new file mode 120000 index 0000000..5f4ff0b --- /dev/null +++ b/models/dist/training/layout/models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_e_l_all_sp_0_1_2_3_4_171024 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/modelens_full_lay_1__4_3_091124 b/models/dist/training/layout/models_eynollah/modelens_full_lay_1__4_3_091124 new file mode 120000 index 0000000..e5f4d77 --- /dev/null +++ b/models/dist/training/layout/models_eynollah/modelens_full_lay_1__4_3_091124 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_full_lay_1__4_3_091124 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/modelens_table_0t4_201124 b/models/dist/training/layout/models_eynollah/modelens_table_0t4_201124 new file mode 120000 index 0000000..60c21ae --- /dev/null +++ b/models/dist/training/layout/models_eynollah/modelens_table_0t4_201124 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_table_0t4_201124 \ No newline at end of file diff --git a/models/dist/training/layout/models_eynollah/modelens_textline_0_1__2_4_16092024 b/models/dist/training/layout/models_eynollah/modelens_textline_0_1__2_4_16092024 new file mode 120000 index 0000000..ec2e51c --- /dev/null +++ b/models/dist/training/layout/models_eynollah/modelens_textline_0_1__2_4_16092024 @@ -0,0 +1 @@ +../../../../models_eynollah/modelens_textline_0_1__2_4_16092024 \ No newline at end of file diff --git a/models/dist/training/ocr/models_eynollah/characters_org.txt b/models/dist/training/ocr/models_eynollah/characters_org.txt new file mode 120000 index 0000000..9b62c85 --- /dev/null +++ b/models/dist/training/ocr/models_eynollah/characters_org.txt @@ -0,0 +1 @@ +../../../../models_eynollah/characters_org.txt \ No newline at end of file diff --git a/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 b/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 new file mode 120000 index 0000000..1e8a14b --- /dev/null +++ b/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_cnnrnn_20250930 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_cnnrnn_20250930 \ No newline at end of file diff --git a/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 b/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 new file mode 120000 index 0000000..0982ba7 --- /dev/null +++ b/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_cnnrnn__degraded_20250805 \ No newline at end of file diff --git a/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_trocr_20250919 b/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_trocr_20250919 new file mode 120000 index 0000000..d5dc0b1 --- /dev/null +++ b/models/dist/training/ocr/models_eynollah/model_eynollah_ocr_trocr_20250919 @@ -0,0 +1 @@ +../../../../models_eynollah/model_eynollah_ocr_trocr_20250919 \ No newline at end of file diff --git a/models/scripts/convert.sh b/models/scripts/convert.sh new file mode 100644 index 0000000..225785f --- /dev/null +++ b/models/scripts/convert.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash + +for i in */models_eynollah/*;do + rslv=$(readlink $i|sed 's,models_eynollah,reloaded/models_eynollah,'); + ln -srf $rslv $i; +done diff --git a/models/scripts/plot.py b/models/scripts/plot.py new file mode 100644 index 0000000..2543860 --- /dev/null +++ b/models/scripts/plot.py @@ -0,0 +1,62 @@ +import os +import sys +import more_itertools +import numpy as np +import argparse + +parser = argparse.ArgumentParser( + prog='plot', + description='Print metrics', + epilog='...') +parser.add_argument('stem1', default='nohup.out.eynollah-206-noautosize-keepseps-refactoring9-crop-after2') +parser.add_argument('stem2', default='eval-gt2-onlyfg-pixel-f1-v0.5-206-noautosize-keepseps-refactoring9-crop-after2') + +def main(): + args = parser.parse_args() + + prefix = os.path.dirname(sys.argv[0]) + files = [path for path in os.listdir(prefix) + if path.startswith(args.stem1) + and path.endswith('.txt')] + + for path in files: + name = path[len(args.stem1):] + metric = name.split('.')[-2] + series = name[:-(len(metric) + 5)] + path = os.path.join(prefix, path) + + if metric == 'peak-vram': + cat, peak = np.loadtxt(path, unpack=True, dtype=np.dtype([("category", "U15"), ("peak", float)])) + peak //= 1024 + peak //= 1024 + print(metric, f"{series: <52}", dict(zip(cat.tolist(), peak.tolist()))) + else: + val = np.loadtxt(path) + print(metric, f"{series: <52}", {'µ': float(np.round(np.mean(val), 2)), + 'M': float(np.round(np.median(val), 2)), + 'min': float(np.min(val)), + 'max': float(np.max(val))}) + + + files = [path for path in os.listdir(prefix) + if path.startswith(args.stem2) + and path.endswith('.log')] + + + for path in files: + name = path[len(args.stem2):] + series = name[:-4] + path = os.path.join(prefix, path) + with open(path, 'r') as fd: + lines = fd.readlines() + result_v0_5 = {} + result_v206 = {} + for metric, v0_5, v206 in more_itertools.chunked(lines, n=3, strict=True): + metric = metric[len(prefix) + 1:-1] + result_v0_5[metric] = float(np.round(float(v0_5.strip()), 3)) + result_v206[metric] = float(np.round(float(v206.strip()), 3)) + print("pixel-f1", f"{series: <52}", result_v206) + + print("pixel-f1", "v0.5 ", result_v0_5) + + diff --git a/models/scripts/run.sh b/models/scripts/run.sh new file mode 100644 index 0000000..79cd4e6 --- /dev/null +++ b/models/scripts/run.sh @@ -0,0 +1,34 @@ +for variant in reloaded.*; do + +for dir in Korrigierte_Layout_GT/*/; do test $dir = ${dir/.} || continue; out=${dir%/}.eynollah-206-noautosize; mkdir -p $out; nohup /usr/bin +/time eynollah -D GPU0 -m $variant layout --dir_in $dir -o + $out -fl -H 1 -O || break; done + +logfile=Korrigierte_Layout_GT/nohup.out.eynollah-206-noautosize-$variant + +mv nohup.out $logfile + +fgrep peaked $logfile | sort -u | cut -d\ -f2,5 > $logfile.peak-vram.txt +fgrep initialization $logfile | sed "s/^.* (\([0-9.]*\)s)/\1/" > $logfile.secs-init.txt +fgrep "Job done in" $logfile | sed "s/^.* //;s/s$//" > $logfile.secs-jobs.txt +fgrep %CPU $logfile | sed "s/^.* \([0-9]*\)%CPU.*/\1/" > $logfile.prct-cpu.txt + +for dir in Korrigierte_Layout_GT/*.eynollah*/; do + test $dir = ${dir/-eval} || continue; + test $dir != ${dir/eynollah-206} || continue; + mets=${dir%.eynollah*}.mets.xml; + dir=$(basename $dir); + gt=${dir%.eynollah*}.gt2; + echo $mets: $dir; + ocrd-segment-evaluate -m $mets -I $gt,$dir -O $dir-eval,$dir-eval2 -P level-of-operation region -P ignore-subtype true -P only-fg true --overwrite; +done + +for mets in Korrigierte_Layout_GT/*.mets.xml; do dir=${mets%.mets.xml}; echo $dir; jq '."by-category".TextRegion."pixel-f1".avg' $dir.eynollah-v0.5-eval/*.json $dir.eynollah-206-noautosize-eval/*.json; done > Korrigierte_Layout_GT/eval-gt2-onlyfg-pixel-f1-v0.5-206-noautosize-$variant.log + + +done + + + +python Korrigierte_Layout_GT/plot.py | sort -V + diff --git a/models/scripts/zenodo_upload.sh b/models/scripts/zenodo_upload.sh new file mode 100755 index 0000000..da67298 --- /dev/null +++ b/models/scripts/zenodo_upload.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# +# Upload big files to Zenodo. +# +# usage: ./zenodo_upload.sh [deposition id] [filename] [--verbose|-v] +# +# Taken from https://github.com/jhpoelen/zenodo-upload +# +# Copyright (c) 2019 Jorrit Poelen +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# + +set -e + +VERBOSE=0 +if [ "$3" == "--verbose" ] || [ "$3" == "-v" ]; then + VERBOSE=1 +fi + +# strip deposition url prefix if provided; see https://github.com/jhpoelen/zenodo-upload/issues/2#issuecomment-797657717 +DEPOSITION=$( echo $1 | sed 's+^http[s]*://zenodo.org/deposit/++g' ) +FILEPATH="$2" +FILENAME=$(echo $FILEPATH | sed 's+.*/++g') +FILENAME=${FILENAME// /%20} +ZENODO_ENDPOINT=${ZENODO_ENDPOINT:-https://zenodo.org} + +DEPOSIT_URL=${ZENODO_ENDPOINT}/api/deposit/depositions/"$DEPOSITION"?access_token="$ZENODO_TOKEN" +BUCKET=$(curl $DEPOSIT_URL | jq --raw-output .links.bucket) + +if [ "$VERBOSE" -eq 1 ]; then + echo "Deposition ID: $DEPOSITION" + echo "File path: $FILEPATH" + echo "File name: $FILENAME" + echo "Bucket URL: $BUCKET" + echo "Upload URL: $BUCKET/$FILENAME" + echo "Deposit URL: $DEPOSIT_URL" + echo "Uploading file..." +fi + +set -x +curl --progress-bar \ + --retry 5 \ + --retry-delay 5 \ + -o ${FILENAME}.upload.log \ + -H "Authorization: Bearer $ZENODO_TOKEN" \ + --upload-file "$FILEPATH" \ + "$BUCKET/$FILENAME" + +# curl -i \ +# -H "Authorization: Bearer $ZENODO_TOKEN" \ +# -X POST \ +# -F name=$FILENAME \ +# -F file=@$FILEPATH \ +# https://zenodo.org/api/deposit/depositions/$DEPOSITION/files diff --git a/pyproject.toml b/pyproject.toml index e6821a5..fdd41db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ source = ["eynollah"] [tool.ruff] line-length = 120 +include = ["pyproject.toml", "src/eynollah/**/*.py"] [tool.ruff.lint] ignore = [ diff --git a/requirements-ocr.txt b/requirements-ocr.txt index dad26f4..0825e18 100644 --- a/requirements-ocr.txt +++ b/requirements-ocr.txt @@ -1,3 +1,5 @@ -torch +torch < 2.11 # avoid pull CUDA 13 (which will clash with 12) transformers <= 4.30.2 ; python_version < '3.10' transformers >= 5 ; python_version >= '3.10' +tensorflow[and-cuda] +tf-keras # avoid keras 3 (also needs TF_USE_LEGACY_KERAS=1) diff --git a/requirements.txt b/requirements.txt index d79853f..a707407 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,7 @@ # ocrd includes opencv, numpy, shapely, click ocrd >= 3.3.0 -numpy < 2.0 +onnxruntime-gpu[cuda,cudnn] # w/ .onnx models +tensorrt_cu12 < 11 # 11 incompatible with CUDA libs from onnxruntime-gpu[cuda,cudnn] scikit-learn >= 0.23.2 -tensorflow -tf-keras # avoid keras 3 (also needs TF_USE_LEGACY_KERAS=1) -numba <= 0.58.1 scikit-image tabulate diff --git a/src/eynollah/Amiri-Regular.ttf b/src/eynollah/Amiri-Regular.ttf new file mode 100644 index 0000000..df5e1df Binary files /dev/null and b/src/eynollah/Amiri-Regular.ttf differ diff --git a/src/eynollah/Charis-Regular.ttf b/src/eynollah/Charis-Regular.ttf deleted file mode 100644 index a4e75a4..0000000 Binary files a/src/eynollah/Charis-Regular.ttf and /dev/null differ diff --git a/src/eynollah/cli/__init__.py b/src/eynollah/cli/__init__.py index 43ed046..1584fa5 100644 --- a/src/eynollah/cli/__init__.py +++ b/src/eynollah/cli/__init__.py @@ -1,7 +1,3 @@ -# NOTE: For predictable order of imports of torch/shapely/tensorflow -# this must be the first import of the CLI! -from ..eynollah_imports import imported_libs - from .cli import main from .cli_binarize import binarize_cli from .cli_enhance import enhance_cli diff --git a/src/eynollah/cli/cli.py b/src/eynollah/cli/cli.py index ace3f1c..2a4c8d1 100644 --- a/src/eynollah/cli/cli.py +++ b/src/eynollah/cli/cli.py @@ -15,6 +15,7 @@ class EynollahCliCtx: Holds options relevant for all eynollah subcommands """ model_zoo: EynollahModelZoo + device: str = '' log_level : Union[str, None] = 'INFO' @@ -35,6 +36,11 @@ class EynollahCliCtx: type=(str, str, str), multiple=True, ) +@click.option( + "--device", + "-D", + help="placement of computations in predictors for each model type; if none (by default), will try to use first available GPU or fall back to CPU; set string to force using a device (e.g. 'GPU0', 'GPU1' or 'CPU'). Can also be a comma-separated list of model category to device mappings (e.g. 'col_classifier:CPU,page:GPU0,*:GPU1')", +) @click.option( "--log_level", "-l", @@ -42,7 +48,7 @@ class EynollahCliCtx: help="Override log level globally to this", ) @click.pass_context -def main(ctx, model_basedir, model_overrides, log_level): +def main(ctx, model_basedir, model_overrides, device, log_level): """ eynollah - Document Layout Analysis, Image Enhancement, OCR """ @@ -58,6 +64,7 @@ def main(ctx, model_basedir, model_overrides, log_level): # Initialize CLI context ctx.obj = EynollahCliCtx( model_zoo=model_zoo, + device=device, log_level=log_level, ) diff --git a/src/eynollah/cli/cli_binarize.py b/src/eynollah/cli/cli_binarize.py index f0e56f5..82209be 100644 --- a/src/eynollah/cli/cli_binarize.py +++ b/src/eynollah/cli/cli_binarize.py @@ -1,6 +1,8 @@ import click -@click.command() +@click.command(context_settings=dict( + help_option_names=['-h', '--help'], + show_default=True)) @click.option( '--patches/--no-patches', default=True, @@ -31,11 +33,6 @@ import click help="overwrite (instead of skipping) if output xml exists", is_flag=True, ) -@click.option( - "--device", - "-D", - help="placement of computations in predictors for each model type; if none (by default), will try to use first available GPU or fall back to CPU; set string to force using a device (e.g. 'GPU0', 'GPU1' or 'CPU'). Can also be a comma-separated list of model category to device mappings (e.g. 'col_classifier:CPU,page:GPU0,*:GPU1')", -) @click.pass_context def binarize_cli( ctx, @@ -44,14 +41,14 @@ def binarize_cli( dir_in, output, overwrite, - device, ): """ Binarize images with a ML model """ from ..sbb_binarize import SbbBinarizer assert bool(input_image) != bool(dir_in), "Either -i (single input) or -di (directory) must be provided, but not both." - binarizer = SbbBinarizer(model_zoo=ctx.obj.model_zoo, device=device) + binarizer = SbbBinarizer(model_zoo=ctx.obj.model_zoo, + device=ctx.obj.device) binarizer.run( image_filename=input_image, use_patches=patches, diff --git a/src/eynollah/cli/cli_enhance.py b/src/eynollah/cli/cli_enhance.py index 517e1e8..bcb8263 100644 --- a/src/eynollah/cli/cli_enhance.py +++ b/src/eynollah/cli/cli_enhance.py @@ -1,6 +1,8 @@ import click -@click.command() +@click.command(context_settings=dict( + help_option_names=['-h', '--help'], + show_default=True)) @click.option( "--image", "-i", @@ -46,13 +48,8 @@ import click is_flag=True, help="save the enhanced image in original image size", ) -@click.option( - "--device", - "-D", - help="placement of computations in predictors for each model type; if none (by default), will try to use first available GPU or fall back to CPU; set string to force using a device (e.g. 'GPU0', 'GPU1' or 'CPU'). Can also be a comma-separated list of model category to device mappings (e.g. 'col_classifier:CPU,page:GPU0,*:GPU1')", -) @click.pass_context -def enhance_cli(ctx, image, out, overwrite, dir_in, num_col_upper, num_col_lower, save_org_scale, device): +def enhance_cli(ctx, image, out, overwrite, dir_in, num_col_upper, num_col_lower, save_org_scale): """ Enhance image """ @@ -60,10 +57,10 @@ def enhance_cli(ctx, image, out, overwrite, dir_in, num_col_upper, num_col_lower from ..image_enhancer import Enhancer enhancer = Enhancer( model_zoo=ctx.obj.model_zoo, + device=ctx.obj.device, num_col_upper=num_col_upper, num_col_lower=num_col_lower, save_org_scale=save_org_scale, - device=device, ) enhancer.run(overwrite=overwrite, dir_in=dir_in, diff --git a/src/eynollah/cli/cli_extract_images.py b/src/eynollah/cli/cli_extract_images.py index 0add5b5..acd31f1 100644 --- a/src/eynollah/cli/cli_extract_images.py +++ b/src/eynollah/cli/cli_extract_images.py @@ -1,6 +1,8 @@ import click -@click.command() +@click.command(context_settings=dict( + help_option_names=['-h', '--help'], + show_default=True)) @click.option( "--image", "-i", @@ -30,36 +32,40 @@ import click @click.option( "--save_images", "-si", - help="if a directory is given, images in documents will be cropped and saved there", + help="if a directory is given, cropped images of pages will be saved there", type=click.Path(exists=True, file_okay=False), ) @click.option( - "--enable-plotting/--disable-plotting", - "-ep/-noep", + "--enable-plotting", + "-ep", is_flag=True, - help="If set, will plot intermediary files and images", + help="plot intermediary diagnostic images to files", ) @click.option( - "--input_binary/--input-RGB", - "-ib/-irgb", + "--input_binary", + "-ib", is_flag=True, - help="In general, eynollah uses RGB as input but if the input document is very dark, very bright or for any other reason you can turn on input binarization. When this flag is set, eynollah will binarize the RGB input document, you should always provide RGB images to eynollah.", + help="In general, eynollah uses RGB as input, but if the input document is very dark, very bright or for any other reason you can turn on internal binarization here. When set, eynollah will binarize the RGB input document first.", ) @click.option( - "--ignore_page_extraction/--extract_page_included", - "-ipe/-epi", + "--ignore_page_extraction", + "-ipe", is_flag=True, - help="if this parameter set to true, this tool would ignore page extraction", + help="ignore page extraction (cropping via page frame detection model)", ) @click.option( "--num_col_upper", "-ncu", - help="lower limit of columns in document image", + default=0, + type=click.IntRange(min=0), + help="lower limit of columns in document image; 0 means autodetected from model", ) @click.option( "--num_col_lower", "-ncl", - help="upper limit of columns in document image", + default=0, + type=click.IntRange(min=0), + help="upper limit of columns in document image; 0 means autodetected from model", ) @click.pass_context def extract_images_cli( diff --git a/src/eynollah/cli/cli_layout.py b/src/eynollah/cli/cli_layout.py index 417b202..0a083d5 100644 --- a/src/eynollah/cli/cli_layout.py +++ b/src/eynollah/cli/cli_layout.py @@ -172,11 +172,6 @@ import click type=click.FloatRange(min=0), help="abort when number of failed images exceeds this value (if >=1) or ratio of failed over total images exceeds this value (if <1); 0 means ignore failures", ) -@click.option( - "--device", - "-D", - help="placement of computations in predictors for each model type; if none (by default), will try to use first available GPU or fall back to CPU; set string to force using a device (e.g. 'GPU0', 'GPU1' or 'CPU'). Can also be a comma-separated list of model category to device mappings (e.g. 'col_classifier:CPU,page:GPU0,*:GPU1')", -) @click.pass_context def layout_cli( ctx, @@ -207,7 +202,6 @@ def layout_cli( ignore_page_extraction, num_jobs, halt_fail, - device, ): """ Detect Layout (with optional image enhancement and reading order detection) @@ -223,7 +217,7 @@ def layout_cli( assert bool(image) != bool(dir_in), "Either -i (single input) or -di (directory) must be provided, but not both." eynollah = Eynollah( model_zoo=ctx.obj.model_zoo, - device=device, + device=ctx.obj.device, enable_plotting=enable_plotting, allow_enhancement=allow_enhancement, curved_line=curved_line, diff --git a/src/eynollah/cli/cli_ocr.py b/src/eynollah/cli/cli_ocr.py index 406af61..09c88db 100644 --- a/src/eynollah/cli/cli_ocr.py +++ b/src/eynollah/cli/cli_ocr.py @@ -1,6 +1,8 @@ import click -@click.command() +@click.command(context_settings=dict( + help_option_names=['-h', '--help'], + show_default=True)) @click.option( "--image", "-i", @@ -16,7 +18,7 @@ import click @click.option( "--dir_in_bin", "-dib", - help=("directory of binarized images (in addition to --dir_in for RGB images; filename stems must match the RGB image files, with '.png' \n Perform prediction using both RGB and binary images. (This does not necessarily improve results, however it may be beneficial for certain document images."), + help=("directory of binarized images (in addition to --dir_in for RGB images; filename stems must match the RGB image files, with '.png'. \n Perform prediction using both RGB and binary images. (This may improve results for certain document images.)"), type=click.Path(exists=True, file_okay=False), ) @click.option( @@ -45,27 +47,45 @@ import click help="overwrite (instead of skipping) if output xml exists", is_flag=True, ) +@click.option( + "--num-jobs", + "-j", + default=0, + type=click.IntRange(min=0), + help="number of parallel images to process (for --dir_in mode; also helps better utilise GPU if available); 0 means based on autodetected number of processor cores", +) +@click.option( + "--halt-fail", + "-H", + default=0, + type=click.FloatRange(min=0), + help="abort when number of failed images exceeds this value (if >=1) or ratio of failed over total images exceeds this value (if <1); 0 means ignore failures", +) @click.option( "--tr_ocr", - "-trocr/-notrocr", + "-trocr", is_flag=True, - help="if this parameter set to true, transformer ocr will be applied, otherwise cnn_rnn model.", + help="use transformer OCR (instead of classic CNN-RNN) model", ) @click.option( "--do_not_mask_with_textline_contour", - "-nmtc/-mtc", + "-nmtc", is_flag=True, - help="if this parameter set to true, cropped textline images will not be masked with textline contour.", + help="skip masking each cropped textline image with its corresponding textline contour", ) @click.option( "--batch_size", "-bs", + default=0, + type=click.IntRange(min=0), help="number of inference batch size. Default b_s for trocr and cnn_rnn models are 2 and 8 respectively", ) @click.option( "--min_conf_value_of_textline_text", "-min_conf", - help="minimum OCR confidence value. Text lines with a confidence value lower than this threshold will not be included in the output XML file.", + default=0.3, + type=click.FloatRange(min=0.0, max=1.0), + help="minimum OCR confidence threshold. Text lines with a lower confidence value will not be included in the output XML file.", ) @click.pass_context def ocr_cli( @@ -77,6 +97,8 @@ def ocr_cli( out, dir_out_image_text, overwrite, + num_jobs, + halt_fail, tr_ocr, do_not_mask_with_textline_contour, batch_size, @@ -85,14 +107,16 @@ def ocr_cli( """ Recognize text with a CNN/RNN or transformer ML model. """ - assert bool(image) ^ bool(dir_in), "Either -i (single image) or -di (directory) must be provided, but not both." + assert bool(image) != bool(dir_in), "Either -i (single image) or -di (directory) must be provided, but not both." from ..eynollah_ocr import Eynollah_ocr eynollah_ocr = Eynollah_ocr( model_zoo=ctx.obj.model_zoo, + device=ctx.obj.device, tr_ocr=tr_ocr, do_not_mask_with_textline_contour=do_not_mask_with_textline_contour, batch_size=batch_size, - min_conf_value_of_textline_text=min_conf_value_of_textline_text) + min_conf_value_of_textline_text=min_conf_value_of_textline_text, + ) eynollah_ocr.run(overwrite=overwrite, dir_in=dir_in, dir_in_bin=dir_in_bin, @@ -100,4 +124,6 @@ def ocr_cli( dir_xmls=dir_xmls, dir_out_image_text=dir_out_image_text, dir_out=out, + num_jobs=num_jobs, + halt_fail=halt_fail, ) diff --git a/src/eynollah/cli/cli_readingorder.py b/src/eynollah/cli/cli_readingorder.py index 0f44b7f..ac52e38 100644 --- a/src/eynollah/cli/cli_readingorder.py +++ b/src/eynollah/cli/cli_readingorder.py @@ -1,6 +1,8 @@ import click -@click.command() +@click.command(context_settings=dict( + help_option_names=['-h', '--help'], + show_default=True)) @click.option( "--input", "-i", @@ -25,9 +27,10 @@ def readingorder_cli(ctx, input, dir_in, out): """ Generate ReadingOrder with a ML model """ - from ..mb_ro_on_layout import machine_based_reading_order_on_layout + from ..mb_ro_on_layout import Reorder assert bool(input) != bool(dir_in), "Either -i (single input) or -di (directory) must be provided, but not both." - orderer = machine_based_reading_order_on_layout(model_zoo=ctx.obj.model_zoo) + orderer = Reorder(model_zoo=ctx.obj.model_zoo, + device=ctx.obj.device) orderer.run(xml_filename=input, dir_in=dir_in, dir_out=out, diff --git a/src/eynollah/extract_images.py b/src/eynollah/extract_images.py index 7a7e3f6..40476a3 100644 --- a/src/eynollah/extract_images.py +++ b/src/eynollah/extract_images.py @@ -9,7 +9,6 @@ import os import time from typing import Optional from pathlib import Path -import tensorflow as tf import numpy as np import cv2 @@ -64,12 +63,6 @@ class EynollahImageExtractor(Eynollah): t_start = time.time() - 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)") diff --git a/src/eynollah/eynollah.py b/src/eynollah/eynollah.py index c632941..becf802 100644 --- a/src/eynollah/eynollah.py +++ b/src/eynollah/eynollah.py @@ -276,7 +276,7 @@ class Eynollah: img_new = np.copy(img) img_is_resized = False #elif conf_col < 0.8 and img_h_new >= 8000: - elif img_h_new >= 8000: + elif conf_col < 0.9 and img_h_new >= 8000: # don't upsample if too large img_new = np.copy(img) img_is_resized = False @@ -901,7 +901,7 @@ class Eynollah: if N > 1: mean_y_diff = np.median(diff_cy) mean_x_diff = np.median(diff_cx) - count_hor = np.count_nonzero(np.diff(w_h_textline) > 0) + count_hor = np.count_nonzero(np.diff(w_h_textline, axis=0) > 0) count_ver = N - count_hor else: mean_y_diff = 0 @@ -1148,7 +1148,6 @@ class Eynollah: boxes, textline_mask_tot ): - assert np.any(textline_mask_tot) self.logger.debug("enter do_order_of_regions") contours_only_text_parent = ensure_array(contours_only_text_parent) contours_only_text_parent_h = ensure_array(contours_only_text_parent_h) @@ -2143,7 +2142,8 @@ class Eynollah: conf_textregions=[0], ) self.logger.info("Basic processing complete") - writer.write_pagexml(pcgts) + if writer.pcgts is None: + writer.write_pagexml(pcgts) self.logger.info("Job done in %.1fs", time.time() - t0) return @@ -2222,7 +2222,8 @@ class Eynollah: cont_page=cont_page, polygons_seplines=[], ) - writer.write_pagexml(pcgts) + if writer.pcgts is None: + writer.write_pagexml(pcgts) self.logger.info("Job done in %.1fs", time.time() - t0) return @@ -2509,6 +2510,7 @@ class Eynollah: conf_tables=conf_tables, ) - writer.write_pagexml(pcgts) + if writer.pcgts is None: + writer.write_pagexml(pcgts) self.logger.info("Job done in %.1fs", time.time() - t0) return diff --git a/src/eynollah/eynollah_imports.py b/src/eynollah/eynollah_imports.py deleted file mode 100644 index 496406c..0000000 --- a/src/eynollah/eynollah_imports.py +++ /dev/null @@ -1,13 +0,0 @@ -""" -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() -import tensorflow.keras -from shapely import * -imported_libs = True -__all__ = ['imported_libs'] diff --git a/src/eynollah/eynollah_ocr.py b/src/eynollah/eynollah_ocr.py index 1b49077..b5bf6d8 100644 --- a/src/eynollah/eynollah_ocr.py +++ b/src/eynollah/eynollah_ocr.py @@ -1,33 +1,39 @@ # FIXME: fix all of those... # pyright: reportOptionalSubscript=false -from logging import Logger, getLogger +import logging +import logging.handlers from typing import List, Optional from pathlib import Path +from itertools import groupby import os import gc import math +import time from dataclasses import dataclass +import multiprocessing as mp +from concurrent.futures import ProcessPoolExecutor, as_completed import cv2 from cv2.typing import MatLike from xml.etree import ElementTree as ET from PIL import Image, ImageDraw import numpy as np -from eynollah.model_zoo import EynollahModelZoo -from eynollah.utils.font import get_font -from eynollah.utils.xml import etree_namespace_for_element_tag -try: - import torch -except ImportError: - torch = None +from ocrd_utils import polygon_from_points, xywh_from_polygon -from .utils import is_image_filename +from .eynollah import Eynollah +from .model_zoo import EynollahModelZoo +from .utils import ( + is_image_filename, + batched, + pairwise, +) +from .utils.font import get_font +from .utils.xml import etree_namespace_for_element_tag from .utils.resize import resize_image from .utils.utils_ocr import ( break_curved_line_into_small_pieces_and_then_merge, - decode_batch_predictions, fit_text_single_line, get_contours_and_bounding_boxes, get_orientation_moments, @@ -36,53 +42,64 @@ from .utils.utils_ocr import ( rotate_image_with_padding, ) + +_instance = None +def _set_instance(instance): + global _instance + _instance = instance +def _run_single(*args, **kwargs): + logq = kwargs.pop('logq') + # replace all inherited handlers with queue handler + logging.root.handlers.clear() + _instance.logger.parent.handlers.clear() + handler = logging.handlers.QueueHandler(logq) + logging.root.addHandler(handler) + return _instance.run_single(*args, **kwargs) + + # TODO: refine typing @dataclass class EynollahOcrResult: extracted_texts_merged: List - extracted_conf_value_merged: Optional[List] + extracted_confs_merged: List cropped_lines_region_indexer: List total_bb_coordinates:List -class Eynollah_ocr: +class Eynollah_ocr(Eynollah): def __init__( self, *, model_zoo: EynollahModelZoo, tr_ocr=False, - batch_size: Optional[int]=None, + batch_size: int=0, do_not_mask_with_textline_contour: bool=False, - min_conf_value_of_textline_text : Optional[float]=None, - logger: Optional[Logger]=None, + min_conf_value_of_textline_text : float=0.3, + logger: Optional[logging.Logger]=None, + device: str = '', ): self.tr_ocr = tr_ocr # masking for OCR and GT generation, relevant for skewed lines and bounding boxes self.do_not_mask_with_textline_contour = do_not_mask_with_textline_contour - self.logger = logger if logger else getLogger('eynollah.ocr') - self.model_zoo = model_zoo + self.logger = logger if logger else logging.getLogger('eynollah.ocr') - self.min_conf_value_of_textline_text = min_conf_value_of_textline_text if min_conf_value_of_textline_text else 0.3 - self.b_s = 2 if batch_size is None and tr_ocr else 8 if batch_size is None else batch_size + self.min_conf_value_of_textline_text = min_conf_value_of_textline_text + self.b_s = batch_size or (2 if tr_ocr else 64) - if tr_ocr: - self.model_zoo.load_models('trocr_processor') - self.model_zoo.load_models(['ocr', 'tr']) - self.model_zoo.get('ocr').to(self.device) + self.model_zoo = model_zoo + self.setup_models(device=device) + + def setup_models(self, device=''): + if self.tr_ocr: + self.model_zoo.load_models(('ocr', 'tr'), + device=device) else: - self.model_zoo.load_models('ocr') - self.model_zoo.load_models('num_to_char') - self.model_zoo.load_models('characters') - self.end_character = len(self.model_zoo.get('characters')) + 2 + self.model_zoo.load_models('ocr', + 'binarization', + device=device) @property def device(self): - assert torch - if torch.cuda.is_available(): - self.logger.info("Using GPU acceleration") - return torch.device("cuda:0") - else: - self.logger.info("Using CPU processing") - return torch.device("cpu") + return self.model_zoo.get('ocr').device def run_trocr( self, @@ -90,178 +107,73 @@ class Eynollah_ocr: img: MatLike, page_tree: ET.ElementTree, page_ns, - tr_ocr_input_height_and_width, ) -> EynollahOcrResult: - - total_bb_coordinates = [] - + total_bb_coordinates = [] cropped_lines = [] cropped_lines_region_indexer = [] cropped_lines_meging_indexing = [] - + + for n_region, region in enumerate(page_tree.getroot().iter('{%s}TextRegion' % page_ns)): + for n_line, line in enumerate(region.iter('{%s}TextLine' % page_ns)): + cropped_lines_region_indexer.append(n_region) + + coords = line.find('{%s}Coords' % page_ns) + if coords is None: + self.logger.warning("region '%s' line '%s' has no Coords", region.attrib['id'], line.attrib['id']) + continue + poly = np.array(polygon_from_points(coords.attrib['points'])).astype(int) + cont = poly[:, np.newaxis] + xywh = xywh_from_polygon(poly) + x, y, w, h = xywh['x'], xywh['y'], xywh['w'], xywh['h'] + + total_bb_coordinates.append([x, y, w, h]) + + img_crop = img[y: y + h, x: x + w] + if not self.do_not_mask_with_textline_contour: + mask_poly = np.zeros(img_crop.shape[:2], dtype=np.uint8) + mask_poly = cv2.fillPoly(mask_poly, pts=[cont - [x, y]], color=1) + img_crop[mask_poly == 0] = 255 # FIXME: or median color? + + if h > 0.1 * w: + cropped_lines.append(img_crop) + cropped_lines_meging_indexing.append(0) + else: + splited_images, _ = return_textlines_split_if_needed(img_crop, None) + if splited_images: + cropped_lines.append(splited_images[0]) + cropped_lines.append(splited_images[1]) + cropped_lines_meging_indexing.append(1) + cropped_lines_meging_indexing.append(-1) + else: + cropped_lines.append(img_crop) + cropped_lines_meging_indexing.append(0) + extracted_texts = [] - - indexer_text_region = 0 - indexer_b_s = 0 - - for nn in page_tree.getroot().iter(f'{{{page_ns}}}TextRegion'): - for child_textregion in nn: - if child_textregion.tag.endswith("TextLine"): - - for child_textlines in child_textregion: - if child_textlines.tag.endswith("Coords"): - cropped_lines_region_indexer.append(indexer_text_region) - p_h=child_textlines.attrib['points'].split(' ') - textline_coords = np.array( [ [int(x.split(',')[0]), - int(x.split(',')[1]) ] - for x in p_h] ) - x,y,w,h = cv2.boundingRect(textline_coords) - - total_bb_coordinates.append([x,y,w,h]) - - h2w_ratio = h/float(w) - - img_poly_on_img = np.copy(img) - mask_poly = np.zeros(img.shape) - mask_poly = cv2.fillPoly(mask_poly, pts=[textline_coords], color=(1, 1, 1)) - - mask_poly = mask_poly[y:y+h, x:x+w, :] - img_crop = img_poly_on_img[y:y+h, x:x+w, :] - img_crop[mask_poly==0] = 255 - - self.logger.debug("processing %d lines for '%s'", - len(cropped_lines), nn.attrib['id']) - if h2w_ratio > 0.1: - cropped_lines.append(resize_image(img_crop, - tr_ocr_input_height_and_width, - tr_ocr_input_height_and_width) ) - cropped_lines_meging_indexing.append(0) - indexer_b_s+=1 - if indexer_b_s==self.b_s: - imgs = cropped_lines[:] - cropped_lines = [] - indexer_b_s = 0 - - pixel_values_merged = self.model_zoo.get('trocr_processor')(imgs, return_tensors="pt").pixel_values - generated_ids_merged = self.model_zoo.get('ocr').generate( - pixel_values_merged.to(self.device)) - generated_text_merged = self.model_zoo.get('trocr_processor').batch_decode( - generated_ids_merged, skip_special_tokens=True) - - extracted_texts = extracted_texts + generated_text_merged - - else: - splited_images, _ = return_textlines_split_if_needed(img_crop, None) - #print(splited_images) - if splited_images: - cropped_lines.append(resize_image(splited_images[0], - tr_ocr_input_height_and_width, - tr_ocr_input_height_and_width)) - cropped_lines_meging_indexing.append(1) - indexer_b_s+=1 - - if indexer_b_s==self.b_s: - imgs = cropped_lines[:] - cropped_lines = [] - indexer_b_s = 0 - - pixel_values_merged = self.model_zoo.get('trocr_processor')(imgs, return_tensors="pt").pixel_values - generated_ids_merged = self.model_zoo.get('ocr').generate( - pixel_values_merged.to(self.device)) - generated_text_merged = self.model_zoo.get('trocr_processor').batch_decode( - generated_ids_merged, skip_special_tokens=True) - - extracted_texts = extracted_texts + generated_text_merged - - - cropped_lines.append(resize_image(splited_images[1], - tr_ocr_input_height_and_width, - tr_ocr_input_height_and_width)) - cropped_lines_meging_indexing.append(-1) - indexer_b_s+=1 - - if indexer_b_s==self.b_s: - imgs = cropped_lines[:] - cropped_lines = [] - indexer_b_s = 0 - - pixel_values_merged = self.model_zoo.get('trocr_processor')(imgs, return_tensors="pt").pixel_values - generated_ids_merged = self.model_zoo.get('ocr').generate( - pixel_values_merged.to(self.device)) - generated_text_merged = self.model_zoo.get('trocr_processor').batch_decode( - generated_ids_merged, skip_special_tokens=True) - - extracted_texts = extracted_texts + generated_text_merged - - else: - cropped_lines.append(img_crop) - cropped_lines_meging_indexing.append(0) - indexer_b_s+=1 - - if indexer_b_s==self.b_s: - imgs = cropped_lines[:] - cropped_lines = [] - indexer_b_s = 0 - - pixel_values_merged = self.model_zoo.get('trocr_processor')(imgs, return_tensors="pt").pixel_values - generated_ids_merged = self.model_zoo.get('ocr').generate( - pixel_values_merged.to(self.device)) - generated_text_merged = self.model_zoo.get('trocr_processor').batch_decode( - generated_ids_merged, skip_special_tokens=True) - - extracted_texts = extracted_texts + generated_text_merged - - - - indexer_text_region = indexer_text_region +1 - - if indexer_b_s!=0: - imgs = cropped_lines[:] - cropped_lines = [] - indexer_b_s = 0 - - pixel_values_merged = self.model_zoo.get('trocr_processor')(imgs, return_tensors="pt").pixel_values - generated_ids_merged = self.model_zoo.get('ocr').generate(pixel_values_merged.to(self.device)) - generated_text_merged = self.model_zoo.get('trocr_processor').batch_decode(generated_ids_merged, skip_special_tokens=True) - - extracted_texts = extracted_texts + generated_text_merged - - ####extracted_texts = [] - ####n_iterations = math.ceil(len(cropped_lines) / self.b_s) - - ####for i in range(n_iterations): - ####if i==(n_iterations-1): - ####n_start = i*self.b_s - ####imgs = cropped_lines[n_start:] - ####else: - ####n_start = i*self.b_s - ####n_end = (i+1)*self.b_s - ####imgs = cropped_lines[n_start:n_end] - ####pixel_values_merged = self.model_zoo.get('trocr_processor')(imgs, return_tensors="pt").pixel_values - ####generated_ids_merged = self.model_ocr.generate( - #### pixel_values_merged.to(self.device)) - ####generated_text_merged = self.model_zoo.get('trocr_processor').batch_decode( - #### generated_ids_merged, skip_special_tokens=True) - - ####extracted_texts = extracted_texts + generated_text_merged - + extracted_confs = [] + self.logger.debug("processing %d lines for %d regions", + len(cropped_lines), len(set(cropped_lines_region_indexer))) + for imgs in batched(cropped_lines, self.b_s): + text, conf = self.model_zoo.get('ocr').predict(imgs) + extracted_confs.extend(conf) + extracted_texts.extend(text) del cropped_lines gc.collect() extracted_texts_merged = [extracted_texts[ind] - if cropped_lines_meging_indexing[ind]==0 - else extracted_texts[ind]+" "+extracted_texts[ind+1] - if cropped_lines_meging_indexing[ind]==1 - else None - for ind in range(len(cropped_lines_meging_indexing))] - - extracted_texts_merged = [ind for ind in extracted_texts_merged if ind is not None] - #print(extracted_texts_merged, len(extracted_texts_merged)) + if cropped_lines_meging_indexing[ind] == 0 + else extracted_texts[ind] + " " + extracted_texts[ind + 1] + for ind in range(len(cropped_lines_meging_indexing)) + if cropped_lines_meging_indexing[ind] >= 0] + extracted_confs_merged = [extracted_confs[ind] + if cropped_lines_meging_indexing[ind] == 0 + else 0.5 * (extracted_confs[ind] + extracted_confs[ind + 1]) + for ind in range(len(cropped_lines_meging_indexing)) + if cropped_lines_meging_indexing[ind] >= 0] return EynollahOcrResult( extracted_texts_merged=extracted_texts_merged, - extracted_conf_value_merged=None, + extracted_confs_merged=extracted_confs_merged, cropped_lines_region_indexer=cropped_lines_region_indexer, total_bb_coordinates=total_bb_coordinates, ) @@ -273,367 +185,164 @@ class Eynollah_ocr: img_bin: Optional[MatLike], page_tree: ET.ElementTree, page_ns, - image_width, - image_height, ) -> EynollahOcrResult: - - total_bb_coordinates = [] + input_shape, _ = self.model_zoo.get('ocr').input_shape + _, image_height, image_width, _ = input_shape - cropped_lines = [] - img_crop_bin = None - imgs_bin = None - imgs_bin_ver_flipped = None + total_bb_coordinates = [] + cropped_lines_rgb = [] cropped_lines_bin = [] cropped_lines_ver_index = [] cropped_lines_region_indexer = [] cropped_lines_meging_indexing = [] - - indexer_text_region = 0 - for nn in page_tree.getroot().iter(f'{{{page_ns}}}TextRegion'): - try: - type_textregion = nn.attrib['type'] - except: - type_textregion = 'paragraph' - for child_textregion in nn: - if child_textregion.tag.endswith("TextLine"): - for child_textlines in child_textregion: - if child_textlines.tag.endswith("Coords"): - cropped_lines_region_indexer.append(indexer_text_region) - p_h=child_textlines.attrib['points'].split(' ') - textline_coords = np.array( [ [int(x.split(',')[0]), - int(x.split(',')[1]) ] - for x in p_h] ) - - x,y,w,h = cv2.boundingRect(textline_coords) - - angle_radians = math.atan2(h, w) - # Convert to degrees - angle_degrees = math.degrees(angle_radians) - if type_textregion=='drop-capital': - angle_degrees = 0 - - total_bb_coordinates.append([x,y,w,h]) - - w_scaled = w * image_height/float(h) - - img_poly_on_img = np.copy(img) - if img_bin: - img_poly_on_img_bin = np.copy(img_bin) - img_crop_bin = img_poly_on_img_bin[y:y+h, x:x+w, :] - - mask_poly = np.zeros(img.shape) - mask_poly = cv2.fillPoly(mask_poly, pts=[textline_coords], color=(1, 1, 1)) - - - mask_poly = mask_poly[y:y+h, x:x+w, :] - img_crop = img_poly_on_img[y:y+h, x:x+w, :] - - # print(file_name, angle_degrees, w*h, - # mask_poly[:,:,0].sum(), - # mask_poly[:,:,0].sum() /float(w*h) , - # 'didi') - - if angle_degrees > 3: - better_des_slope = get_orientation_moments(textline_coords) - - img_crop = rotate_image_with_padding(img_crop, better_des_slope) - if img_bin: - img_crop_bin = rotate_image_with_padding(img_crop_bin, better_des_slope) - - mask_poly = rotate_image_with_padding(mask_poly, better_des_slope) - mask_poly = mask_poly.astype('uint8') - - #new bounding box - x_n, y_n, w_n, h_n = get_contours_and_bounding_boxes(mask_poly[:,:,0]) - - mask_poly = mask_poly[y_n:y_n+h_n, x_n:x_n+w_n, :] - img_crop = img_crop[y_n:y_n+h_n, x_n:x_n+w_n, :] - - if not self.do_not_mask_with_textline_contour: - img_crop[mask_poly==0] = 255 - if img_bin: - img_crop_bin = img_crop_bin[y_n:y_n+h_n, x_n:x_n+w_n, :] - if not self.do_not_mask_with_textline_contour: - img_crop_bin[mask_poly==0] = 255 - - if mask_poly[:,:,0].sum() /float(w_n*h_n) < 0.50 and w_scaled > 90: - if img_bin: - img_crop, img_crop_bin = \ - break_curved_line_into_small_pieces_and_then_merge( - img_crop, mask_poly, img_crop_bin) - else: - img_crop, _ = \ - break_curved_line_into_small_pieces_and_then_merge( - img_crop, mask_poly) - else: - better_des_slope = 0 - if not self.do_not_mask_with_textline_contour: - img_crop[mask_poly==0] = 255 - if img_bin: - if not self.do_not_mask_with_textline_contour: - img_crop_bin[mask_poly==0] = 255 - if type_textregion=='drop-capital': - pass - else: - if mask_poly[:,:,0].sum() /float(w*h) < 0.50 and w_scaled > 90: - if img_bin: - img_crop, img_crop_bin = \ - break_curved_line_into_small_pieces_and_then_merge( - img_crop, mask_poly, img_crop_bin) - else: - img_crop, _ = \ - break_curved_line_into_small_pieces_and_then_merge( - img_crop, mask_poly) - - if w_scaled < 750:#1.5*image_width: - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - img_crop, image_height, image_width) - cropped_lines.append(img_fin) - if abs(better_des_slope) > 45: - cropped_lines_ver_index.append(1) - else: - cropped_lines_ver_index.append(0) - - cropped_lines_meging_indexing.append(0) - if img_bin: - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - img_crop_bin, image_height, image_width) - cropped_lines_bin.append(img_fin) - else: - splited_images, splited_images_bin = return_textlines_split_if_needed( - img_crop, img_crop_bin if img_bin else None) - if splited_images: - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - splited_images[0], image_height, image_width) - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(1) - - if abs(better_des_slope) > 45: - cropped_lines_ver_index.append(1) - else: - cropped_lines_ver_index.append(0) - - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - splited_images[1], image_height, image_width) - - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(-1) - - if abs(better_des_slope) > 45: - cropped_lines_ver_index.append(1) - else: - cropped_lines_ver_index.append(0) - - if img_bin: - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - splited_images_bin[0], image_height, image_width) - cropped_lines_bin.append(img_fin) - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - splited_images_bin[1], image_height, image_width) - cropped_lines_bin.append(img_fin) - - else: - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - img_crop, image_height, image_width) - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(0) - - if abs(better_des_slope) > 45: - cropped_lines_ver_index.append(1) - else: - cropped_lines_ver_index.append(0) - - if img_bin: - img_fin = preprocess_and_resize_image_for_ocrcnn_model( - img_crop_bin, image_height, image_width) - cropped_lines_bin.append(img_fin) - + img_rgb = img # cosmetic + if img_bin is None: + # run ad-hoc binarization + self.logger.info("running binarization for ensemble input") + img_bin = self.do_prediction(True, img, self.model_zoo.get("binarization"), + n_batch_inference=5) + img_bin = np.repeat(img_bin[:, :, np.newaxis], 3, axis=2) + img_bin = 255 * (img_bin == 0).astype(np.uint8) + + for n_region, region in enumerate(page_tree.getroot().iter('{%s}TextRegion' % page_ns)): + type_textregion = region.attrib.get('type', 'paragraph') + for n_line, line in enumerate(region.iter('{%s}TextLine' % page_ns)): + cropped_lines_region_indexer.append(n_region) + + coords = line.find('{%s}Coords' % page_ns) + if coords is None: + self.logger.warning("region '%s' line '%s' has no Coords", region.attrib['id'], line.attrib['id']) + continue + poly = np.array(polygon_from_points(coords.attrib['points'])).astype(int) + cont = poly[:, np.newaxis] + xywh = xywh_from_polygon(poly) + x, y, w, h = xywh['x'], xywh['y'], xywh['w'], xywh['h'] + + angle_radians = math.atan2(h, w) + angle_degrees = math.degrees(angle_radians) + if type_textregion=='drop-capital': + angle_degrees = 0 + + total_bb_coordinates.append([x, y, w, h]) + + w_scaled = w * image_height / float(h) + + img_crop_rgb = img_rgb[y: y + h, x: x + w] + img_crop_bin = img_bin[y: y + h, x: x + w] + + mask_poly = np.zeros(img_crop_rgb.shape[:2], dtype=np.uint8) + mask_poly = cv2.fillPoly(mask_poly, pts=[cont - [x, y]], color=1) + + if angle_degrees > 3: + better_des_slope = get_orientation_moments(cont) + img_crop_rgb = rotate_image_with_padding(img_crop_rgb, better_des_slope) + img_crop_bin = rotate_image_with_padding(img_crop_bin, better_des_slope) + mask_poly = rotate_image_with_padding(mask_poly, better_des_slope) + # get new bounding box + x_n, y_n, w_n, h_n = get_contours_and_bounding_boxes(mask_poly) + img_crop_rgb = img_crop_rgb[y_n: y_n + h_n, x_n: x_n + w_n] + img_crop_bin = img_crop_bin[y_n: y_n + h_n, x_n: x_n + w_n] + mask_poly = mask_poly[y_n: y_n + h_n, x_n: x_n + w_n] + else: + better_des_slope = 0 + + if not self.do_not_mask_with_textline_contour: + img_crop_rgb[mask_poly == 0] = 255 # FIXME: or median color? + img_crop_bin[mask_poly == 0] = 255 + + if (type_textregion !='drop-capital' and + mask_poly.sum() < 0.50 * mask_poly.size and + w_scaled > 90): + + img_crop_rgb, img_crop_bin = \ + break_curved_line_into_small_pieces_and_then_merge( + img_crop_rgb, img_crop_bin, mask_poly) + + if w_scaled < 750:#1.5*image_width: + img_crop_split_rgb = img_crop_split_bin = None + else: + img_crop_split_rgb, img_crop_split_bin = return_textlines_split_if_needed( + img_crop_rgb, img_crop_bin) + if img_crop_split_rgb: + cropped_lines_rgb.extend(img_crop_split_rgb) + cropped_lines_bin.extend(img_crop_split_bin) + if abs(better_des_slope) > 45: + cropped_lines_ver_index.append(1) + cropped_lines_ver_index.append(1) + else: + cropped_lines_ver_index.append(0) + cropped_lines_ver_index.append(0) + cropped_lines_meging_indexing.append(1) + cropped_lines_meging_indexing.append(-1) + else: + cropped_lines_rgb.append(img_crop_rgb) + cropped_lines_bin.append(img_crop_bin) + if abs(better_des_slope) > 45: + cropped_lines_ver_index.append(1) + else: + cropped_lines_ver_index.append(0) + cropped_lines_meging_indexing.append(0) + + cropped_lines_rgb = [preprocess_and_resize_image_for_ocrcnn_model(img, image_height, image_width) + for img in cropped_lines_rgb] + cropped_lines_bin = [preprocess_and_resize_image_for_ocrcnn_model(img, image_height, image_width) + for img in cropped_lines_bin] - indexer_text_region = indexer_text_region +1 - extracted_texts = [] - extracted_conf_value = [] + extracted_confs = [] + self.logger.debug("processing %d lines for %d regions", + len(cropped_lines_rgb), len(set(cropped_lines_region_indexer))) + cropped_lines = zip(cropped_lines_rgb, cropped_lines_bin, cropped_lines_ver_index) + for batch in batched(cropped_lines, self.b_s): + imgs_rgb, imgs_bin, ver_index = zip(*batch) + ver_index = np.array(ver_index) + imgs_rgb = np.stack(imgs_rgb) + imgs_bin = np.stack(imgs_bin) + if ver_index.any(): + imgs_rgb = np.append(imgs_rgb, imgs_rgb[ver_index > 0, ::-1, ::-1], axis=0) + imgs_bin = np.append(imgs_bin, imgs_bin[ver_index > 0, ::-1, ::-1], axis=0) - n_iterations = math.ceil(len(cropped_lines) / self.b_s) - - # FIXME: copy pasta - for i in range(n_iterations): - if i==(n_iterations-1): - n_start = i*self.b_s - imgs = cropped_lines[n_start:] - imgs = np.array(imgs) - imgs = imgs.reshape(imgs.shape[0], image_height, image_width, 3) - - ver_imgs = np.array( cropped_lines_ver_index[n_start:] ) - indices_ver = np.where(ver_imgs == 1)[0] - - #print(indices_ver, 'indices_ver') - if len(indices_ver)>0: - imgs_ver_flipped = imgs[indices_ver, : ,: ,:] - imgs_ver_flipped = imgs_ver_flipped[:,::-1,::-1,:] - #print(imgs_ver_flipped, 'imgs_ver_flipped') - - else: - imgs_ver_flipped = None - - if img_bin: - imgs_bin = cropped_lines_bin[n_start:] - imgs_bin = np.array(imgs_bin) - imgs_bin = imgs_bin.reshape(imgs_bin.shape[0], image_height, image_width, 3) - - if len(indices_ver)>0: - imgs_bin_ver_flipped = imgs_bin[indices_ver, : ,: ,:] - imgs_bin_ver_flipped = imgs_bin_ver_flipped[:,::-1,::-1,:] - #print(imgs_ver_flipped, 'imgs_ver_flipped') - - else: - imgs_bin_ver_flipped = None - else: - n_start = i*self.b_s - n_end = (i+1)*self.b_s - imgs = cropped_lines[n_start:n_end] - imgs = np.array(imgs).reshape(self.b_s, image_height, image_width, 3) - - ver_imgs = np.array( cropped_lines_ver_index[n_start:n_end] ) - indices_ver = np.where(ver_imgs == 1)[0] - #print(indices_ver, 'indices_ver') - - if len(indices_ver)>0: - imgs_ver_flipped = imgs[indices_ver, : ,: ,:] - imgs_ver_flipped = imgs_ver_flipped[:,::-1,::-1,:] - #print(imgs_ver_flipped, 'imgs_ver_flipped') - else: - imgs_ver_flipped = None - - - if img_bin: - imgs_bin = cropped_lines_bin[n_start:n_end] - imgs_bin = np.array(imgs_bin).reshape(self.b_s, image_height, image_width, 3) - - - if len(indices_ver)>0: - imgs_bin_ver_flipped = imgs_bin[indices_ver, : ,: ,:] - imgs_bin_ver_flipped = imgs_bin_ver_flipped[:,::-1,::-1,:] - #print(imgs_ver_flipped, 'imgs_ver_flipped') - else: - imgs_bin_ver_flipped = None - - - self.logger.debug("processing next %d lines", len(imgs)) - preds = self.model_zoo.get('ocr').predict(imgs, verbose=0) + # inference model now yields (char-bytes, line-prob) instead of vocidx-softmax + # (so ctc_decode and inverse StringLookup are included) + # also, the model now expects a secondary binary input image + preds, probs = self.model_zoo.get('ocr').predict((imgs_rgb, imgs_bin), verbose=0) - if len(indices_ver)>0: - preds_flipped = self.model_zoo.get('ocr').predict(imgs_ver_flipped, verbose=0) - preds_max_fliped = np.max(preds_flipped, axis=2 ) - preds_max_args_flipped = np.argmax(preds_flipped, axis=2 ) - pred_max_not_unk_mask_bool_flipped = preds_max_args_flipped[:,:]!=self.end_character - masked_means_flipped = \ - np.sum(preds_max_fliped * pred_max_not_unk_mask_bool_flipped, axis=1) / \ - np.sum(pred_max_not_unk_mask_bool_flipped, axis=1) - masked_means_flipped[np.isnan(masked_means_flipped)] = 0 - - preds_max = np.max(preds, axis=2 ) - preds_max_args = np.argmax(preds, axis=2 ) - pred_max_not_unk_mask_bool = preds_max_args[:,:]!=self.end_character - - masked_means = \ - np.sum(preds_max * pred_max_not_unk_mask_bool, axis=1) / \ - np.sum(pred_max_not_unk_mask_bool, axis=1) - masked_means[np.isnan(masked_means)] = 0 - - masked_means_ver = masked_means[indices_ver] - #print(masked_means_ver, 'pred_max_not_unk') - - indices_where_flipped_conf_value_is_higher = \ - np.where(masked_means_flipped > masked_means_ver)[0] - - #print(indices_where_flipped_conf_value_is_higher, 'indices_where_flipped_conf_value_is_higher') - if len(indices_where_flipped_conf_value_is_higher)>0: - indices_to_be_replaced = indices_ver[indices_where_flipped_conf_value_is_higher] - preds[indices_to_be_replaced,:,:] = \ - preds_flipped[indices_where_flipped_conf_value_is_higher, :, :] + if ver_index.any(): + preds, preds_ver = np.split(preds, [-np.count_nonzero(ver_index)], axis=0) + probs, probs_ver = np.split(probs, [-np.count_nonzero(ver_index)], axis=0) + flipped_ver_is_better = np.flatnonzero(probs_ver > probs[ver_index > 0]) + if len(flipped_ver_is_better): + self.logger.info("%d skewed lines perform better when flipped", len(flipped_ver_is_better)) + preds[ver_index > 0][flipped_ver_is_better] = preds_ver[flipped_ver_is_better] + probs[ver_index > 0][flipped_ver_is_better] = probs_ver[flipped_ver_is_better] - if img_bin: - preds_bin = self.model_zoo.get('ocr').predict(imgs_bin, verbose=0) - - if len(indices_ver)>0: - preds_flipped = self.model_zoo.get('ocr').predict(imgs_bin_ver_flipped, verbose=0) - preds_max_fliped = np.max(preds_flipped, axis=2 ) - preds_max_args_flipped = np.argmax(preds_flipped, axis=2 ) - pred_max_not_unk_mask_bool_flipped = preds_max_args_flipped[:,:]!=self.end_character - masked_means_flipped = \ - np.sum(preds_max_fliped * pred_max_not_unk_mask_bool_flipped, axis=1) / \ - np.sum(pred_max_not_unk_mask_bool_flipped, axis=1) - masked_means_flipped[np.isnan(masked_means_flipped)] = 0 - - preds_max = np.max(preds, axis=2 ) - preds_max_args = np.argmax(preds, axis=2 ) - pred_max_not_unk_mask_bool = preds_max_args[:,:]!=self.end_character - - masked_means = \ - np.sum(preds_max * pred_max_not_unk_mask_bool, axis=1) / \ - np.sum(pred_max_not_unk_mask_bool, axis=1) - masked_means[np.isnan(masked_means)] = 0 - - masked_means_ver = masked_means[indices_ver] - #print(masked_means_ver, 'pred_max_not_unk') - - indices_where_flipped_conf_value_is_higher = \ - np.where(masked_means_flipped > masked_means_ver)[0] - - #print(indices_where_flipped_conf_value_is_higher, 'indices_where_flipped_conf_value_is_higher') - if len(indices_where_flipped_conf_value_is_higher)>0: - indices_to_be_replaced = indices_ver[indices_where_flipped_conf_value_is_higher] - preds_bin[indices_to_be_replaced,:,:] = \ - preds_flipped[indices_where_flipped_conf_value_is_higher, :, :] - - preds = (preds + preds_bin) / 2. - - pred_texts = decode_batch_predictions(preds, self.model_zoo.get('num_to_char')) - - preds_max = np.max(preds, axis=2 ) - preds_max_args = np.argmax(preds, axis=2 ) - pred_max_not_unk_mask_bool = preds_max_args[:,:]!=self.end_character - masked_means = \ - np.sum(preds_max * pred_max_not_unk_mask_bool, axis=1) / \ - np.sum(pred_max_not_unk_mask_bool, axis=1) - - for ib in range(imgs.shape[0]): - pred_texts_ib = pred_texts[ib].replace("[UNK]", "") - if masked_means[ib] >= self.min_conf_value_of_textline_text: - extracted_texts.append(pred_texts_ib) - extracted_conf_value.append(masked_means[ib]) - else: - extracted_texts.append("") - extracted_conf_value.append(0) - del cropped_lines + def nooov(x): + if x == b'[UNK]': + return b'' + return x + for pred, prob in zip(preds, probs): + text = b''.join(map(nooov, pred.tolist())).decode('utf-8') + extracted_texts.append(text) + extracted_confs.append(prob) + del cropped_lines_rgb del cropped_lines_bin gc.collect() extracted_texts_merged = [extracted_texts[ind] - if cropped_lines_meging_indexing[ind]==0 - else extracted_texts[ind]+" "+extracted_texts[ind+1] - if cropped_lines_meging_indexing[ind]==1 - else None - for ind in range(len(cropped_lines_meging_indexing))] - - extracted_conf_value_merged = [extracted_conf_value[ind] # type: ignore - if cropped_lines_meging_indexing[ind]==0 - else (extracted_conf_value[ind]+extracted_conf_value[ind+1])/2. - if cropped_lines_meging_indexing[ind]==1 - else None - for ind in range(len(cropped_lines_meging_indexing))] - - extracted_conf_value_merged: List[float] = [extracted_conf_value_merged[ind_cfm] - for ind_cfm in range(len(extracted_texts_merged)) - if extracted_texts_merged[ind_cfm] is not None] - - extracted_texts_merged = [ind for ind in extracted_texts_merged if ind is not None] + if cropped_lines_meging_indexing[ind] == 0 + else extracted_texts[ind] + " " + extracted_texts[ind + 1] + for ind in range(len(cropped_lines_meging_indexing)) + if cropped_lines_meging_indexing[ind] >= 0] + extracted_confs_merged = [extracted_confs[ind] + if cropped_lines_meging_indexing[ind] == 0 + else 0.5 * (extracted_confs[ind] + extracted_confs[ind + 1]) + for ind in range(len(cropped_lines_meging_indexing)) + if cropped_lines_meging_indexing[ind] >= 0] return EynollahOcrResult( extracted_texts_merged=extracted_texts_merged, - extracted_conf_value_merged=extracted_conf_value_merged, + extracted_confs_merged=extracted_confs_merged, cropped_lines_region_indexer=cropped_lines_region_indexer, total_bb_coordinates=total_bb_coordinates, ) @@ -651,13 +360,12 @@ class Eynollah_ocr: cropped_lines_region_indexer = result.cropped_lines_region_indexer total_bb_coordinates = result.total_bb_coordinates extracted_texts_merged = result.extracted_texts_merged - extracted_conf_value_merged = result.extracted_conf_value_merged + extracted_confs_merged = result.extracted_confs_merged - unique_cropped_lines_region_indexer = np.unique(cropped_lines_region_indexer) if out_image_with_text: image_text = Image.new("RGB", (img.shape[1], img.shape[0]), "white") draw = ImageDraw.Draw(image_text) - font = get_font() + font = get_font(font_size=40) for indexer_text, bb_ind in enumerate(total_bb_coordinates): x_bb = bb_ind[0] @@ -681,91 +389,66 @@ class Eynollah_ocr: draw.text((text_x, text_y), extracted_texts_merged[indexer_text], fill="black", font=font) image_text.save(out_image_with_text) - text_by_textregion = [] - for ind in unique_cropped_lines_region_indexer: - ind = np.array(cropped_lines_region_indexer)==ind - extracted_texts_merged_un = np.array(extracted_texts_merged)[ind] - if len(extracted_texts_merged_un)>1: - text_by_textregion_ind = "" - next_glue = "" - for indt in range(len(extracted_texts_merged_un)): - if (extracted_texts_merged_un[indt].endswith('⸗') or - extracted_texts_merged_un[indt].endswith('-') or - extracted_texts_merged_un[indt].endswith('¬')): - text_by_textregion_ind += next_glue + extracted_texts_merged_un[indt][:-1] - next_glue = "" - else: - text_by_textregion_ind += next_glue + extracted_texts_merged_un[indt] - next_glue = " " - text_by_textregion.append(text_by_textregion_ind) - else: - text_by_textregion.append(" ".join(extracted_texts_merged_un)) + cropped_lines_region_indexer = np.array(cropped_lines_region_indexer) + for n_region, region in enumerate(page_tree.getroot().iter('{%s}TextRegion' % page_ns)): + lines_indexer = np.flatnonzero(cropped_lines_region_indexer == n_region) + if not len(lines_indexer): + continue - indexer = 0 - indexer_textregion = 0 - for nn in page_tree.getroot().iter(f'{{{page_ns}}}TextRegion'): - - is_textregion_text = False - for childtest in nn: - if childtest.tag.endswith("TextEquiv"): - is_textregion_text = True - - if not is_textregion_text: - text_subelement_textregion = ET.SubElement(nn, 'TextEquiv') - unicode_textregion = ET.SubElement(text_subelement_textregion, 'Unicode') - - - has_textline = False - for child_textregion in nn: - if child_textregion.tag.endswith("TextLine"): - - is_textline_text = False - for childtest2 in child_textregion: - if childtest2.tag.endswith("TextEquiv"): - is_textline_text = True - - - if not is_textline_text: - text_subelement = ET.SubElement(child_textregion, 'TextEquiv') - if extracted_conf_value_merged: - text_subelement.set('conf', f"{extracted_conf_value_merged[indexer]:.2f}") - unicode_textline = ET.SubElement(text_subelement, 'Unicode') - unicode_textline.text = extracted_texts_merged[indexer] - else: - for childtest3 in child_textregion: - if childtest3.tag.endswith("TextEquiv"): - for child_uc in childtest3: - if child_uc.tag.endswith("Unicode"): - if extracted_conf_value_merged: - childtest3.set('conf', f"{extracted_conf_value_merged[indexer]:.2f}") - child_uc.text = extracted_texts_merged[indexer] - - indexer = indexer + 1 - has_textline = True - if has_textline: - if is_textregion_text: - for child4 in nn: - if child4.tag.endswith("TextEquiv"): - for childtr_uc in child4: - if childtr_uc.tag.endswith("Unicode"): - childtr_uc.text = text_by_textregion[indexer_textregion] + text_region = "" + next_glue = "" + for line_idx in lines_indexer: + if extracted_confs_merged[line_idx] < self.min_conf_value_of_textline_text: + continue + text_line = extracted_texts_merged[line_idx] + if (text_line.endswith(('⸗', '-', '¬')) and + # last line of a region can still be wrapped + # around columns or pages + line_idx < len(lines_indexer) - 1): + text_region += next_glue + text_line[:-1] + next_glue = "" else: - unicode_textregion.text = text_by_textregion[indexer_textregion] - indexer_textregion = indexer_textregion + 1 + text_region += next_glue + text_line + next_glue = " " + + region_textequiv = region.find('{%s}TextEquiv' % page_ns) + if region_textequiv is None: + region_textequiv = ET.SubElement(region, 'TextEquiv') + region_teunicode = region_textequiv.find('{%s}Unicode' % page_ns) + if region_teunicode is None: + region_teunicode = ET.SubElement(region_textequiv, 'Unicode') + region_teunicode.text = text_region + + for n_line, line in enumerate(region.iter('{%s}TextLine' % page_ns)): + line_textequiv = line.find('{%s}TextEquiv' % page_ns) + if line_textequiv is None: + line_textequiv = ET.SubElement(line, 'TextEquiv') + line_teunicode = line_textequiv.find('{%s}Unicode' % page_ns) + if line_teunicode is None: + line_teunicode = ET.SubElement(line_textequiv, 'Unicode') + + line_idx = lines_indexer[n_line] + if extracted_confs_merged[line_idx] < self.min_conf_value_of_textline_text: + line.remove(line_textequiv) + else: + line_textequiv.set('conf', str(round(extracted_confs_merged[line_idx], 2))) + line_teunicode.text = extracted_texts_merged[line_idx] ET.register_namespace("",page_ns) + self.logger.info("output filename: '%s'", out_file_ocr) page_tree.write(out_file_ocr, xml_declaration=True, method='xml', encoding="utf-8", default_namespace=None) - def run( - self, - *, - overwrite: bool = False, - dir_in: Optional[str] = None, - dir_in_bin: Optional[str] = None, - image_filename: Optional[str] = None, - dir_xmls: str, - dir_out_image_text: Optional[str] = None, - dir_out: str, + def run(self, + *, + overwrite: bool = False, + dir_in: str = "", + dir_in_bin: str = "", + image_filename: str = "", + dir_xmls: str, + dir_out_image_text: str = "", + dir_out: str, + num_jobs: int = 0, + halt_fail: float = 0, ): """ Run OCR. @@ -775,63 +458,129 @@ class Eynollah_ocr: dir_in_bin (str): Prediction with RGB and binarized images for selected pages, should not be the default """ if dir_in: + t0_tot = time.time() ls_imgs = [os.path.join(dir_in, image_filename) - for image_filename in filter(is_image_filename, + for image_filename in filter(is_image_filename, os.listdir(dir_in))] + if dir_in_bin and dir_in_bin == dir_in: + # try filtering PNGs from rest + def pathstem(filename): + return os.path.splitext(filename)[0] + def notpng(filenames): + for filename in filenames: + if not filename.lower().endswith(".png"): + return filename + return filenames[0] + ls_imgs = [notpng(files) + for _, files in groupby(sorted(ls_imgs), + key=pathstem)] + with ProcessPoolExecutor(max_workers=num_jobs or None, + mp_context=mp.get_context('fork'), + initializer=_set_instance, + initargs=(self,) + ) as exe: + jobs = {} + mngr = mp.get_context('fork').Manager() + n_success = n_fail = 0 + for img_filename in ls_imgs: + logq = mngr.Queue() + jobs[exe.submit(_run_single, img_filename, + dir_out=dir_out, + dir_xmls=dir_xmls, + dir_in_bin=dir_in_bin, + dir_out_image_text=dir_out_image_text, + overwrite=overwrite, + logq=logq)] = img_filename, logq + for job in as_completed(list(jobs)): + img_filename, logq = jobs[job] + loglistener = logging.handlers.QueueListener( + logq, *self.logger.handlers, + *self.logger.parent.handlers, + respect_handler_level=False) + try: + loglistener.start() + job.result() + n_success += 1 + except: + self.logger.exception("Job %s failed", img_filename) + n_fail += 1 + if (halt_fail and + n_fail >= halt_fail * (len(jobs) if halt_fail < 1 else 1)): + self.logger.fatal("terminating after %d failures", n_fail) + for job in jobs: + job.cancel() + break + finally: + loglistener.stop() + self.logger.info("%d of %d jobs successful", n_success, len(jobs)) + self.logger.info("All jobs done in %.1fs", time.time() - t0_tot) else: assert image_filename - ls_imgs = [image_filename] + self.run_single(image_filename, + dir_xmls=dir_xmls, + dir_out=dir_out, + dir_in_bin=dir_in_bin, + dir_out_image_text=dir_out_image_text, + overwrite=overwrite) - for img_filename in ls_imgs: - file_stem = Path(img_filename).stem - page_file_in = os.path.join(dir_xmls, file_stem+'.xml') - out_file_ocr = os.path.join(dir_out, file_stem+'.xml') - - if os.path.exists(out_file_ocr): - if overwrite: - self.logger.warning("will overwrite existing output file '%s'", out_file_ocr) - else: - self.logger.warning("will skip input for existing output file '%s'", out_file_ocr) - return - - img = cv2.imread(img_filename) + def run_single(self, + img_filename: str, + dir_xmls: str, + dir_out: str = "", + dir_in_bin: str = "", + dir_out_image_text: str = "", + overwrite: bool = False, + ): + file_stem = Path(img_filename).stem + page_file_in = os.path.join(dir_xmls, file_stem + '.xml') + out_file_ocr = os.path.join(dir_out, file_stem + '.xml') - page_tree = ET.parse(page_file_in, parser = ET.XMLParser(encoding="utf-8")) - page_ns = etree_namespace_for_element_tag(page_tree.getroot().tag) - - out_image_with_text = None - if dir_out_image_text: - out_image_with_text = os.path.join(dir_out_image_text, file_stem + '.png') - - img_bin = None - if dir_in_bin: - img_bin = cv2.imread(os.path.join(dir_in_bin, file_stem+'.png')) - - - if self.tr_ocr: - result = self.run_trocr( - img=img, - page_tree=page_tree, - page_ns=page_ns, - - tr_ocr_input_height_and_width = 384 - ) + if os.path.exists(out_file_ocr): + if overwrite: + self.logger.warning("will overwrite existing output file '%s'", out_file_ocr) else: - result = self.run_cnn( - img=img, - page_tree=page_tree, - page_ns=page_ns, + self.logger.warning("will skip input for existing output file '%s'", out_file_ocr) + return + if not os.path.exists(page_file_in): + self.logger.error("will skip missing input file '%s'", page_file_in) + return - img_bin=img_bin, - image_width=512, - image_height=32, - ) + t0 = time.time() - self.write_ocr( - result=result, + img = cv2.imread(img_filename) + self.logger.info(img_filename) + page_tree = ET.parse(page_file_in, parser = ET.XMLParser(encoding="utf-8")) + page_ns = etree_namespace_for_element_tag(page_tree.getroot().tag) + + out_image_with_text = None + if dir_out_image_text: + out_image_with_text = os.path.join(dir_out_image_text, file_stem + '.png') + + img_bin = None + if dir_in_bin: + img_bin = cv2.imread(os.path.join(dir_in_bin, file_stem+'.png')) + + + if self.tr_ocr: + result = self.run_trocr( img=img, page_tree=page_tree, page_ns=page_ns, - out_file_ocr=out_file_ocr, - out_image_with_text=out_image_with_text, ) + else: + result = self.run_cnn( + img=img, + page_tree=page_tree, + page_ns=page_ns, + img_bin=img_bin, + ) + + self.write_ocr( + result=result, + img=img, + page_tree=page_tree, + page_ns=page_ns, + out_file_ocr=out_file_ocr, + out_image_with_text=out_image_with_text, + ) + self.logger.info("Job done in %.1fs", time.time() - t0) diff --git a/src/eynollah/mb_ro_on_layout.py b/src/eynollah/mb_ro_on_layout.py index b0b5910..5bab608 100644 --- a/src/eynollah/mb_ro_on_layout.py +++ b/src/eynollah/mb_ro_on_layout.py @@ -17,9 +17,7 @@ import cv2 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 .eynollah import Eynollah from .model_zoo import EynollahModelZoo from .utils.resize import resize_image from .utils.contour import ( @@ -33,23 +31,26 @@ DPI_THRESHOLD = 298 KERNEL = np.ones((5, 5), np.uint8) -class machine_based_reading_order_on_layout: +class Reorder(Eynollah): def __init__( - self, - *, - model_zoo: EynollahModelZoo, - logger : Optional[logging.Logger] = None, + self, + *, + model_zoo: EynollahModelZoo, + logger : Optional[logging.Logger] = None, + device: str = '', ): self.logger = logger or logging.getLogger('eynollah.mbreorder') self.model_zoo = model_zoo - 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.model_zoo.load_models('reading_order') + self.setup_models(device=device) + + def setup_models(self, device=''): + loadable = ['reading_order'] + self.model_zoo.load_models(*loadable, device=device) + for model in loadable: + self.logger.debug("model %s has input shape %s", model, + self.model_zoo.get(model).input_shape) + def read_xml(self, xml_file): tree1 = ET.parse(xml_file, parser = ET.XMLParser(encoding='utf-8')) @@ -675,7 +676,7 @@ class machine_based_reading_order_on_layout: tot_counter += 1 batch.append(j) if tot_counter % inference_bs == 0 or tot_counter == len(ij_list): - y_pr = self.model_zoo.get('reading_order').predict(input_1 , verbose='0') + y_pr = self.model_zoo.get('reading_order').predict(input_1, verbose=0) for jb, j in enumerate(batch): if y_pr[jb][0]>=0.5: post_list.append(j) diff --git a/src/eynollah/model_zoo/default_specs.py b/src/eynollah/model_zoo/default_specs.py index dc725e4..7fe91f8 100644 --- a/src/eynollah/model_zoo/default_specs.py +++ b/src/eynollah/model_zoo/default_specs.py @@ -1,8 +1,8 @@ from .specs import EynollahModelSpec, EynollahModelSpecSet # NOTE: This needs to change whenever models/versions change -ZENODO = "https://zenodo.org/records/17727267" -MODELS_VERSION = "v0_8_0" +ZENODO = "https://zenodo.org/records/21381102" +MODELS_VERSION = "v0_9_1" def dist_url(dist_name: str="layout") -> str: return f'{ZENODO}/models_{dist_name}_{MODELS_VERSION}.zip' @@ -208,22 +208,6 @@ DEFAULT_MODEL_SPECS = EynollahModelSpecSet([ type='Keras', ), - EynollahModelSpec( - category="num_to_char", - variant='', - filename="characters_org.txt", - dist_url=dist_url("ocr"), - type='decoder', - ), - - EynollahModelSpec( - category="characters", - variant='', - filename="characters_org.txt", - dist_url=dist_url("ocr"), - type='List[str]', - ), - EynollahModelSpec( category="ocr", variant='tr', @@ -233,20 +217,4 @@ DEFAULT_MODEL_SPECS = EynollahModelSpecSet([ type='Keras', ), - EynollahModelSpec( - category="trocr_processor", - variant='', - filename="models_eynollah/model_eynollah_ocr_trocr_20250919", - dist_url=dist_url("ocr"), - type='TrOCRProcessor', - ), - - EynollahModelSpec( - category="trocr_processor", - variant='htr', - filename="models_eynollah/microsoft/trocr-base-handwritten", - dist_url=dist_url("extra"), - type='TrOCRProcessor', - ), - ]) diff --git a/src/eynollah/model_zoo/model_zoo.py b/src/eynollah/model_zoo/model_zoo.py index 9611388..5291bf8 100644 --- a/src/eynollah/model_zoo/model_zoo.py +++ b/src/eynollah/model_zoo/model_zoo.py @@ -14,6 +14,19 @@ from .default_specs import DEFAULT_MODEL_SPECS from .types import AnyModel, T +MODEL_VRAM_LIMITS = { + "binarization": 868, # due to bs 5 + "enhancement": 980, # due to bs 3 + "col_classifier": 210, + "page": 618, + "textline": 1880, # 954 for bs 1 + "region_1_2": 1580, + "region_fl_np": 1756, + "table": 1818, + "reading_order": 632, + "ocr": 2600, # 850 for bs 8 +} + class EynollahModelZoo: """ Wrapper class that handles storage and loading of models for all eynollah runners. @@ -35,7 +48,7 @@ class EynollahModelZoo: self._overrides = [] if model_overrides: self.override_models(*model_overrides) - self._loaded: Dict[str, Predictor] = {} + self._loaded: Dict[str, Union[Predictor, AnyModel]] = {} @property def model_overrides(self): @@ -70,6 +83,13 @@ class EynollahModelZoo: model_path = Path(self.model_basedir).joinpath(spec.filename) else: model_path = Path(spec.filename) + if model_path.suffix == '.h5' and Path(model_path.stem).exists(): + # prefer SavedModel over HDF5 format if it exists + model_path = Path(model_path.stem) + if model_path.with_suffix('.onnx').exists(): + # prefer ONNX over SavedModel format if it exists + model_path = model_path.with_suffix('.onnx') + return model_path def load_models( @@ -82,63 +102,90 @@ class EynollahModelZoo: """ ret = {} # cannot use self._loaded here, yet – first spawn all predictors for load_args in all_load_args: + load_kwargs = dict(device=device) if isinstance(load_args, str): - model_category = load_args - load_args = [model_category] + model_category, model_variant = load_args, "" + elif len(load_args) > 2: + # for calls to self.model_path + self.override_models(load_args) + # for calls to Predictor.load_model + model_category, model_variant, model_path = load_args + load_kwargs["model_variant"] = model_variant + load_kwargs["model_path_override"] = model_path else: - model_category = load_args[0] - load_kwargs = {} - if model_category.endswith('_resized'): - load_args[0] = model_category[:-8] - load_kwargs["resized"] = True - elif model_category.endswith('_patched'): - load_args[0] = model_category[:-8] - load_kwargs["patched"] = True - spec = self.specs.get(model_category, load_args[1] if len(load_args) > 1 else '') - if spec.type in ['Keras'] and spec.category != 'ocr': - ret[model_category] = Predictor(self.logger, self) - ret[model_category].load_model(*load_args, **load_kwargs, device=device) - else: - ret[model_category] = self.load_model(*load_args, **load_kwargs, device=device) + model_category, model_variant = load_args + load_kwargs["model_variant"] = model_variant + + # if model_category.endswith('_resized'): + # model_category = model_category[:-8] + # load_kwargs["resized"] = True + # elif model_category.endswith('_patched'): + # model_category = model_category[:-8] + # load_kwargs["patched"] = True + + model = Predictor(self.logger, self) + model.load_model(model_category, **load_kwargs) + + ret[model_category] = model self._loaded.update(ret) return self._loaded def load_model( - self, - model_category: str, - model_variant: str = '', - model_path_override: Optional[str] = None, - patched: bool = False, - resized: bool = False, + self, + model_category: str, + model_variant: str = '', + model_path_override: Optional[str] = None, + # patched: bool = False, + # resized: bool = False, device: str = '', ) -> AnyModel: """ Load any model """ - os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 + if model_path_override: + self.override_models((model_category, model_variant, model_path_override)) + model_path = self.model_path(model_category, model_variant) + + if model_category == 'ocr' and model_variant == 'tr': + model = self._load_trocr_model(model_path, device=device) + elif model_path.is_dir() and (model_path / "keras_metadata.pb").exists(): + # Keras model + model = self._load_keras_model(model_category, model_path, device=device) + elif model_path.is_dir(): + # TF-Serving model + model = self._load_serving_model(model_category, model_path, device=device) + elif model_path.suffix == '.onnx': + # ONNX model + model = self._load_onnx_model(model_category, model_path, device=device) + else: + raise ValueError("unknown model type for '%s'" % str(model_path)) + model._name = model_category + return model + + def get(self, model_category: str) -> Union[Predictor, AnyModel]: + if model_category not in self._loaded: + raise ValueError(f'Model "{model_category}" not previously loaded with "load_model(..)"') + return self._loaded[model_category] + + def _configure_tf_device(self, model_category, device=''): from ocrd_utils import tf_disable_interactive_logs tf_disable_interactive_logs() - import tensorflow as tf - from tensorflow.keras.models import load_model - from ..patch_encoder import ( - PatchEncoder, - Patches, - wrap_layout_model_patched, - wrap_layout_model_resized, - ) cuda = False try: gpus = tf.config.list_physical_devices('GPU') if device: - if ',' in device: + if ':' in device: for spec in device.split(','): cat, dev = spec.split(':') if fnmatchcase(model_category, cat): device = dev break - if device == 'CPU': + if ':' in device: + self.logger.warning("missing device specification for model type %s", model_category) + gpus = gpus[:1] + elif device == 'CPU': gpus = [] else: assert device.startswith('GPU') @@ -147,7 +194,14 @@ class EynollahModelZoo: gpus = gpus[:1] # TF will always use first allowable tf.config.set_visible_devices(gpus, 'GPU') for device in gpus: - tf.config.experimental.set_memory_growth(device, True) + # tf.config.experimental.set_memory_growth(device, True) + # dynamic growth never frees memory (to avoid fragmentation), + # so the VRAM requirements end up much larger than feasible + # (for small GPUs); so try hard (calibrated) limits instead: + tf.config.set_logical_device_configuration( + device, + [tf.config.LogicalDeviceConfiguration( + memory_limit=MODEL_VRAM_LIMITS[model_category])]) vendor_name = ( tf.config.experimental.get_device_details(device) .get('device_name', 'unknown')) @@ -155,95 +209,263 @@ class EynollahModelZoo: self.logger.info("using GPU %s (%s) for model %s", device.name, vendor_name, - model_category + ( - "_patched" if patched else - "_resized" if resized else "")) + model_category # + ( + # "_patched" if patched else + # "_resized" if resized else "") + ) except RuntimeError: self.logger.exception("cannot configure GPU devices") if not cuda: self.logger.warning("no GPU device available") - if model_path_override: - self.override_models((model_category, model_variant, model_path_override)) - model_path = self.model_path(model_category, model_variant) - if model_path.suffix == '.h5' and Path(model_path.stem).exists(): - # prefer SavedModel over HDF5 format if it exists - model_path = Path(model_path.stem) - if model_category == 'ocr': - model = self._load_ocr_model(variant=model_variant) - elif model_category == 'num_to_char': - model = self._load_num_to_char() - elif model_category == 'characters': - model = self._load_characters() - elif model_category == 'trocr_processor': - from transformers import TrOCRProcessor - model = TrOCRProcessor.from_pretrained(model_path) - else: + def _configure_torch_device(self, model_category, device=''): + import torch + + device0 = torch.device('cpu') + if not device and torch.cuda.is_available(): + device = 'GPU' # try + if device and ':' in device: + for spec in device.split(','): + cat, dev = spec.split(':') + if fnmatchcase('ocr', cat): + device = dev + break + if ':' in device: + self.logger.warning("missing device specification for model type %s", model_category) + device = 'GPU' + if device and device.startswith('GPU'): try: - # avoid wasting VRAM on non-transformer models - model = load_model(model_path, compile=False) - except Exception as e: - self.logger.error(e) - model = load_model( - model_path, compile=False, - custom_objects=dict(PatchEncoder=PatchEncoder, - Patches=Patches)) - model._name = model_category - if resized: - model = wrap_layout_model_resized(model) - model._name = model_category + '_resized' - elif patched: - model = wrap_layout_model_patched(model) - model._name = model_category + '_patched' - else: - model.jit_compile = True - model.make_predict_function() + device0 = torch.device('cuda', int(device[3:] or 0)) + name = torch.cuda.get_device_name(device0) + self.logger.info("using GPU %s (%s) for model ocr:tr", device0, name) + except: + self.logger.exception("cannot configure GPU device") + device0 = torch.device('cpu') + if device0.type != 'cuda': + self.logger.warning("no GPU device available") + return device0 + + def _load_keras_model(self, model_category, model_path, device=''): + 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 tensorflow.keras.models import load_model + from tensorflow.keras.models import Model as KerasModel + + from ..training.models import cnn_rnn_ocr_model4inference + + self._configure_tf_device(model_category, device=device) + + model = load_model(model_path, compile=False) + assert isinstance(model, KerasModel) + + # from ..patch_encoder import ( + # wrap_layout_model_patched, + # wrap_layout_model_resized, + # ) + # if resized: + # model = wrap_layout_model_resized(model) + # model._name = model_category + '_resized' + # elif patched: + # model = wrap_layout_model_patched(model) + # model._name = model_category + '_patched' + + if model_category == 'ocr': + # cnn-rnn-ocr task model may not be in inference mode, yet + model = cnn_rnn_ocr_model4inference(model, model_path) + + model.make_predict_function() + return model - def get(self, model_category: str) -> Predictor: - if model_category not in self._loaded: - raise ValueError(f'Model "{model_category}" not previously loaded with "load_model(..)"') - return self._loaded[model_category] + def _load_serving_model(self, model_category, model_path, device=''): + from ocrd_utils import tf_disable_interactive_logs + tf_disable_interactive_logs() + import tensorflow as tf - def _load_ocr_model(self, variant: str) -> AnyModel: + self._configure_tf_device(model_category, device=device) + model = tf.saved_model.load(model_path) + model.predict_on_batch = model.serve + spec = model.signatures['serving_default'] + # some models receive lots of additional/internal + # (unknown) captured inputs polluting .inputs + # TF>=2.16 has spec.function_type.flat_inputs + # this non-public API works: + # input_spec = spec.inputs[:len(spec._arg_keywords)] + # but perhaps this is most reliable: + input_spec = tf.nest.flatten(spec.structured_input_signature, True) + input_spec = [tuple(i.shape) for i in input_spec] + if len(input_spec) > 1: + model.input_shape = tuple(input_spec) + else: + model.input_shape = input_spec[0] + + return model + + def _load_onnx_model(self, model_category, model_path, device=''): + import onnxruntime as ort + import numpy as np + from ocrd_utils import config + + ort.set_default_logger_severity(3) + + providers = ort.get_available_providers() + if device: + if ':' in device: + for spec in device.split(','): + cat, dev = spec.split(':') + if fnmatchcase(model_category, cat): + device = dev + break + if ':' in device: + self.logger.warning("missing device specification for model type %s", model_category) + gpu = 0 + elif device == 'CPU': + gpu = -1 + else: + assert device.startswith('GPU') + gpu = int(device[3:] or "0") + else: + gpu = 0 # try first allowable + # make runtime-configurable + if override_providers := os.environ.get('EYNOLLAH_ONNX_EP', ''): + override_providers = override_providers.split(',') + providers = [provider for provider in providers + if provider[:-17] in override_providers] + # configure and prioritise + if 'AzureExecutionProvider' in providers: + providers.remove('AzureExecutionProvider') + if 'CUDAExecutionProvider' in providers: + providers.remove('CUDAExecutionProvider') + if gpu >= 0: + providers = [('CUDAExecutionProvider', { + 'device_id': gpu, + # 'arena_extend_strategy': 'kNextPowerOfTwo', + 'gpu_mem_limit': MODEL_VRAM_LIMITS[model_category] * 1024 * 1024, + # 'cudnn_conv_algo_search': 'EXHAUSTIVE', + #'cudnn_conv_use_max_workspace': 0, + # 'do_copy_in_default_stream': True, + # enable_cuda_graph + # cudnn_conv1d_pad_to_nc1d + # prefer_nhwc + # tunable_op_enable + # tunable_op_tuning_enable + # tunable_op_max_tuning_duration_ms + # use_ep_level_unified_stream + # enable_skip_layer_norm_strict_mode + # ... + })] + providers + if 'TensorrtExecutionProvider' in providers: + providers.remove('TensorrtExecutionProvider') + if gpu >= 0: + providers = [('TensorrtExecutionProvider', { + 'device_id': gpu, + 'trt_max_workspace_size': MODEL_VRAM_LIMITS[model_category] * 1024 * 1024, + # 'trt_fp16_enable': True, + # trt_bf16_enable + 'trt_engine_cache_enable': True, + 'trt_timing_cache_enable': True, + 'trt_engine_cache_path': config.XDG_CONFIG_HOME, + 'trt_timing_cache_path': config.XDG_CONFIG_HOME, + # ... + # trt_engine_hw_compatible + # trt_engine_cache_prefix + # trt_onnx_model_folder_path + # trt_ep_context_file_path + # trt_cuda_graph_enable + # trt_profile_opt_shapes + # trt_profile_min_shapes + # trt_profile_max_shapes + # trt_builder_optimization_level + # trt_build_heuristics_enable + # trt_sparsity_enable + # trt_weight_stripped_engine_enable + # trt_dla_core + # trt_dla_enable + # trt_min_subgraph_size + # trt_ep_context_embed_mode + })] + providers + provider0 = providers[0] + if isinstance(provider0, tuple): + provider0 = provider0[0] + self.logger.info("using %s with ONNX provider %s for model %s", + "GPU %d" % gpu if (gpu >= 0 and not + provider0.startswith("CPU")) + else "CPU", + provider0[:-17], model_category) + model = ort.InferenceSession( + model_path, + providers=providers) + model_inputs = [model_input.name + for model_input in model.get_inputs()] + model_outputs = [model_output.name + for model_output in model.get_outputs()] + def predict_onnx(inputs): + if len(model_inputs) == 1: + inputs = [inputs] + outputs = model.run(model_outputs, { + model_input: + input_data.astype( + # models expect data_type() == 'tensor(float)', but np.float16 is 'tensor(float16)' + # FIXME: do this dynamically (but how to convert .type to np.dtype?) + np.float32 if input_data.dtype in [np.float16, np.float64] else + input_data.dtype) + for model_input, input_data in zip(model_inputs, inputs) + }) + if len(model_outputs) == 1: + outputs = outputs[0] + return outputs + model.predict_on_batch = predict_onnx + input_spec = model.get_inputs() + input_spec = [i.shape for i in input_spec] + if len(input_spec) > 1: + model.input_shape = tuple(input_spec) + else: + model.input_shape = input_spec[0] + + return model + + def _load_trocr_model(self, model_path, device: str = "") -> AnyModel: """ Load OCR model """ - from tensorflow.keras.models import Model as KerasModel - from tensorflow.keras.models import load_model + from transformers import VisionEncoderDecoderModel, TrOCRProcessor + import numpy as np - ocr_model_dir = self.model_path('ocr', variant) - if variant == 'tr': - from transformers import VisionEncoderDecoderModel - ret = VisionEncoderDecoderModel.from_pretrained(ocr_model_dir) - assert isinstance(ret, VisionEncoderDecoderModel) - return ret - else: - ocr_model = load_model(ocr_model_dir, compile=False) - assert isinstance(ocr_model, KerasModel) - return KerasModel( - ocr_model.get_layer(name="image").input, # type: ignore - ocr_model.get_layer(name="dense2").output, # type: ignore - ) + device = self._configure_torch_device('ocr', device=device) + proc = TrOCRProcessor.from_pretrained(model_path) + model = VisionEncoderDecoderModel.from_pretrained(model_path) + assert isinstance(model, VisionEncoderDecoderModel) - def _load_characters(self) -> List[str]: - """ - Load encoding for OCR - """ - with open(self.model_path('num_to_char'), "r") as config_file: - return json.load(config_file) - - def _load_num_to_char(self) -> 'StringLookup': - """ - Load decoder for OCR - """ - from tensorflow.keras.layers import StringLookup - - characters = self._load_characters() - # Mapping characters to integers. - char_to_num = StringLookup(vocabulary=characters, mask_token=None) - # Mapping integers back to original characters. - return StringLookup(vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True) + model.to(device) + def predict_torch(inputs): + output = model.generate( + proc(inputs, return_tensors="pt").pixel_values.to(device), + # beam search instead of greedy decoding: + num_beams=4, + # also return probability + output_scores=True, + return_dict_in_generate=True) + if output.sequences_scores is not None: + # log-prob averaged over length + conf = output.sequences_scores.exp().clamp(0.0, 1.0).cpu().numpy() + else: + conf = np.ones(len(output.sequences), dtype=float) + text = proc.batch_decode( + output.sequences, + skip_special_tokens=True, + clean_up_tokenization_spaces=False) + # we must convert to ndarray for Predictor resultq to work + text = np.array(text) + return text, conf + model.predict_on_batch = predict_torch + # not actually needed (image processor does resize itself) + # no batch dimension (images passed as list w/ varying shapes) + model.input_shape = (None, + None, + len(proc.image_processor.image_mean)) + return model def __str__(self): return tabulate( @@ -277,5 +499,5 @@ class EynollahModelZoo: """ if hasattr(self, '_loaded') and getattr(self, '_loaded'): for needle in list(self._loaded.keys()): - self._loaded[needle].shutdown() - del self._loaded[needle] + if isinstance(self._loaded[needle], Predictor): + self._loaded[needle].shutdown() diff --git a/src/eynollah/ocrd-tool.json b/src/eynollah/ocrd-tool.json index 89ed0da..8d9b901 100644 --- a/src/eynollah/ocrd-tool.json +++ b/src/eynollah/ocrd-tool.json @@ -6,23 +6,61 @@ "ocrd-eynollah-segment": { "executable": "ocrd-eynollah-segment", "categories": ["Layout analysis"], - "description": "Segment page into regions and lines and do reading order detection with eynollah", + "description": "Segment page into regions and lines and detect reading order with Eynollah", "input_file_grp_cardinality": 1, "output_file_grp_cardinality": 1, "steps": ["layout/segmentation/region", "layout/segmentation/line"], "parameters": { + "device": { + "type": "string", + "default": "", + "description": "Allocate models to computation device; can be a single device name for all models, or a comma-delimited, colon-tagged mapping from model category to device name (where model categories are binarization, col_classifier, page, textline, region_1_2, region_fl_np, table, reading_order, but can be abbreviated by glob expressions), e.g. col*:CPU,page:GPU0,*:GPU1. If empty, selects the first available GPU for all models." + }, "models": { "type": "string", "format": "uri", "content-type": "text/directory", "cacheable": true, - "description": "Directory containing models to be used (See https://qurator-data.de/eynollah)", + "description": "Directory containing models to be used (See https://huggingface.co/collections/SBB/eynollah-models)", "required": true }, + "model_overrides": { + "type": "object", + "properties": { + "binarization": { + "type": "object" + }, + "enhancement": { + "type": "object" + }, + "col_classifier": { + "type": "object" + }, + "page": { + "type": "object" + }, + "textline": { + "type": "object" + }, + "region_1_2": { + "type": "object" + }, + "region_fl_np": { + "type": "object" + }, + "table": { + "type": "object" + }, + "reading_order": { + "type": "object" + } + }, + "description": "Map model categories to mappings from model variant to model path, e.g. {'ocr': {'tr': 'path/to/my/trocr-model'}}." + }, "dpi": { "type": "number", "format": "float", - "description": "pixel density in dots per inch (overrides any meta-data in the images); ignored if <= 0 (with fall-back 230)", + "description": "ignored (only for backwards-compatibility)", "default": 0 }, "full_layout": { @@ -57,10 +95,20 @@ "default": false, "description": "if true, do not attempt page frame detection (cropping)" }, + "skip_layout_and_reading_order": { + "type": "boolean", + "default": false, + "description": "skip regions, only run textlines, stuffing them in a single text region for the entire page" + }, + "binarize": { + "type": "boolean", + "default": false, + "description": "run adaptive thresholding on input image with dedicated model before analysis (may be useful for degraded input)." + }, "allow_scaling": { "type": "boolean", "default": false, - "description": "check the resolution against the number of detected columns and if needed, scale the image up or down during layout detection (heuristic to improve quality and performance)" + "description": "ignored (only for backwards-compatibility)" }, "allow_enhancement": { "type": "boolean", @@ -72,6 +120,20 @@ "default": false, "description": "if true, return reading order in right-to-left reading direction." }, + "num_col_upper": { + "type": "number", + "format": "integer", + "minimum": 0, + "default": 0, + "description": "Constrain detection of number of columns by this upper boundary (ignored when zero)." + }, + "num_col_lower": { + "type": "number", + "format": "integer", + "minimum": 0, + "default": 0, + "description": "Constrain detection of number of columns by this lower boundary (ignored when zero)." + }, "headers_off": { "type": "boolean", "default": false, @@ -84,6 +146,22 @@ } }, "resources": [ + { + "url": "https://zenodo.org/records/21381102/files/models_inference_layout_v0_9_1.zip", + "name": "models_inference_layout_v0_9_1", + "type": "archive", + "size": 1847700967, + "description": "Models for layout detection, reading order detection, textline detection, page extraction, column classification, table detection, binarization and image enhancement", + "version_range": ">= v0.9.0" + }, + { + "url": "https://zenodo.org/records/21362927/files/models_inference_layout_v0_9_0.zip", + "name": "models_inference_layout_v0_9_0", + "type": "archive", + "size": 1572255489, + "description": "Models for layout detection, reading order detection, textline detection, page extraction, column classification, table detection, binarization and image enhancement", + "version_range": ">= v0.9.0" + }, { "url": "https://zenodo.org/records/17727267/files/models_all_v0_8_0.zip", "name": "models_all_v0_8_0", @@ -143,6 +221,22 @@ } }, "resources": [ + { + "url": "https://zenodo.org/records/21381102/files/models_inference_layout_v0_9_1.zip", + "name": "models_inference_layout_v0_9_1", + "type": "archive", + "size": 1847700967, + "description": "Models for layout detection, reading order detection, textline detection, page extraction, column classification, table detection, binarization and image enhancement", + "version_range": ">= v0.9.0" + }, + { + "url": "https://zenodo.org/records/21381102/files/models_inference_all_v0_9_1.zip", + "name": "models_inference_all_v0_9_1", + "type": "archive", + "size": 3761011909, + "description": "Models for layout detection, reading order detection, textline detection, page extraction, column classification, table detection, binarization and image enhancement", + "version_range": ">= v0.9.0" + }, { "url": "https://zenodo.org/records/17727267/files/models_all_v0_8_0.zip", "name": "models_all_v0_8_0", diff --git a/src/eynollah/ocrd_cli.py b/src/eynollah/ocrd_cli.py index acd8d4e..effecb2 100644 --- a/src/eynollah/ocrd_cli.py +++ b/src/eynollah/ocrd_cli.py @@ -1,10 +1,8 @@ -# NOTE: For predictable order of imports of torch/shapely/tensorflow -# this must be the first import of the CLI! -from .eynollah_imports import imported_libs -from .processor import EynollahProcessor from click import command from ocrd.decorators import ocrd_cli_options, ocrd_cli_wrap_processor +from .processor import EynollahProcessor + @command() @ocrd_cli_options def main(*args, **kwargs): diff --git a/src/eynollah/patch_encoder.py b/src/eynollah/patch_encoder.py index f163132..556a4c9 100644 --- a/src/eynollah/patch_encoder.py +++ b/src/eynollah/patch_encoder.py @@ -6,8 +6,8 @@ from tensorflow.keras import layers, models class PatchEncoder(layers.Layer): # 441=21*21 # 14*14 # 28*28 - def __init__(self, num_patches=441, projection_dim=64): - super().__init__() + def __init__(self, num_patches=441, projection_dim=64, name='encode_patches'): + super().__init__(name=name) self.num_patches = num_patches self.projection_dim = projection_dim self.projection = layers.Dense(self.projection_dim) @@ -23,13 +23,19 @@ class PatchEncoder(layers.Layer): **super().get_config()) class Patches(layers.Layer): - def __init__(self, patch_size_x=1, patch_size_y=1): - super().__init__() + def __init__(self, patch_size_x=1, patch_size_y=1, name='extract_patches'): + super().__init__(name=name) self.patch_size_x = patch_size_x self.patch_size_y = patch_size_y def call(self, images): - batch_size = tf.shape(images)[0] + #batch_size = tf.shape(images)[0] + return tf.map_fn(self.call_single, images) + + def call_single(self, image): + # avoid batched extract_patches: too much memory, + # and variable batch dim not supported by ONNX implementation + images = tf.expand_dims(image, axis=0) patches = tf.image.extract_patches( images=images, sizes=[1, self.patch_size_y, self.patch_size_x, 1], @@ -37,8 +43,9 @@ class Patches(layers.Layer): rates=[1, 1, 1, 1], padding="VALID", ) - patch_dims = patches.shape[-1] - return tf.reshape(patches, [batch_size, -1, patch_dims]) + _, n_rows, n_cols, patch_dims = patches.shape + n_tiles = patches.shape[1] * patches.shape[2] #-1 + return tf.reshape(patches, [1, n_tiles, patch_dims]) def get_config(self): return dict(patch_size_x=self.patch_size_x, diff --git a/src/eynollah/predictor.py b/src/eynollah/predictor.py index e1159e7..0641e38 100644 --- a/src/eynollah/predictor.py +++ b/src/eynollah/predictor.py @@ -1,5 +1,5 @@ from contextlib import ExitStack -from typing import List, Dict +from typing import List, Dict, Sequence, Tuple, Union import logging import logging.handlers import multiprocessing as mp @@ -8,6 +8,7 @@ import numpy as np from .utils.shm import share_ndarray, ndarray_shared QSIZE = 200 +ArrayT = Union[np.ndarray, Tuple[np.ndarray]] class Predictor(mp.context.SpawnProcess): @@ -40,10 +41,10 @@ class Predictor(mp.context.SpawnProcess): def input_shape(self): return self({}) - def predict(self, data: dict, verbose=0): + def predict(self, data: Union[Sequence[ArrayT], ArrayT], verbose=0) -> Union[Sequence[ArrayT], ArrayT]: return self(data) - def __call__(self, data: dict): + def __call__(self, data: Union[Sequence[ArrayT], ArrayT, Dict]) -> Union[ArrayT, Tuple]: # unusable as per python/cpython#79967 #with self.jobid.get_lock(): # would work, but not public: @@ -55,7 +56,15 @@ class Predictor(mp.context.SpawnProcess): self.taskq.put((jobid, data)) #self.logger.debug("sent shape query task '%d' for model '%s'", jobid, self.name) return self.result(jobid) - with share_ndarray(data) as shared_data: + with ExitStack() as stack: + if isinstance(data, tuple): + # multi-input + shared_data = [] + for data0 in data: + shared_data.append(stack.enter_context(share_ndarray(data0))) + shared_data = tuple(shared_data) + else: + shared_data = stack.enter_context(share_ndarray(data)) self.taskq.put((jobid, shared_data)) #self.logger.debug("sent prediction task '%d' for model '%s': %s", jobid, self.name, shared_data) return self.result(jobid) @@ -67,6 +76,14 @@ class Predictor(mp.context.SpawnProcess): result = self.results.pop(jobid) if isinstance(result, Exception): raise Exception(f"predictor {self.name} failed for {jobid}") from result + elif isinstance(result, tuple) and isinstance(result[0], dict): + # multi-output + result1 = [] + for result0 in result: + with ndarray_shared(result0) as shared_result0: + result1.append(np.copy(shared_result0)) + result = result1 + self.closable.append(jobid) elif isinstance(result, dict): with ndarray_shared(result) as shared_result: result = np.copy(shared_result) @@ -111,6 +128,8 @@ class Predictor(mp.context.SpawnProcess): "binarization": 4, "enhancement": 4, "reading_order": 4, + "ocr": 8, + "ocr_tr": 2, # medium size (672x672x3)... "textline": 2, # large models... @@ -126,8 +145,20 @@ class Predictor(mp.context.SpawnProcess): self.resultq.put((jobid, result)) #self.logger.debug("sent result for '%d': %s", jobid, result) else: + if self.name == 'ocr_tr': + # this model takes a list of (image) tensors + # of heterogeneous shape as input, + # resizing them internally; + # so this looks like multi-input + multi_input = True + batch_size = len(shared_data) + elif isinstance(shared_data, tuple): + multi_input = True + batch_size = shared_data[0]['shape'][0] + else: + multi_input = False + batch_size = shared_data['shape'][0] tasks = [(jobid, shared_data)] - batch_size = shared_data['shape'][0] while (not self.taskq.empty() and # climb to target batch size batch_size * len(tasks) < REBATCH_SIZE): @@ -136,7 +167,7 @@ class Predictor(mp.context.SpawnProcess): # add to our batch tasks.append((jobid0, shared_data0)) else: - # immediately anser + # immediately answer self.resultq.put((jobid0, self.model.input_shape)) if len(tasks) > 1: self.logger.debug("rebatching %d '%s' tasks of batch size %d", @@ -147,25 +178,53 @@ class Predictor(mp.context.SpawnProcess): for jobid, shared_data in tasks: #self.logger.debug("predicting '%d' with model '%s': %s", jobid, self.name, shared_data) jobs.append(jobid) - data.append(stack.enter_context(ndarray_shared(shared_data))) - data = np.concatenate(data) + if multi_input: + data.append(tuple(stack.enter_context(ndarray_shared(shared_data0)) + for shared_data0 in shared_data)) + else: + data.append(stack.enter_context(ndarray_shared(shared_data))) + if multi_input: + data = list(np.concatenate(data0) + for data0 in zip(*data)) + else: + data = np.concatenate(data) #result = self.model.predict(data, verbose=0) # faster, less VRAM result = self.model.predict_on_batch(data) - results = np.split(result, len(jobs)) - #self.logger.debug("sharing result array for '%d'", jobid) + def make_shareable(x): + # convert tf.string/np.object to fixed-length bytes + # (because object segfaults in shm) + if x.dtype is np.dtype(object): + # ONNX conversion for some reason decodes bytes into str already + # so here we undo this, too + if x[0].dtype is np.dtype(object) and isinstance(x[0, 0], str): + x = np.char.encode(x.astype(str), 'utf-8') + return x.astype(bytes) + return x + if isinstance(result, (list, tuple)): + multi_output = True + results = zip(*(np.split(make_shareable(result0), len(jobs)) + for result0 in result)) + else: + multi_output = False + results = np.split(make_shareable(result), len(jobs)) with ExitStack() as stack: for jobid, result in zip(jobs, results): + #self.logger.debug("sharing result array for '%d'", jobid) # we don't know when the result will be received, # but don't want to wait either, so track closing # context per job, and wait for closable signal # from client - result = stack.enter_context(share_ndarray(result)) + if multi_output: + result = tuple(stack.enter_context(share_ndarray(result0)) + for result0 in result) + else: + result = stack.enter_context(share_ndarray(result)) closing[jobid] = stack.pop_all() self.resultq.put((jobid, result)) #self.logger.debug("sent result for '%d': %s", jobid, result) except Exception as e: - self.logger.error("prediction for %s failed: %s", self.name, e.__class__.__name__) + self.logger.exception("prediction for %s failed: %s", self.name, e.__class__.__name__) result = e self.resultq.put((jobid, result)) close_all() @@ -174,8 +233,11 @@ class Predictor(mp.context.SpawnProcess): def load_model(self, *load_args, **load_kwargs): assert len(load_args) self.name = '_'.join(list(load_args[:1]) + + list(load_kwargs[key] for key in load_kwargs + if key == 'model_variant') + list(key for key in load_kwargs - if key != 'device')) + if key in ['patched', 'resized'] + and load_kwargs[key])) self.load_args = load_args self.load_kwargs = load_kwargs self.start() # call run() in subprocess @@ -194,17 +256,18 @@ class Predictor(mp.context.SpawnProcess): def shutdown(self): # do not terminate from forked processor instances - if mp.parent_process() is None: + if not hasattr(self, 'model'): self.stopped.set() + self.join() self.taskq.close() self.taskq.cancel_join_thread() self.resultq.close() self.resultq.cancel_join_thread() self.logq.close() - self.terminate() + #self.terminate() else: del self.model def __del__(self): - #self.logger.debug(f"deinit of {self} in {mp.current_process().name}") + #self.logger.debug(f"deinit of {self.name} in {mp.current_process().name}") self.shutdown() diff --git a/src/eynollah/processor.py b/src/eynollah/processor.py index 47fa770..264d364 100644 --- a/src/eynollah/processor.py +++ b/src/eynollah/processor.py @@ -14,7 +14,13 @@ class EynollahProcessor(Processor): def setup(self) -> None: assert self.parameter - model_zoo = EynollahModelZoo(basedir=self.parameter['models']) + basedir = self.resolve_resource(self.parameter['models']) + overrides = [] + for category, override in self.parameter.get('model_overrides', {}).items(): + for variant, path in override.items(): + path = self.resolve_resource(path) + overrides.append((category, variant, path)) + model_zoo = EynollahModelZoo(basedir, model_overrides=overrides) self.eynollah = Eynollah( model_zoo=model_zoo, allow_enhancement=self.parameter['allow_enhancement'], @@ -22,10 +28,15 @@ class EynollahProcessor(Processor): right2left=self.parameter['right_to_left'], reading_order_machine_based=self.parameter['reading_order_machine_based'], ignore_page_extraction=self.parameter['ignore_page_extraction'], + skip_layout_and_reading_order=self.parameter['skip_layout_and_reading_order'], full_layout=self.parameter['full_layout'], allow_scaling=self.parameter['allow_scaling'], headers_off=self.parameter['headers_off'], tables=self.parameter['tables'], + device=self.parameter['device'], + input_binary=self.parameter['binarize'], + num_col_upper=self.parameter['num_col_upper'], + num_col_lower=self.parameter['num_col_lower'], logger=self.logger ) self.eynollah.plotter = None @@ -48,13 +59,17 @@ class EynollahProcessor(Processor): \b - If ``tables``, try to detect table blocks and add them as TableRegion. - - If ``full_layout``, then in addition to paragraphs and marginals, also - try to detect drop capitals and headings. + - If ``full_layout`` (the default), then in addition to paragraphs and marginals, + also try to detect drop capitals and headings. + - If ``ignore_page_extraction``, then attempt no cropping of the page. - If ``ignore_page_extraction``, then attempt no cropping of the page. - If ``curved_line``, then compute contour polygons for text lines instead of simple bounding boxes. - If ``reading_order_machine_based``, then detect reading order via data-driven model instead of geometrical heuristics. + - If ``binarize``, then run internal binarization on the raw image. + - If ``num_col_upper`` or ``num_col_lower`` are non-zero, these will + constrain the column detection (upper or lower bound, respectively). Produce a new output file by serialising the resulting hierarchy. """ diff --git a/src/eynollah/training/cli.py b/src/eynollah/training/cli.py index ae14f04..4b1b391 100644 --- a/src/eynollah/training/cli.py +++ b/src/eynollah/training/cli.py @@ -7,16 +7,11 @@ import sys from .build_model_load_pretrained_weights_and_save import build_model_load_pretrained_weights_and_save from .generate_gt_for_training import main as generate_gt_cli from .inference import main as inference_cli -from .train import ex +from .train import train_cli +from .convert import convert_cli from .extract_line_gt import linegt_cli from .weights_ensembling import ensemble_cli - -@click.command(context_settings=dict( - ignore_unknown_options=True, -)) -@click.argument('SACRED_ARGS', nargs=-1, type=click.UNPROCESSED) -def train_cli(sacred_args): - ex.run_commandline([sys.argv[0]] + list(sacred_args)) +from .generate_or_update_cnn_rnn_ocr_character_list import main as update_ocr_characters_cli @click.group('training') def main(): @@ -26,5 +21,7 @@ main.add_command(build_model_load_pretrained_weights_and_save) main.add_command(generate_gt_cli, 'generate-gt') main.add_command(inference_cli, 'inference') main.add_command(train_cli, 'train') +main.add_command(convert_cli, 'convert') main.add_command(linegt_cli, 'export_textline_images_and_text') main.add_command(ensemble_cli, 'ensembling') +main.add_command(update_ocr_characters_cli, 'generate_or_update_cnn_rnn_ocr_character_list') diff --git a/src/eynollah/training/convert.py b/src/eynollah/training/convert.py new file mode 100644 index 0000000..7351356 --- /dev/null +++ b/src/eynollah/training/convert.py @@ -0,0 +1,115 @@ +import os +from pathlib import Path +from shutil import copy2 +import logging +import json + +import click + +@click.command(context_settings=dict( + help_option_names=['-h', '--help'], + show_default=True)) +@click.option( + "--rebuild", + "-r", + help="build new model from code and then load existing weights (requires input in SavedModel directory format with config.json present)", + is_flag=True +) +@click.option( + "--format", + "-f", + "format_", + help="data format to convert to", + type=click.Choice(["hdf5", "keras", "tf", "tf-serving", "onnx"]), + default="tf" +) +@click.option( + "--in", + "-i", + "in_", + help="path to input model (file in hdf5 / keras format, or directory in tf format)", + required=True, + type=click.Path(exists=True, dir_okay=True) +) +@click.option( + "--out", + "-o", + help="path to output model (file in hdf5 / keras / onnx format, or directory in tf / tf-serving format)", + required=True, + type=click.Path(exists=False, dir_okay=True) +) +def convert_cli(rebuild, format_, in_, out): + """ + convert models for inference + + Load model from path, optionally by rebuilding, convert to output format and write model to path. + """ + 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.keras.models import load_model + from tensorflow.keras.models import Model as KerasModel + + model_path = Path(in_) + config_path = model_path / "config.json" + if model_path.is_dir(): + assert (model_path / "keras_metadata.pb").exists(), ( + "input directory must be Keras model in SavedModel format") + if rebuild: + from .train import ex + from .models import get_model + + assert config_path.exists(), ( + "rebuilding requires input model in SavedModel format with config.json") + + # merge defaults with existing config file + ex.add_config(str(config_path)) + # some models deviate between training and inference + ex.add_config(inference=True) + # OCR models: make sure the local vocab file gets re-used, if available + characters_txt_file = model_path / "characters_org.txt" + if characters_txt_file.exists(): + with open(characters_txt_file, "r") as voc_file: + voc = json.load(voc_file) + ex.add_config(characters_txt_file=characters_txt_file) + ex.add_config(n_classes=len(voc) + 3) + # just retrieve final config (via pseudo-run) + ex.main(lambda: 0) + config = ex.run(options={'--loglevel': 'ERROR'}).config + # use the config to capture the model builder + model = get_model(config, logging.root) + model.load_weights(model_path).assert_existing_objects_matched().expect_partial() + else: + from .models import cnn_rnn_ocr_model4inference + + model = load_model(model_path, compile=False) + + if isinstance(model, KerasModel): + # cnn-rnn-ocr task deviates between training and inference + model = cnn_rnn_ocr_model4inference(model, model_path) + + if format_ in ["hdf5", "keras", "tf"]: + kwargs = {"save_format": {"hdf5": "h5"}.get(format_, format_)} + if format_ != "keras": + kwargs["include_optimizer"] = False + model.save(out, **kwargs) + elif format_ == "tf-serving": + model.export(out) + elif format_ == "onnx": + import tf2onnx + import onnx + tf2onnx.convert.from_keras(model, opset=18, output_path=out) + model = onnx.load(out) + model = onnx.shape_inference.infer_shapes(model, strict_mode=True) + onnx.checker.check_model(model, full_check=True) + onnx.save(model, out) + else: + raise ValueError("unknown output format '%s'" % format_) + + # copy config.json if possible + if config_path.exists() and format_ in ['tf', 'tf-serving']: + copy2(config_path, Path(out) / config_path.name) + + diff --git a/src/eynollah/training/extract_line_gt.py b/src/eynollah/training/extract_line_gt.py index 58fc253..fe9f60d 100644 --- a/src/eynollah/training/extract_line_gt.py +++ b/src/eynollah/training/extract_line_gt.py @@ -50,6 +50,12 @@ from ..utils import is_image_filename is_flag=True, help="if this parameter set to true, cropped textline images will not be masked with textline contour.", ) +@click.option( + "--exclude_vertical_lines", + "-exv", + is_flag=True, + help="if this parameter set to true, vertical textline images will be excluded.", +) def linegt_cli( image, dir_in, @@ -57,6 +63,7 @@ def linegt_cli( dir_out, pref_of_dataset, do_not_mask_with_textline_contour, + exclude_vertical_lines, ): assert bool(dir_in) ^ bool(image), "Set --dir-in or --image-filename, not both" if dir_in: @@ -100,6 +107,9 @@ def linegt_cli( x, y, w, h = cv2.boundingRect(textline_coords) + if exclude_vertical_lines and h > 2 * w: + continue + total_bb_coordinates.append([x, y, w, h]) img_poly_on_img = np.copy(img) diff --git a/src/eynollah/training/generate_gt_for_training.py b/src/eynollah/training/generate_gt_for_training.py index cc5a1b2..a848b65 100644 --- a/src/eynollah/training/generate_gt_for_training.py +++ b/src/eynollah/training/generate_gt_for_training.py @@ -6,6 +6,7 @@ from pathlib import Path from PIL import Image, ImageDraw, ImageFont import cv2 import numpy as np +from eynollah.utils.font import get_font from .gt_gen_utils import ( filter_contours_area_of_image, @@ -393,11 +394,15 @@ def visualize_reading_order(xml_file, dir_xml, dir_out, dir_imgs): layout = np.zeros( (y_len,x_len,3) ) layout = cv2.fillPoly(layout, pts =co_text_all, color=(1,1,1)) - img_file_name_with_format = find_format_of_given_filename_in_dir(dir_imgs, f_name) - img = cv2.imread(os.path.join(dir_imgs, img_file_name_with_format)) - - overlayed = overlay_layout_on_image(layout, img, cx_ordered, cy_ordered, color, thickness) - cv2.imwrite(os.path.join(dir_out, f_name+'.png'), overlayed) + try: + img_file_name_with_format = find_format_of_given_filename_in_dir(dir_imgs, f_name) + img = cv2.imread(os.path.join(dir_imgs, img_file_name_with_format)) + + overlayed = overlay_layout_on_image(layout, img, cx_ordered, cy_ordered, color, thickness) + cv2.imwrite(os.path.join(dir_out, f_name+'.png'), overlayed) + except: + pass + else: img = np.zeros( (y_len,x_len,3) ) @@ -452,14 +457,17 @@ def visualize_textline_segmentation(xml_file, dir_xml, dir_out, dir_imgs): xml_file = os.path.join(dir_xml,ind_xml ) f_name = Path(ind_xml).stem - img_file_name_with_format = find_format_of_given_filename_in_dir(dir_imgs, f_name) - img = cv2.imread(os.path.join(dir_imgs, img_file_name_with_format)) + try: + img_file_name_with_format = find_format_of_given_filename_in_dir(dir_imgs, f_name) + img = cv2.imread(os.path.join(dir_imgs, img_file_name_with_format)) + + co_tetxlines, y_len, x_len = get_textline_contours_for_visualization(xml_file) - co_tetxlines, y_len, x_len = get_textline_contours_for_visualization(xml_file) - - added_image = visualize_image_from_contours(co_tetxlines, img) - - cv2.imwrite(os.path.join(dir_out, f_name+'.png'), added_image) + added_image = visualize_image_from_contours(co_tetxlines, img) + + cv2.imwrite(os.path.join(dir_out, f_name+'.png'), added_image) + except: + pass @@ -509,15 +517,17 @@ def visualize_layout_segmentation(xml_file, dir_xml, dir_out, dir_imgs): f_name = Path(ind_xml).stem print(f_name, 'f_name') - img_file_name_with_format = find_format_of_given_filename_in_dir(dir_imgs, f_name) - img = cv2.imread(os.path.join(dir_imgs, img_file_name_with_format)) + try: + img_file_name_with_format = find_format_of_given_filename_in_dir(dir_imgs, f_name) + img = cv2.imread(os.path.join(dir_imgs, img_file_name_with_format)) + + co_text, co_graphic, co_sep, co_img, co_table, co_map, co_music, co_noise, y_len, x_len = get_layout_contours_for_visualization(xml_file) - co_text, co_graphic, co_sep, co_img, co_table, co_map, co_noise, y_len, x_len = get_layout_contours_for_visualization(xml_file) - - - added_image = visualize_image_from_contours_layout(co_text['paragraph'], co_text['header']+co_text['heading'], co_text['drop-capital'], co_sep, co_img, co_text['marginalia'], co_table, img) + added_image = visualize_image_from_contours_layout(co_text['paragraph'], co_text['header']+co_text['heading'], co_text['drop-capital'], co_sep, co_img, co_text['marginalia'], co_table, co_map, co_music, img) - cv2.imwrite(os.path.join(dir_out, f_name+'.png'), added_image) + cv2.imwrite(os.path.join(dir_out, f_name+'.png'), added_image) + except: + pass @@ -552,8 +562,8 @@ def visualize_ocr_text(xml_file, dir_xml, dir_out): else: xml_files_ind = [xml_file] - font_path = "Charis-7.000/Charis-Regular.ttf" # Make sure this file exists! - font = ImageFont.truetype(font_path, 40) + ###font_path = "Charis-7.000/Charis-Regular.ttf" # Make sure this file exists! + font = get_font(font_size=40)#ImageFont.truetype(font_path, 40) for ind_xml in tqdm(xml_files_ind): indexer = 0 @@ -590,11 +600,11 @@ def visualize_ocr_text(xml_file, dir_xml, dir_out): is_vertical = h > 2*w # Check orientation - font = fit_text_single_line(draw, ocr_texts[index], font_path, w, int(h*0.4) ) + font = fit_text_single_line(draw, ocr_texts[index], w, int(h*0.4) ) if is_vertical: - vertical_font = fit_text_single_line(draw, ocr_texts[index], font_path, h, int(w * 0.8)) + vertical_font = fit_text_single_line(draw, ocr_texts[index], h, int(w * 0.8)) text_img = Image.new("RGBA", (h, w), (255, 255, 255, 0)) # Note: dimensions are swapped text_draw = ImageDraw.Draw(text_img) diff --git a/src/eynollah/training/generate_or_update_cnn_rnn_ocr_character_list.py b/src/eynollah/training/generate_or_update_cnn_rnn_ocr_character_list.py new file mode 100644 index 0000000..8620515 --- /dev/null +++ b/src/eynollah/training/generate_or_update_cnn_rnn_ocr_character_list.py @@ -0,0 +1,59 @@ +import os +import numpy as np +import json +import click +import logging + + + +def run_character_list_update(dir_labels, out, current_character_list): + ls_labels = os.listdir(dir_labels) + ls_labels = [ind for ind in ls_labels if ind.endswith('.txt')] + + if current_character_list: + with open(current_character_list, 'r') as f_name: + characters = json.load(f_name) + + characters = set(characters) + else: + characters = set() + + + for ind in ls_labels: + label = open(os.path.join(dir_labels,ind),'r').read().split('\n')[0] + + for char in label: + characters.add(char) + + + characters = sorted(list(set(characters))) + + with open(out, 'w') as f_name: + json.dump(characters, f_name) + + +@click.command() +@click.option( + "--dir_labels", + "-dl", + help="directory of labels which are .txt files", + type=click.Path(exists=True, file_okay=False), + required=True, +) +@click.option( + "--current_character_list", + "-ccl", + help="existing character list in a .txt file that needs to be updated with a set of labels", + type=click.Path(exists=True, file_okay=True), + required=False, +) +@click.option( + "--out", + "-o", + help="An output .txt file where the generated or updated character list will be written", + type=click.Path(exists=False, file_okay=True), +) + +def main(dir_labels, out, current_character_list): + run_character_list_update(dir_labels, out, current_character_list) + diff --git a/src/eynollah/training/gt_gen_utils.py b/src/eynollah/training/gt_gen_utils.py index 796e896..e085fa0 100644 --- a/src/eynollah/training/gt_gen_utils.py +++ b/src/eynollah/training/gt_gen_utils.py @@ -8,7 +8,7 @@ from shapely import geometry from pathlib import Path from PIL import ImageFont from ocrd_utils import bbox_from_points - +from eynollah.utils.font import get_font KERNEL = np.ones((5, 5), np.uint8) NS = { 'pc': 'http://schema.primaresearch.org/PAGE/gts/pagecontent/2019-07-15' @@ -18,7 +18,7 @@ with warnings.catch_warnings(): warnings.simplefilter("ignore") -def visualize_image_from_contours_layout(co_par, co_header, co_drop, co_sep, co_image, co_marginal, co_table, co_map, img): +def visualize_image_from_contours_layout(co_par, co_header, co_drop, co_sep, co_image, co_marginal, co_table, co_map, co_music, img): alpha = 0.5 blank_image = np.ones( (img.shape[:]), dtype=np.uint8) * 255 @@ -32,6 +32,7 @@ def visualize_image_from_contours_layout(co_par, co_header, co_drop, co_sep, co_ col_marginal = (106, 90, 205) col_table = (0, 90, 205) col_map = (90, 90, 205) + col_music = (90, 90, 0) if len(co_image)>0: cv2.drawContours(blank_image, co_image, -1, col_image, thickness=cv2.FILLED) # Fill the contour @@ -59,6 +60,9 @@ def visualize_image_from_contours_layout(co_par, co_header, co_drop, co_sep, co_ if len(co_map)>0: cv2.drawContours(blank_image, co_map, -1, col_map, thickness=cv2.FILLED) # Fill the contour + + if len(co_music)>0: + cv2.drawContours(blank_image, co_music, -1, col_music, thickness=cv2.FILLED) # Fill the contour img_final =cv2.cvtColor(blank_image, cv2.COLOR_BGR2RGB) @@ -352,11 +356,11 @@ def get_textline_contours_and_ocr_text(xml_file): ocr_textlines.append(ocr_text_in[0]) return co_use_case, y_len, x_len, ocr_textlines -def fit_text_single_line(draw, text, font_path, max_width, max_height): +def fit_text_single_line(draw, text, max_width, max_height): initial_font_size = 50 font_size = initial_font_size while font_size > 10: # Minimum font size - font = ImageFont.truetype(font_path, font_size) + font = get_font(font_size=font_size)# ImageFont.truetype(font_path, font_size) text_bbox = draw.textbbox((0, 0), text, font=font) # Get text bounding box text_width = text_bbox[2] - text_bbox[0] text_height = text_bbox[3] - text_bbox[1] @@ -366,7 +370,7 @@ def fit_text_single_line(draw, text, font_path, max_width, max_height): font_size -= 2 # Reduce font size and retry - return ImageFont.truetype(font_path, 10) # Smallest font fallback + return get_font(font_size=10)#ImageFont.truetype(font_path, 10) # Smallest font fallback def get_layout_contours_for_visualization(xml_file): tree1 = ET.parse(xml_file, parser = ET.XMLParser(encoding='utf-8')) @@ -389,6 +393,7 @@ def get_layout_contours_for_visualization(xml_file): co_img=[] co_table=[] co_map=[] + co_music=[] co_noise=[] types_text = [] @@ -630,6 +635,31 @@ def get_layout_contours_for_visualization(xml_file): elif vv.tag!=link+'Point' and sumi>=1: break co_map.append(np.array(c_t_in)) + + if tag.endswith('}MusicRegion') or tag.endswith('}musicregion'): + #print('sth') + 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 + #print(vv.tag,'in') + elif vv.tag!=link+'Point' and sumi>=1: + break + co_music.append(np.array(c_t_in)) if tag.endswith('}NoiseRegion') or tag.endswith('}noiseregion'): @@ -656,7 +686,7 @@ def get_layout_contours_for_visualization(xml_file): elif vv.tag!=link+'Point' and sumi>=1: break co_noise.append(np.array(c_t_in)) - return co_text, co_graphic, co_sep, co_img, co_table, co_map, co_noise, y_len, x_len + return co_text, co_graphic, co_sep, co_img, co_table, co_map, co_music, co_noise, y_len, x_len def get_images_of_ground_truth( gt_list, @@ -870,7 +900,7 @@ def get_images_of_ground_truth( types_graphic_label = list(types_graphic_dict.values()) - labels_rgb_color = [ (0,0,0), (255,0,0), (255,125,0), (255,0,125), (125,255,125), (125,125,0), (0,125,255), (0,125,0), (125,125,125), (255,0,255), (125,0,125), (0,255,0),(0,0,255), (0,255,255), (255,125,125), (0,125,125), (0,255,125), (255,125,255), (125,255,0), (125,255,255)] + labels_rgb_color = [ (0,0,0), (255,0,0), (255,125,0), (255,0,125), (125,255,125), (125,125,0), (0,125,255), (0,125,0), (125,125,125), (255,0,255), (125,0,125), (0,255,0),(0,0,255), (0,255,255), (255,125,125), (0,125,125), (0,255,125), (255,125,255), (125,255,0), (125,255,255), (125,125,255)] region_tags=np.unique([x for x in alltags if x.endswith('Region')]) @@ -882,6 +912,7 @@ def get_images_of_ground_truth( co_img=[] co_table=[] co_map=[] + co_music=[] co_noise=[] for tag in region_tags: @@ -966,19 +997,21 @@ def get_images_of_ground_truth( if "rest_as_decoration" in types_graphic: types_graphic_without_decoration = [element for element in types_graphic if element!='rest_as_decoration' and element!='decoration'] if len(types_graphic_without_decoration) == 0: - if "type" in nn.attrib: - c_t_in_graphic['decoration'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + #if "type" in nn.attrib: + c_t_in_graphic['decoration'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) elif len(types_graphic_without_decoration) >= 1: if "type" in nn.attrib: if nn.attrib['type'] in types_graphic_without_decoration: c_t_in_graphic[nn.attrib['type']].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) else: c_t_in_graphic['decoration'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) - + else: + c_t_in_graphic['decoration'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) else: if "type" in nn.attrib: if nn.attrib['type'] in all_defined_graphic_types: - c_t_in_graphic[nn.attrib['type']].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + c_t_in_graphic[nn.attrib['type']].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ] for x in p_h] ) ) + break else: @@ -989,9 +1022,9 @@ def get_images_of_ground_truth( if "rest_as_decoration" in types_graphic: types_graphic_without_decoration = [element for element in types_graphic if element!='rest_as_decoration' and element!='decoration'] if len(types_graphic_without_decoration) == 0: - if "type" in nn.attrib: - c_t_in_graphic['decoration'].append( [ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ] ) - sumi+=1 + #if "type" in nn.attrib: + c_t_in_graphic['decoration'].append( [ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ] ) + sumi+=1 elif len(types_graphic_without_decoration) >= 1: if "type" in nn.attrib: if nn.attrib['type'] in types_graphic_without_decoration: @@ -1000,6 +1033,9 @@ def get_images_of_ground_truth( else: c_t_in_graphic['decoration'].append( [ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ] ) sumi+=1 + else: + c_t_in_graphic['decoration'].append( [ int(float(vv.attrib['x'])) , int(float(vv.attrib['y'])) ] ) + sumi+=1 else: if "type" in nn.attrib: @@ -1118,6 +1154,32 @@ def get_images_of_ground_truth( elif vv.tag!=link+'Point' and sumi>=1: break co_map.append(np.array(c_t_in)) + + if 'musicregion' in keys: + if tag.endswith('}MusicRegion') or tag.endswith('}musicregion'): + #print('sth') + 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 + #print(vv.tag,'in') + elif vv.tag!=link+'Point' and sumi>=1: + break + co_music.append(np.array(c_t_in)) if 'noiseregion' in keys: if tag.endswith('}NoiseRegion') or tag.endswith('}noiseregion'): @@ -1195,6 +1257,10 @@ def get_images_of_ground_truth( erosion_rate = 0#2 dilation_rate = 3#4 co_map, img_boundary = update_region_contours(co_map, img_boundary, erosion_rate, dilation_rate, y_len, x_len ) + if "musicregion" in elements_with_artificial_class: + erosion_rate = 0#2 + dilation_rate = 3#4 + co_music, img_boundary = update_region_contours(co_music, img_boundary, erosion_rate, dilation_rate, y_len, x_len ) @@ -1222,6 +1288,8 @@ def get_images_of_ground_truth( img_poly=cv2.fillPoly(img, pts =co_table, color=labels_rgb_color[ config_params['tableregion']]) if 'mapregion' in keys: img_poly=cv2.fillPoly(img, pts =co_map, color=labels_rgb_color[ config_params['mapregion']]) + if 'musicregion' in keys: + img_poly=cv2.fillPoly(img, pts =co_music, color=labels_rgb_color[ config_params['musicregion']]) if 'noiseregion' in keys: img_poly=cv2.fillPoly(img, pts =co_noise, color=labels_rgb_color[ config_params['noiseregion']]) @@ -1286,6 +1354,9 @@ def get_images_of_ground_truth( if 'mapregion' in keys: color_label = config_params['mapregion'] img_poly=cv2.fillPoly(img, pts =co_map, color=(color_label,color_label,color_label)) + if 'musicregion' in keys: + color_label = config_params['musicregion'] + img_poly=cv2.fillPoly(img, pts =co_music, color=(color_label,color_label,color_label)) if 'noiseregion' in keys: color_label = config_params['noiseregion'] img_poly=cv2.fillPoly(img, pts =co_noise, color=(color_label,color_label,color_label)) diff --git a/src/eynollah/training/inference.py b/src/eynollah/training/inference.py index 2be937d..58ccb22 100644 --- a/src/eynollah/training/inference.py +++ b/src/eynollah/training/inference.py @@ -132,15 +132,31 @@ class SBBPredict: self.model = Model( self.model.get_layer(name = "image").input, self.model.get_layer(name = "dense2").output) + assert isinstance(self.model, Model) + elif self.task == "transformer-ocr": + import torch + from transformers import VisionEncoderDecoderModel, TrOCRProcessor + + self.model = VisionEncoderDecoderModel.from_pretrained(self.model_dir) + self.processor = TrOCRProcessor.from_pretrained(self.model_dir) + + if self.cpu: + self.device = torch.device('cpu') + else: + self.device = torch.device('cuda:0') + + self.model.to(self.device) + + assert isinstance(self.model, torch.nn.Module) else: self.model = load_model(self.model_dir, compile=False, custom_objects={"PatchEncoder": PatchEncoder, "Patches": Patches}) + assert isinstance(self.model, Model) ##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': last = self.model.layers[-1] self.img_height = last.output_shape[1] @@ -230,6 +246,13 @@ class SBBPredict: pred_texts = decode_batch_predictions(preds, num_to_char) pred_texts = pred_texts[0].replace("[UNK]", "") return pred_texts + + elif self.task == "transformer-ocr": + from PIL import Image + image = Image.open(image_dir).convert("RGB") + pixel_values = self.processor(image, return_tensors="pt").pixel_values + generated_ids = self.model.generate(pixel_values.to(self.device)) + return self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0] elif self.task == 'reading_order': @@ -566,6 +589,8 @@ class SBBPredict: cv2.imwrite(self.save,res) elif self.task == "cnn-rnn-ocr": print(f"Detected text: {res}") + elif self.task == "transformer-ocr": + print(f"Detected text: {res}") else: img_seg_overlayed, only_layout = self.visualize_model_output(res, self.img_org, self.task) if self.save: @@ -672,7 +697,7 @@ def main(image, dir_in, model, patches, save, save_layout, ground_truth, xml_fil with open(os.path.join(model,'config.json')) as f: config_params_model = json.load(f) task = config_params_model['task'] - if task not in ['classification', 'reading_order', "cnn-rnn-ocr"]: + if task not in ['classification', 'reading_order', "cnn-rnn-ocr", "transformer-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, diff --git a/src/eynollah/training/metrics.py b/src/eynollah/training/metrics.py index caa0e65..2b4fc4f 100644 --- a/src/eynollah/training/metrics.py +++ b/src/eynollah/training/metrics.py @@ -5,7 +5,10 @@ import tensorflow as tf from tensorflow.keras import backend as K from tensorflow.keras.metrics import Metric, MeanMetricWrapper, get from tensorflow.keras.initializers import Zeros -from tensorflow_addons.image import connected_components +try: + from tensorflow_addons.image import connected_components +except ModuleNotFoundError: + pass # n/a beyond TF 2.15 (and only needed for training) ... import numpy as np diff --git a/src/eynollah/training/models.py b/src/eynollah/training/models.py index 3494249..9076f0a 100644 --- a/src/eynollah/training/models.py +++ b/src/eynollah/training/models.py @@ -1,4 +1,5 @@ import os +import json os.environ['TF_USE_LEGACY_KERAS'] = '1' # avoid Keras 3 after TF 2.15 import tensorflow as tf @@ -20,9 +21,11 @@ from tensorflow.keras.layers import ( LSTM, MaxPooling2D, MultiHeadAttention, + Permute, Reshape, UpSampling2D, ZeroPadding2D, + StringLookup, add, concatenate ) @@ -57,6 +60,51 @@ class CTCLayer(Layer): # At test time, just return the computed predictions. return y_pred + +class CTCDecoder(Layer): + def call(self, inputs): + n_samples = tf.shape(inputs)[0] + n_steps = inputs.shape[1] + n_classes = inputs.shape[2] + lengths = tf.ones(n_samples, dtype=tf.int32) * n_steps + ## Keras beam search seems to mess with double letters + ## but Keras greedy sometimes removes arbitrary letters + # outputs, logits = tf.keras.backend.ctc_decode(inputs, + # lengths, + # beam_width=20, + # greedy=False, # True, + # # backend does not allow these kwargs + # #merge_repeated=False, + # #mask_index=inputs.shape[2]-1, + # ) + # tf.nn.ctc_*_decoder (in contrast to tf.keras.backend.ctc_decode) + # needs logits instead of probs and time-major (batch 2nd dim) + inputs = tf.math.log( + tf.transpose(inputs, perm=[1, 0, 2]) + tf.keras.backend.epsilon() + ) + # tf.compat.v1.nn.ctc_beam_search_decoder() also needs merge_repeated=False + # tf.nn.ctc_beam_search_decoder() is not supported by ONNX, yet + # tf.nn.ctc_greedy_decoder() is not as precise, though: + decoded, logits = tf.nn.ctc_beam_search_decoder( + inputs, + lengths, + beam_width=10, + top_paths=1 + ) + # get top path for all sequences in batch + decoded = decoded[0] + logits = logits[:, 0] + probs = tf.exp(-logits / n_steps) + # convert to dense + outputs = tf.SparseTensor(decoded.indices, decoded.values, + (n_samples, n_steps)) + outputs = tf.sparse.to_dense(sp_input=outputs, default_value=n_classes-1) + # # drop non-tokens (-1) and OOV (0) + # result = [] + # for output in outputs: + # result.append(tf.gather(output, tf.where(output > 0))) + # outputs = tf.stack(result) + return outputs, probs def mlp(x, hidden_units, dropout_rate): for units in hidden_units: @@ -309,11 +357,10 @@ def transformer_block(img, # Skip connection 2. encoded_patches = Add()([x3, x2]) - encoded_patches = tf.reshape(encoded_patches, - [-1, - img.shape[1], - img.shape[2], - projection_dim // (patchsize_x * patchsize_y)]) + encoded_patches = Reshape((img.shape[1], + img.shape[2], + projection_dim // (patchsize_x * patchsize_y)), + name="reshape_patches")(encoded_patches) return encoded_patches def vit_resnet50_unet(num_patches, @@ -423,11 +470,11 @@ 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 = Input(shape=(image_height, image_width, 3), name="image") +def cnn_rnn_ocr_model(input_height=None, input_width=None, n_classes=None, max_len=None, inference=False, characters_txt_file=None): + inputs = Input(shape=(input_height, input_width, 3), name="image") labels = Input(name="label", shape=(None,)) - x = Conv2D(64,kernel_size=(3,3),padding="same")(input_img) + x = Conv2D(64,kernel_size=(3,3),padding="same")(inputs) x = BatchNormalization(name="bn1")(x) x = Activation("relu", name="relu1")(x) x = Conv2D(64,kernel_size=(3,3),padding="same")(x) @@ -451,52 +498,142 @@ def cnn_rnn_ocr_model(image_height=None, image_width=None, n_classes=None, max_s x = Activation("relu", name="relu6")(x) x = MaxPooling2D(pool_size=(2,2),strides=(2,2))(x) - x = Conv2D(image_width,kernel_size=(3,3),padding="same")(x) + x = Conv2D(input_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 = Conv2D(input_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 = 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) - + + x = Reshape(new_shape, name="reshape")(x) + x2d = Reshape(new_shape2, name="reshape2")(x2d) + x4d = Reshape(new_shape4, name="reshape4")(x4d) + + xrnnorg = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(x) + xrnn2d = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(x2d) + xrnn4d = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(x4d) + + xrnn2d = Reshape((1, xrnn2d.shape[1], xrnn2d.shape[2]), name="reshape6")(xrnn2d) + xrnn4d = Reshape((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 = Reshape((xrnn2dup.shape[2], xrnn2dup.shape[3]), name="reshape10")(xrnn2dup) + xrnn4dup = Reshape((xrnn4dup.shape[2], xrnn4dup.shape[3]), name="reshape12")(xrnn4dup) addition = Add()([xrnnorg, xrnn2dup, 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) + + addition_rnn = Bidirectional(LSTM(input_width, return_sequences=True, dropout=0.25))(addition) + + #out = Conv1D(max_len, 1, data_format="channels_first")(addition_rnn) + out = Permute((2, 1))(addition_rnn) + out = Conv1D(max_len, 1, data_format="channels_last")(out) + out = Permute((2, 1))(out) out = BatchNormalization(name="bn9")(out) out = Activation("relu", name="relu9")(out) - #out = Conv1D(n_classes, 1, activation='relu', data_format="channels_last")(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 = Model(inputs=(input_img, labels), outputs=output, name="handwriting_recognizer") + if inference: + # add second path for binarization + inputs_bin = Input(shape=(input_height, input_width, 3), name="image_bin") + out_bin = Model(inputs, out)(inputs_bin) + # ensemble raw results + out = 0.5 * (out + out_bin) + # get tf.string batch + out, prob = CTCDecoder()(out) + # decode int to str + with open(characters_txt_file, "r") as voc_file: + voc = json.load(voc_file) + char2num = StringLookup(vocabulary=voc) + voc = char2num.get_vocabulary() + num2char = StringLookup(vocabulary=voc, invert=True) + output = num2char(out) - return model + return Model((inputs, inputs_bin), (output, prob)) + + # Add CTC layer for calculating CTC loss at each step. + out = CTCLayer(name="ctc_loss")(labels, out) + + return Model((inputs, labels), out) + +def cnn_rnn_ocr_model4inference(model, model_path): + """convert trained cnn-rnn-ocr model to inference model post-hoc""" + try: + model.get_layer(name='ctc_loss') + except ValueError: + # likely already converted + return model + else: + inputs = model.get_layer(name='image').input + output = model.get_layer(name='dense2').output + inputs_bin = Input(inputs.shape[1:], name='image_bin') + output_bin = Model(inputs, output)(inputs_bin) + output = 0.5 * (output + output_bin) + output, prob = CTCDecoder()(output) + with open(model_path / "characters_org.txt", "r") as voc_file: + voc = json.load(voc_file) + char2num = StringLookup(vocabulary=voc) + voc = char2num.get_vocabulary() + num2char = StringLookup(vocabulary=voc, invert=True) + output = num2char(output) + inputs = (inputs, inputs_bin) + outputs = (output, prob) + return Model(inputs, outputs) + +def get_model(config, logger): + from sacred.config import create_captured_function + + task = config['task'] + if task in ["segmentation", "enhancement", "binarization"]: + if config['backbone_type'] == 'nontransformer': + builder = resnet50_unet + else: + num_patches_x, num_patches_y = config['transformer_num_patches_xy'] + num_patches = num_patches_x * num_patches_y + + if config['transformer_cnn_first']: + builder = vit_resnet50_unet + multiple = 32 + else: + builder = vit_resnet50_unet_transformer_before_cnn + multiple = 1 + + assert config['input_height'] == ( + num_patches_y * config['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 config['input_width'] == ( + num_patches_x * config['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 == (config['transformer_projection_dim'] % + (config['transformer_patchsize_y'] * + config['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") + + config['num_patches'] = num_patches + elif task == "cnn-rnn-ocr": + builder = cnn_rnn_ocr_model + elif task=='classification': + builder = resnet50_classifier + elif task=='reading_order': + builder = machine_based_reading_order_model + else: + raise ValueError("unknown model task '%s'" % task) + + builder = create_captured_function(builder) + builder.config = config + builder.logger = logger + return builder() diff --git a/src/eynollah/training/reload-models-v0.8.mk b/src/eynollah/training/reload-models-v0.8.mk index b7a38dd..eab5ffb 100644 --- a/src/eynollah/training/reload-models-v0.8.mk +++ b/src/eynollah/training/reload-models-v0.8.mk @@ -4,38 +4,71 @@ MODELS_SRC = models_eynollah MODELS_DST = reloaded/models_eynollah -# $(MODELS_DST)/eynollah-binarization_20210425 \ -# $(MODELS_DST)/eynollah-column-classifier_20210425 \ -# $(MODELS_DST)/eynollah-enhancement_20210425 \ -# $(MODELS_DST)/eynollah-main-regions-aug-rotation_20210425 \ -# $(MODELS_DST)/eynollah-main-regions-aug-scaling_20210425 \ -# $(MODELS_DST)/eynollah-main-regions-ensembled_20210425 \ -# $(MODELS_DST)/eynollah-main-regions_20220314 \ -# $(MODELS_DST)/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 \ -# $(MODELS_DST)/eynollah-tables_20210319 \ -# $(MODELS_DST)/model_eynollah_ocr_cnnrnn_20250930 \ +# eynollah-main-regions-aug-rotation_20210425 +# eynollah-main-regions-aug-scaling_20210425 +# eynollah-main-regions-ensembled_20210425 +# eynollah-main-regions_20220314 +# eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18 +# eynollah-tables_20210319 -RELOADABLE_MODELS = \ - $(MODELS_DST)/model_eynollah_page_extraction_20250915 \ - $(MODELS_DST)/model_eynollah_reading_order_20250824 \ - $(MODELS_DST)/modelens_e_l_all_sp_0_1_2_3_4_171024 \ - $(MODELS_DST)/modelens_full_lay_1__4_3_091124 \ - $(MODELS_DST)/modelens_table_0t4_201124 \ - $(MODELS_DST)/modelens_textline_0_1__2_4_16092024 +CURRENT_MODELS := +CURRENT_MODELS += model_eynollah_page_extraction_20250915 +CURRENT_MODELS += model_eynollah_reading_order_20250824 +CURRENT_MODELS += modelens_e_l_all_sp_0_1_2_3_4_171024 +CURRENT_MODELS += modelens_full_lay_1__4_3_091124 +CURRENT_MODELS += modelens_table_0t4_201124 +CURRENT_MODELS += modelens_textline_0_1__2_4_16092024 +CURRENT_MODELS += model_eynollah_ocr_cnnrnn_20250930 +CURRENT_MODELS += eynollah-binarization_20210425 +CURRENT_MODELS += eynollah-column-classifier_20210425 +CURRENT_MODELS += eynollah-enhancement_20210425 -all: $(RELOADABLE_MODELS) +# tf (SavedModel format) for training +# onnx conversion for fast inference +all: tf onnx +tf-serving: $(CURRENT_MODELS:%=$(MODELS_DST)/%) +tf: $(CURRENT_MODELS:%=$(MODELS_DST)/%) +keras: $(CURRENT_MODELS:%=$(MODELS_DST)/%.keras) +hdf5: $(CURRENT_MODELS:%=$(MODELS_DST)/%.h5) +onnx: $(CURRENT_MODELS:%=$(MODELS_DST)/%.onnx) + +# distinguish tf from tf-serving: target pattern is the same, +# so check if either is current goal, otherwise assumg tf +$(MODELS_DST)/%: FORMAT = $(or $(filter tf-serving,$(MAKECMDGOALS)), tf) $(MODELS_DST)/%: $(MODELS_SRC)/% - mkdir -p $@ - test -e $&1 | tee $(notdir $<).log - cp $ $(notdir $<).$(FORMAT).log 2>&1 || { cat $(notdir $<).$(FORMAT).log; false; } + +$(MODELS_DST)/%.keras: $(MODELS_SRC)/% + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).keras.log 2>&1 || { cat $(notdir $<).keras.log; false; } + +$(MODELS_DST)/%.h5: $(MODELS_SRC)/% + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).hdf5.log 2>&1 || { cat $(notdir $<).hdf5.log; false; } + +$(MODELS_DST)/%.onnx: $(MODELS_SRC)/% + $(if $(findstring _ocr,$@),$(warning skipping $@: OCR CTC decoder is buggy in ONNX) : )\ + eynollah-training convert \ + $(and $(wildcard $ $(notdir $<).onnx.log 2>&1 || { cat $(notdir $<).onnx.log; false; } compare: for i in `find $(MODELS_DST) -mindepth 2`;do \ @@ -43,6 +76,5 @@ compare: du -bs $$n $$i ; \ done - clear: rm -rf $(MODELS_DST) diff --git a/src/eynollah/training/train.py b/src/eynollah/training/train.py index de998fd..a6b5331 100644 --- a/src/eynollah/training/train.py +++ b/src/eynollah/training/train.py @@ -2,6 +2,8 @@ import os import sys import io import json +import click +from typing import Optional from tqdm import tqdm import requests @@ -17,7 +19,6 @@ 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 @@ -32,16 +33,9 @@ from .metrics import ( connected_components_loss, ) from .models import ( - PatchEncoder, - Patches, - machine_based_reading_order_model, - resnet50_classifier, - resnet50_unet, - vit_resnet50_unet, - vit_resnet50_unet_transformer_before_cnn, - cnn_rnn_ocr_model, RESNET50_WEIGHTS_PATH, - RESNET50_WEIGHTS_URL + RESNET50_WEIGHTS_URL, + get_model ) from .utils import ( generate_arrays_from_folder_reading_order, @@ -355,10 +349,9 @@ def config_params(): 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}") - reload_weights = False # Set true to build new model from config, load weights from dir_of_start_model, save under dir_output and exit. continue_training = False # Whether to continue training an existing model. + dir_of_start_model = '' # Directory of model checkpoint to load to continue training or load weights from. (E.g. if you already trained for 3 epochs, set "dir_of_start_model=dir_output/model_03".) 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). @@ -379,7 +372,6 @@ def run(_config, weight_decay, learning_rate, continue_training, - reload_weights, save_interval, augmentation, # dependent config keys need a default, @@ -406,7 +398,7 @@ def run(_config, f1_threshold_classification=None, classification_classes_name=None, ## if task=cnn-rnn-ocr - characters_txt_file=None, + characters_txt_file: Optional[str]=None, color_padding_rotation=False, thetha_padd=None, bin_deg=False, @@ -477,58 +469,15 @@ def run(_config, 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: - 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) + model = load_model(dir_of_start_model, compile=False) else: index_start = 0 - 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: - model_builder = vit_resnet50_unet - multiple = 32 - else: - 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) + model = get_model(_config, _log) + if dir_of_start_model: + model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() + _log.info("reloaded weights from %s", dir_of_start_model) assert model is not None #if you want to see the model structure just uncomment model summary. @@ -559,15 +508,6 @@ def run(_config, optimizer=Adam(learning_rate=learning_rate), metrics=metrics) - if reload_weights: - model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() - dir_save = os.path.join(dir_output, os.path.basename(os.path.normpath(dir_of_start_model))) - model.save(dir_save, include_optimizer=False) - with open(os.path.join(dir_save, "config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - _log.info("reloaded model from %s to %s", dir_of_start_model, dir_save) - return - if not data_is_provided: # first create a directory in output for both training and evaluations # in order to flow data from these directories. @@ -708,10 +648,11 @@ def run(_config, 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) + model = get_model(_config, _log) + if dir_of_start_model: + model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() + _log.info("reloaded weights from %s", dir_of_start_model) + #initial_learning_rate = 1e-4 #decay_steps = int (n_epochs * ( len_dataset / n_batch )) #alpha = 0.01 @@ -722,15 +663,6 @@ def run(_config, #print(model.summary()) - if reload_weights: - model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() - dir_save = os.path.join(dir_output, os.path.basename(os.path.normpath(dir_of_start_model))) - model.save(dir_save, include_optimizer=False) - with open(os.path.join(dir_save, "config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - _log.info("reloaded model from %s to %s", dir_of_start_model, dir_save) - return - # todo: use Dataset.map() on Dataset.list_files() def get_dataset(dir_img, dir_lab): def gen(): @@ -767,30 +699,93 @@ def run(_config, callbacks=callbacks, initial_epoch=index_start) + elif task=="transformer-ocr": + import torch + from torch.utils.data import Dataset as TorchDataset + from transformers import TrOCRProcessor, VisionEncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator + dir_img, dir_lab = get_dirs_or_files(dir_train) + + if continue_training: + model = VisionEncoderDecoderModel.from_pretrained(dir_of_start_model) + else: + model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed") + + processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed") + + + # Create a DataLoader + class TransformerOCRTorchDataset(TorchDataset): + """ + Wraps preprocess_imgs in a format consumable by torch + """ + def __init__(self, config, dir_img, dir_lab): + self.samples = preprocess_imgs( + config, + dir_img, + dir_lab, + processor=processor, + ) + + def __len__(self): + return len(self.samples) + + def __iter__(self): + yield from self.samples + + dataset = TransformerOCRTorchDataset(_config, dir_img, dir_lab) + data_loader = torch.utils.data.DataLoader(dataset, batch_size=1) + train_dataset = data_loader.dataset + + # set special tokens used for creating the decoder_input_ids from the labels + model.config.decoder_start_token_id = processor.tokenizer.cls_token_id + model.config.pad_token_id = processor.tokenizer.pad_token_id + # make sure vocab size is set correctly + model.config.vocab_size = model.config.decoder.vocab_size + + # set beam search parameters + model.config.eos_token_id = processor.tokenizer.sep_token_id + model.config.max_length = max_len + model.config.early_stopping = True + model.config.no_repeat_ngram_size = 3 + model.config.length_penalty = 2.0 + model.config.num_beams = 4 + + + training_args = Seq2SeqTrainingArguments( + predict_with_generate=True, + num_train_epochs=n_epochs, + learning_rate=learning_rate, + per_device_train_batch_size=n_batch, + fp16=True, + output_dir=dir_output, + logging_steps=2, + save_steps=save_interval, + ) + + # instantiate trainer + trainer = Seq2SeqTrainer( + model=model, + tokenizer=processor.feature_extractor, + args=training_args, + train_dataset=train_dataset, + data_collator=default_data_collator, + ) + trainer.train() + elif task=='classification': 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) + model = get_model(_config, _log) + if dir_of_start_model: + model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() + _log.info("reloaded weights from %s", dir_of_start_model) model.compile(loss='categorical_crossentropy', optimizer=Adam(learning_rate=0.001), # rs: why not learning_rate? metrics=['accuracy', F1Score(average='macro', name='f1')]) - if reload_weights: - model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() - dir_save = os.path.join(dir_output, os.path.basename(os.path.normpath(dir_of_start_model))) - model.save(dir_save, include_optimizer=False) - with open(os.path.join(dir_save, "config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - _log.info("reloaded model from %s to %s", dir_of_start_model, dir_save) - return - list_classes = list(classification_classes_name.values()) data_args = dict(label_mode="categorical", class_names=list_classes, @@ -820,7 +815,7 @@ def run(_config, 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) + run_ensembling(usable_checkpoints, ens_path, framework='tensorflow') _log.info("ensemble model saved under '%s'", ens_path) elif task=='reading_order': @@ -828,11 +823,10 @@ def run(_config, 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) + model = get_model(_config, _log) + if dir_of_start_model: + model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() + _log.info("reloaded weights from %s", dir_of_start_model) #f1score_tot = [0] model.compile(loss="binary_crossentropy", @@ -840,15 +834,6 @@ def run(_config, optimizer=Adam(learning_rate=0.0001), # rs: why not learning_rate? metrics=['accuracy']) - if reload_weights: - model.load_weights(dir_of_start_model).assert_existing_objects_matched().expect_partial() - dir_save = os.path.join(dir_output, os.path.basename(os.path.normpath(dir_of_start_model))) - model.save(dir_save, include_optimizer=False) - with open(os.path.join(dir_save, "config.json"), "w") as fp: - json.dump(_config, fp) # encode dict into JSON - _log.info("reloaded model from %s to %s", dir_of_start_model, dir_save) - return - dir_flow_train_imgs = os.path.join(dir_train, 'images') dir_flow_train_labels = os.path.join(dir_train, 'labels') @@ -881,3 +866,23 @@ def run(_config, model_dir = os.path.join(dir_out,'model_best') model.save(model_dir) ''' + +@click.command(context_settings=dict( + ignore_unknown_options=True, +)) +@click.argument('SACRED_ARGS', nargs=-1, type=click.UNPROCESSED) +def train_cli(sacred_args): + """ + train model on extracted GT + + SACRED_ARGS as per CLI interface of Sacred, cf. + https://sacred.readthedocs.io/en/stable/command_line.html: + + \b + To configure the learning task, pass the string `with`, + followed by any number of + - config JSON file paths + - parameter overrides in the form of key=value + (where the later settings will override the former). + """ + ex.run_commandline([sys.argv[0]] + list(sacred_args)) diff --git a/src/eynollah/training/utils.py b/src/eynollah/training/utils.py index 33a1fd2..d04a17b 100644 --- a/src/eynollah/training/utils.py +++ b/src/eynollah/training/utils.py @@ -789,7 +789,7 @@ def preprocess_imgs(config, 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": + elif config['task'] in ["cnn-rnn-ocr", "transformer-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] @@ -797,7 +797,7 @@ def preprocess_imgs(config, lab = None try: - if config['task'] == "cnn-rnn-ocr": + if config['task'] in ["cnn-rnn-ocr", "transformer-ocr"]: yield from preprocess_img_ocr(img, img_name, lab, **config) continue else: @@ -1116,14 +1116,33 @@ def preprocess_img_ocr( number_of_backgrounds_per_image=None, list_all_possible_background_images=None, list_all_possible_foreground_rgbs=None, + task=None, + processor=None, **kwargs ): 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 + if task == 'cnn-rnn-ocr': + assert char_to_num, 'task is cnn-rnn-ocr, so preprocess_imgs_ocr should be passed "char_to_num"' + lab = char_to_num(tf.strings.unicode_split(lab, input_encoding="UTF-8")) + def yield_encoder(x): + return x + elif task == 'transformer-ocr': + import torch + assert processor, 'task is transformer-ocr, so preprocess_imgs_ocr should be passed "processor"' + # TODO make max_length configurable again, if deemed sensible + lab = [tok if tok != processor.tokenizer.pad_token_id else -100 + for tok in processor.tokenizer(lab, + padding="max_length", + max_length=128 + ).input_ids] + def yield_encoder(img_, lab_): + return {"pixel_values": processor(Image.fromarray(img_), + return_tensors="pt").pixel_values.squeeze(), + "labels": torch.tensor(lab_)} + yield yield_encoder(scale_image(img), lab) #to_yield = {"image": ret_x, "label": ret_y} if dir_img_bin: @@ -1139,32 +1158,32 @@ def preprocess_img_ocr( 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 + yield yield_encoder(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 + yield yield_encoder(scale_image(img_rot), lab) if blur_aug: for blur_type in blur_k: img_blur = bluring(img, blur_type) - yield scale_image(img_blur), lab + yield yield_encoder(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 + yield yield_encoder(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 + yield yield_encoder(scale_image(img_deg), lab) if brightening: for bright_scale_ind in brightness: img_bright = do_brightening(img, bright_scale_ind) - yield scale_image(img_bright), lab + yield yield_encoder(scale_image(img_bright), lab) if padding_white: 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 + yield yield_encoder(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) @@ -1178,7 +1197,7 @@ def preprocess_img_ocr( 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 + yield yield_encoder(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) @@ -1186,59 +1205,59 @@ def preprocess_img_ocr( 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 + yield yield_encoder(scale_image(img_bg), lab) if binarization: - yield scale_image(img_bin_corr), lab + yield yield_encoder(scale_image(img_bin_corr), lab) if image_inversion: img_inv = invert_image(img_bin_corr) - yield scale_image(img_inv), lab + yield yield_encoder(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 + yield yield_encoder(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 + yield yield_encoder(scale_image(img_red), lab) if white_noise_strap: img_noisy = return_image_with_strapped_white_noises(img) - yield scale_image(img_noisy), lab + yield yield_encoder(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 + yield yield_encoder(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 + yield yield_encoder(scale_image(img_rot), lab) if textline_left_in_depth: img_warp = do_direction_in_depth(img, 'left') - yield scale_image(img_warp), lab + yield yield_encoder(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 + yield yield_encoder(scale_image(img_warp), lab) if textline_right_in_depth: img_warp = do_direction_in_depth(img, 'right') - yield scale_image(img_warp), lab + yield yield_encoder(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 + yield yield_encoder(scale_image(img_warp), lab) if textline_up_in_depth: img_warp = do_direction_in_depth(img, 'up') - yield scale_image(img_warp), lab + yield yield_encoder(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 + yield yield_encoder(scale_image(img_warp), lab) if textline_down_in_depth: img_warp = do_direction_in_depth(img, 'down') - yield scale_image(img_warp), lab + yield yield_encoder(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 + yield yield_encoder(scale_image(img_warp), lab) if pepper_aug: for pepper_ind in pepper_indexes: img_noisy = add_salt_and_pepper_noise(img, pepper_ind, pepper_ind) - yield scale_image(img_noisy), lab + yield yield_encoder(scale_image(img_noisy), lab) if pepper_bin_aug: 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 + yield yield_encoder(scale_image(img_noisy), lab) diff --git a/src/eynollah/training/weights_ensembling.py b/src/eynollah/training/weights_ensembling.py index e3ede24..1c175c0 100644 --- a/src/eynollah/training/weights_ensembling.py +++ b/src/eynollah/training/weights_ensembling.py @@ -1,4 +1,5 @@ import os +from typing import Optional from warnings import catch_warnings, simplefilter import click @@ -11,38 +12,62 @@ from ocrd_utils import tf_disable_interactive_logs tf_disable_interactive_logs() import tensorflow as tf from tensorflow.keras.models import load_model +import torch +from transformers import VisionEncoderDecoderModel from ..patch_encoder import ( PatchEncoder, Patches, ) -def run_ensembling(model_dirs, out_dir): - all_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()) +def run_ensembling(dir_models, out, framework): + ls_models = os.listdir(dir_models) + # model: Optional[VisionEncoderDecoderModel] = None + # model_name: Optional[str] = None + if framework=="torch": + models = [] + sd_models = [] - 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 model_name in ls_models: + model = VisionEncoderDecoderModel.from_pretrained(os.path.join(dir_models, model_name)) + models.append(model) + sd_models.append(model.state_dict()) + for key in sd_models[0]: + sd_models[0][key] = sum(sd[key] for sd in sd_models) / len(sd_models) + + model.load_state_dict(sd_models[0]) + os.system("mkdir "+out) + torch.save(model.state_dict(), os.path.join(out, "pytorch_model.bin")) + os.system('cp ' + os.path.join(os.path.join(dir_models, model_name), "config.json") + " " + out) + + else: + weights=[] - #model = tf.keras.models.clone_model(model) - model.set_weights(new_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()) + + new_weights = list() - model.save(out_dir) - os.system('cp ' + os.path.join(model_dirs[0], "config.json ") + out_dir + "/") + 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.set_weights(new_weights) + model.save(out) + os.system('cp '+os.path.join(os.path.join(dir_models, model_name), "config.json") + " " + out) + os.system('cp '+os.path.join(os.path.join(dir_models, model_name), "characters_org.txt") + " " + out) @click.command() @click.option( "--in", "-i", + "in_", help="input directory of checkpoint models to be read", multiple=True, required=True, @@ -55,12 +80,19 @@ def run_ensembling(model_dirs, out_dir): required=True, type=click.Path(exists=False, file_okay=False), ) -def ensemble_cli(in_, out): +@click.option( + "--framework", + "-fw", + help="this parameter gets tensorflow or torch as model framework", + type=click.Choice(['torch', 'tensorflow']), + default="tensorflow" +) + +def ensemble_cli(in_, out, framework): """ 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) - + run_ensembling(in_, out, framework) diff --git a/src/eynollah/utils/__init__.py b/src/eynollah/utils/__init__.py index 47a765c..15e8110 100644 --- a/src/eynollah/utils/__init__.py +++ b/src/eynollah/utils/__init__.py @@ -2,6 +2,7 @@ from typing import Iterable, List, Tuple from logging import getLogger import time import math +from itertools import islice try: import matplotlib.pyplot as plt @@ -33,6 +34,11 @@ def pairwise(iterable): yield a, b a = b +def batched(iterable, n): + iterator = iter(iterable) + while batch := tuple(islice(iterator, n)): + yield batch + 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): @@ -1660,37 +1666,6 @@ def return_boxes_of_images_by_order_of_reading_new( #print(peaks_neg_tot,'peaks_neg_tot') peaks_neg_tot_tables.append(peaks_neg_tot) - all_columns = set(range(len(peaks_neg_tot) - 1)) - #print("all_columns", all_columns) - - # 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 = sep_mask[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 - args_hor = matrix_new[:, 9] == 0 x_min_hor_some = matrix_new[:, 2][args_hor] x_max_hor_some = matrix_new[:, 3][args_hor] @@ -1821,8 +1796,8 @@ def return_boxes_of_images_by_order_of_reading_new( (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(text_mask[ - y_top: y_bot, peaks_neg_tot[start + 1]]) > 0.1 * (y_bot - y_top)), + and text_mask[y_top: y_bot, + peaks_neg_tot[start + 1]].sum() > 0.1 * (y_bot - y_top)), # Otherwise advance only 1 column. default=start + 1) def add_sep(cur): diff --git a/src/eynollah/utils/contour.py b/src/eynollah/utils/contour.py index f1a7a8e..83243ca 100644 --- a/src/eynollah/utils/contour.py +++ b/src/eynollah/utils/contour.py @@ -11,7 +11,7 @@ from shapely.geometry.polygon import orient from shapely import set_precision, affinity from shapely.ops import unary_union, nearest_points -from .rotate import rotate_image, rotation_image_new +from .rotate import rotate_image def contours_in_same_horizon(cy_main_hor): """ @@ -120,94 +120,6 @@ def return_contours_of_interested_region(region_pre_p, label, min_area=0.0002, d dilate=dilate) return contours_imgs -def do_work_of_contours_in_image(contour, index_r_con, img, slope_first): - img_copy = np.zeros(img.shape[:2], dtype=np.uint8) - img_copy = cv2.fillPoly(img_copy, pts=[contour], color=1) - - img_copy = rotation_image_new(img_copy, -slope_first) - _, thresh = cv2.threshold(img_copy, 0, 255, 0) - - cont_int, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - - cont_int[0][:, 0, 0] = cont_int[0][:, 0, 0] + np.abs(img_copy.shape[1] - img.shape[1]) - cont_int[0][:, 0, 1] = cont_int[0][:, 0, 1] + np.abs(img_copy.shape[0] - img.shape[0]) - - return cont_int[0], index_r_con - -def get_textregion_contours_in_org_image_multi(cnts, img, slope_first, map=map): - if not len(cnts): - return [], [] - results = map(partial(do_work_of_contours_in_image, - img=img, - slope_first=slope_first, - ), - cnts, range(len(cnts))) - return tuple(zip(*results)) - -def get_textregion_contours_in_org_image(cnts, img, slope_first): - cnts_org = [] - # print(cnts,'cnts') - for i in range(len(cnts)): - img_copy = np.zeros(img.shape[:2], dtype=np.uint8) - img_copy = cv2.fillPoly(img_copy, pts=[cnts[i]], color=1) - - # plt.imshow(img_copy) - # plt.show() - - # print(img.shape,'img') - img_copy = rotation_image_new(img_copy, -slope_first) - ##print(img_copy.shape,'img_copy') - # plt.imshow(img_copy) - # plt.show() - - _, thresh = cv2.threshold(img_copy, 0, 255, 0) - - cont_int, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - cont_int[0][:, 0, 0] = cont_int[0][:, 0, 0] + np.abs(img_copy.shape[1] - img.shape[1]) - cont_int[0][:, 0, 1] = cont_int[0][:, 0, 1] + np.abs(img_copy.shape[0] - img.shape[0]) - # print(np.shape(cont_int[0])) - cnts_org.append(cont_int[0]) - - return cnts_org - -def get_textregion_confidences_old(cnts, img, slope_first): - zoom = 3 - img = cv2.resize(img, (img.shape[1] // zoom, - img.shape[0] // zoom), - interpolation=cv2.INTER_NEAREST) - cnts_org = [] - for cnt in cnts: - img_copy = np.zeros(img.shape[:2], dtype=np.uint8) - img_copy = cv2.fillPoly(img_copy, pts=[cnt // zoom], color=1) - - img_copy = rotation_image_new(img_copy, -slope_first).astype(np.uint8) - _, thresh = cv2.threshold(img_copy, 0, 255, 0) - - cont_int, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - cont_int[0][:, 0, 0] = cont_int[0][:, 0, 0] + np.abs(img_copy.shape[1] - img.shape[1]) - cont_int[0][:, 0, 1] = cont_int[0][:, 0, 1] + np.abs(img_copy.shape[0] - img.shape[0]) - cnts_org.append(cont_int[0] * zoom) - - return cnts_org - -def do_back_rotation_and_get_cnt_back(contour_par, index_r_con, img, slope_first, confidence_matrix): - img_copy = np.zeros(img.shape[:2], dtype=np.uint8) - img_copy = cv2.fillPoly(img_copy, pts=[contour_par], color=1) - confidence_matrix_mapped_with_contour = confidence_matrix * img_copy - confidence_contour = np.sum(confidence_matrix_mapped_with_contour) / float(np.sum(img_copy)) - - img_copy = rotation_image_new(img_copy, -slope_first).astype(np.uint8) - _, thresh = cv2.threshold(img_copy, 0, 255, 0) - - cont_int, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) - if len(cont_int)==0: - cont_int = [contour_par] - confidence_contour = 0 - else: - cont_int[0][:, 0, 0] = cont_int[0][:, 0, 0] + np.abs(img_copy.shape[1] - img.shape[1]) - cont_int[0][:, 0, 1] = cont_int[0][:, 0, 1] + np.abs(img_copy.shape[0] - img.shape[0]) - return cont_int[0], index_r_con, confidence_contour - def get_region_confidences(cnts, confidence_matrix): if not len(cnts): return [] @@ -418,7 +330,7 @@ def estimate_skew_contours(contours): if not np.any(usable): raise ValueError("not enough contours with consistent length") if np.count_nonzero(usable) == 1: - return angle_in[usable] + return angle_in[usable][0] # 4. there is no way to distinguish between +90 and -89.9 here, # so map to [0,180] when calculating averages, then map back to [-90,90] # (we don't want -90 and +89 to average zero, or +1 and +179 to average 90) @@ -442,13 +354,7 @@ def contour2polygon(contour: Union[np.ndarray, Sequence[Sequence[Sequence[Number polygon = Polygon([point[0] for point in contour]) if dilate: polygon = polygon.buffer(dilate) - if polygon.geom_type == 'GeometryCollection': - # heterogeneous result: filter zero-area shapes (LineString, Point) - polygon = unary_union([geom for geom in polygon.geoms if geom.area > 0]) - if polygon.geom_type == 'MultiPolygon': - # homogeneous result: construct convex hull to connect - polygon = join_polygons(polygon.geoms) - return make_valid(polygon) + return ensure_polygon(make_valid(polygon)) def polygon2contour(polygon: Polygon) -> np.ndarray: polygon = np.array(polygon.exterior.coords[:-1], dtype=int) @@ -459,16 +365,22 @@ def make_intersection(poly1, poly2): # post-process if interp.is_empty or interp.area == 0.0: return None - if interp.geom_type == 'GeometryCollection': - # heterogeneous result: filter zero-area shapes (LineString, Point) - interp = unary_union([geom for geom in interp.geoms if geom.area > 0]) - if interp.geom_type == 'MultiPolygon': - # homogeneous result: construct convex hull to connect - interp = join_polygons(interp.geoms) - assert interp.geom_type == 'Polygon', interp.wkt + interp = ensure_polygon(interp) interp = make_valid(interp) + interp = ensure_polygon(interp) return interp +def ensure_polygon(geometry): + if geometry.geom_type == 'GeometryCollection': + # heterogeneous result: filter zero-area shapes (LineString, Point) + geometry = unary_union([geom for geom in geometry.geoms if geom.area > 0]) + if geometry.geom_type == 'MultiPolygon': + # homogeneous result: construct convex hull to connect + geometry = join_polygons(geometry.geoms) + poly = Polygon(geometry) + assert poly.geom_type == 'Polygon', poly.wkt + return poly + def make_valid(polygon: Polygon) -> Polygon: """Ensures shapely.geometry.Polygon object is valid by repeated rearrangement/simplification/enlargement.""" def isint(x): diff --git a/src/eynollah/utils/font.py b/src/eynollah/utils/font.py index 939933e..3e9e588 100644 --- a/src/eynollah/utils/font.py +++ b/src/eynollah/utils/font.py @@ -9,8 +9,8 @@ else: import importlib.resources as importlib_resources -def get_font(): +def get_font(font_size): #font_path = "Charis-7.000/Charis-Regular.ttf" # Make sure this file exists! - font = importlib_resources.files(__package__) / "../Charis-Regular.ttf" + font = importlib_resources.files(__package__) / "../Amiri-Regular.ttf" with importlib_resources.as_file(font) as font: - return ImageFont.truetype(font=font, size=40) + return ImageFont.truetype(font=font, size=font_size) diff --git a/src/eynollah/utils/rotate.py b/src/eynollah/utils/rotate.py index 6651c4e..e45a438 100644 --- a/src/eynollah/utils/rotate.py +++ b/src/eynollah/utils/rotate.py @@ -2,10 +2,6 @@ import math import cv2 -def rotation_image_new(img, thetha): - rotated = rotate_image(img, thetha) - return rotate_max_area_new(img, rotated, thetha) - def rotate_image(img_patch, slope): (h, w) = img_patch.shape[:2] center = (w // 2, h // 2) diff --git a/src/eynollah/utils/utils_ocr.py b/src/eynollah/utils/utils_ocr.py index 93d1137..14cbb79 100644 --- a/src/eynollah/utils/utils_ocr.py +++ b/src/eynollah/utils/utils_ocr.py @@ -3,15 +3,20 @@ import copy import numpy as np import cv2 -import tensorflow as tf +# avoid module-level import: +# import tensorflow as tf +# (wait for tf-keras and logging setup in ModelZoo.load_model) from scipy.signal import find_peaks from scipy.ndimage import gaussian_filter1d from PIL import Image, ImageDraw, ImageFont +from . import pairwise from .resize import resize_image def decode_batch_predictions(pred, num_to_char, max_len = 128): + import tensorflow as tf + # input_len is the product of the batch size and the # number of time steps. input_len = np.ones(pred.shape[0]) * pred.shape[1] @@ -36,43 +41,6 @@ def decode_batch_predictions(pred, num_to_char, max_len = 128): d = d.numpy().decode("utf-8") output.append(d) return output - - -def distortion_free_resize(image, img_size): - w, h = img_size - image = tf.image.resize(image, size=(h, w), preserve_aspect_ratio=True) - - # Check tha amount of padding needed to be done. - pad_height = h - tf.shape(image)[0] - pad_width = w - tf.shape(image)[1] - - # Only necessary if you want to do same amount of padding on both sides. - if pad_height % 2 != 0: - height = pad_height // 2 - pad_height_top = height + 1 - pad_height_bottom = height - else: - pad_height_top = pad_height_bottom = pad_height // 2 - - if pad_width % 2 != 0: - width = pad_width // 2 - pad_width_left = width + 1 - pad_width_right = width - else: - pad_width_left = pad_width_right = pad_width // 2 - - image = tf.pad( - image, - paddings=[ - [pad_height_top, pad_height_bottom], - [pad_width_left, pad_width_right], - [0, 0], - ], - ) - - image = tf.transpose(image, (1, 0, 2)) - image = tf.image.flip_left_right(image) - return image def return_start_and_end_of_common_text_of_textline_ocr_without_common_section(textline_image): width = np.shape(textline_image)[1] @@ -214,10 +182,11 @@ def get_orientation_moments(contour): def get_orientation_moments_of_mask(mask): mask=mask.astype('uint8') - contours, _ = cv2.findContours(mask[:,:,0], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - - largest_contour = max(contours, key=cv2.contourArea) if contours else None - + contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if not len(contours): + return 0 + largest_contour = max(contours, key=cv2.contourArea) moments = cv2.moments(largest_contour) if moments["mu20"] - moments["mu02"] == 0: # Avoid division by zero return 90 if moments["mu11"] > 0 else -90 @@ -256,249 +225,58 @@ def return_splitting_point_of_image(image_to_spliited): return np.sort(peaks_sort_4) -def break_curved_line_into_small_pieces_and_then_merge(img_curved, mask_curved, img_bin_curved=None): - peaks_4 = return_splitting_point_of_image(img_curved) - if len(peaks_4)>0: +def break_curved_line_into_small_pieces_and_then_merge(img_rgb_curved, img_bin_curved, mask_curved): + peaks_4 = return_splitting_point_of_image(img_rgb_curved) + if len(peaks_4): imgs_tot = [] - - for ind in range(len(peaks_4)+1): - if ind==0: - img = img_curved[:, :peaks_4[ind], :] - if img_bin_curved is not None: - img_bin = img_bin_curved[:, :peaks_4[ind], :] - mask = mask_curved[:, :peaks_4[ind], :] - elif ind==len(peaks_4): - img = img_curved[:, peaks_4[ind-1]:, :] - if img_bin_curved is not None: - img_bin = img_bin_curved[:, peaks_4[ind-1]:, :] - mask = mask_curved[:, peaks_4[ind-1]:, :] - else: - img = img_curved[:, peaks_4[ind-1]:peaks_4[ind], :] - if img_bin_curved is not None: - img_bin = img_bin_curved[:, peaks_4[ind-1]:peaks_4[ind], :] - mask = mask_curved[:, peaks_4[ind-1]:peaks_4[ind], :] - + for left, right in pairwise([None] + list(peaks_4) + [None]): + img_rgb = img_rgb_curved[:, left: right] + img_bin = img_bin_curved[:, left: right] + mask = mask_curved[:, left: right] or_ma = get_orientation_moments_of_mask(mask) - - if img_bin_curved is not None: - imgs_tot.append([img, mask, or_ma, img_bin] ) - else: - imgs_tot.append([img, mask, or_ma] ) - + imgs_tot.append([img_rgb, img_bin, mask, or_ma]) w_tot_des_list = [] - w_tot_des = 0 - imgs_deskewed_list = [] + imgs_rgb_deskewed_list = [] imgs_bin_deskewed_list = [] - for ind in range(len(imgs_tot)): - img_in = imgs_tot[ind][0] - mask_in = imgs_tot[ind][1] - ori_in = imgs_tot[ind][2] - if img_bin_curved is not None: - img_bin_in = imgs_tot[ind][3] - - if abs(ori_in)<45: - img_in_des = rotate_image_with_padding(img_in, ori_in, border_value=(255,255,255) ) - if img_bin_curved is not None: - img_bin_in_des = rotate_image_with_padding(img_bin_in, ori_in, border_value=(255,255,255) ) + for img_rgb_in, img_bin_in, mask_in, ori_in in imgs_tot: + if abs(ori_in) < 45: + img_rgb_in_des = rotate_image_with_padding(img_rgb_in, ori_in, border_value=(255,255,255) ) + img_bin_in_des = rotate_image_with_padding(img_bin_in, ori_in, border_value=(255,255,255) ) mask_in_des = rotate_image_with_padding(mask_in, ori_in) - mask_in_des = mask_in_des.astype('uint8') - - #new bounding box - x_n, y_n, w_n, h_n = get_contours_and_bounding_boxes(mask_in_des[:,:,0]) - - if w_n==0 or h_n==0: - img_in_des = np.copy(img_in) - if img_bin_curved is not None: - img_bin_in_des = np.copy(img_bin_in) - w_relative = int(32 * img_in_des.shape[1]/float(img_in_des.shape[0]) ) - if w_relative==0: - w_relative = img_in_des.shape[1] - img_in_des = resize_image(img_in_des, 32, w_relative) - if img_bin_curved is not None: - img_bin_in_des = resize_image(img_bin_in_des, 32, w_relative) + # get new bounding box + x_n, y_n, w_n, h_n = get_contours_and_bounding_boxes(mask_in_des) + if w_n and h_n: + img_rgb_in_des = img_rgb_in_des[y_n: y_n + h_n, x_n: x_n + w_n] + img_bin_in_des = img_bin_in_des[y_n: y_n + h_n, x_n: x_n + w_n] else: - mask_in_des = mask_in_des[y_n:y_n+h_n, x_n:x_n+w_n, :] - img_in_des = img_in_des[y_n:y_n+h_n, x_n:x_n+w_n, :] - if img_bin_curved is not None: - img_bin_in_des = img_bin_in_des[y_n:y_n+h_n, x_n:x_n+w_n, :] - - w_relative = int(32 * img_in_des.shape[1]/float(img_in_des.shape[0]) ) - if w_relative==0: - w_relative = img_in_des.shape[1] - img_in_des = resize_image(img_in_des, 32, w_relative) - if img_bin_curved is not None: - img_bin_in_des = resize_image(img_bin_in_des, 32, w_relative) - - - else: - img_in_des = np.copy(img_in) - if img_bin_curved is not None: + img_rgb_in_des = np.copy(img_rgb_in) img_bin_in_des = np.copy(img_bin_in) - w_relative = int(32 * img_in_des.shape[1]/float(img_in_des.shape[0]) ) - if w_relative==0: - w_relative = img_in_des.shape[1] - img_in_des = resize_image(img_in_des, 32, w_relative) - if img_bin_curved is not None: - img_bin_in_des = resize_image(img_bin_in_des, 32, w_relative) - - w_tot_des+=img_in_des.shape[1] - w_tot_des_list.append(img_in_des.shape[1]) - imgs_deskewed_list.append(img_in_des) - if img_bin_curved is not None: - imgs_bin_deskewed_list.append(img_bin_in_des) - - - + else: + img_rgb_in_des = np.copy(img_rgb_in) + img_bin_in_des = np.copy(img_bin_in) - img_final_deskewed = np.zeros((32, w_tot_des, 3))+255 - if img_bin_curved is not None: - img_bin_final_deskewed = np.zeros((32, w_tot_des, 3))+255 - else: - img_bin_final_deskewed = None + h, w = img_rgb_in_des.shape[:2] + new_h = 32 + new_w = 32 * w // h + new_w = new_w or w + img_rgb_in_des = resize_image(img_rgb_in_des, new_h, new_w) + img_bin_in_des = resize_image(img_bin_in_des, new_h, new_w) + + w_tot_des_list.append(new_w) + imgs_rgb_deskewed_list.append(img_rgb_in_des) + imgs_bin_deskewed_list.append(img_bin_in_des) + + img_rgb_final_deskewed = np.ones((new_h, sum(w_tot_des_list), 3)) * 255 + img_bin_final_deskewed = np.ones((new_h, sum(w_tot_des_list), 3)) * 255 w_indexer = 0 for ind in range(len(w_tot_des_list)): - img_final_deskewed[:,w_indexer:w_indexer+w_tot_des_list[ind],:] = imgs_deskewed_list[ind][:,:,:] - if img_bin_curved is not None: - img_bin_final_deskewed[:,w_indexer:w_indexer+w_tot_des_list[ind],:] = imgs_bin_deskewed_list[ind][:,:,:] - w_indexer = w_indexer+w_tot_des_list[ind] - return img_final_deskewed, img_bin_final_deskewed + w_indexer2 = w_indexer + w_tot_des_list[ind] + img_rgb_final_deskewed[:, w_indexer: w_indexer2] = imgs_rgb_deskewed_list[ind] + img_bin_final_deskewed[:, w_indexer: w_indexer2] = imgs_bin_deskewed_list[ind] + w_indexer = w_indexer2 + return img_rgb_final_deskewed, img_bin_final_deskewed else: - return img_curved, img_bin_curved - -def return_textline_contour_with_added_box_coordinate(textline_contour, box_ind): - textline_contour[:,:,0] += box_ind[2] - textline_contour[:,:,1] += box_ind[0] - return textline_contour - - -def return_rnn_cnn_ocr_of_given_textlines(image, - all_found_textline_polygons, - all_box_coord, - prediction_model, - b_s_ocr, num_to_char, - curved_line=False): - max_len = 512 - padding_token = 299 - image_width = 512#max_len * 4 - image_height = 32 - ind_tot = 0 - #cv2.imwrite('./img_out.png', image_page) - ocr_all_textlines = [] - cropped_lines_region_indexer = [] - cropped_lines_meging_indexing = [] - cropped_lines = [] - indexer_text_region = 0 - - for indexing, ind_poly_first in enumerate(all_found_textline_polygons): - #ocr_textline_in_textregion = [] - if len(ind_poly_first)==0: - cropped_lines_region_indexer.append(indexer_text_region) - cropped_lines_meging_indexing.append(0) - img_fin = np.ones((image_height, image_width, 3))*1 - cropped_lines.append(img_fin) - - else: - for indexing2, ind_poly in enumerate(ind_poly_first): - cropped_lines_region_indexer.append(indexer_text_region) - if not curved_line: - ind_poly = copy.deepcopy(ind_poly) - box_ind = all_box_coord[indexing] - - ind_poly = return_textline_contour_with_added_box_coordinate(ind_poly, box_ind) - #print(ind_poly_copy) - ind_poly[ind_poly<0] = 0 - x, y, w, h = cv2.boundingRect(ind_poly) - - w_scaled = w * image_height/float(h) - - mask_poly = np.zeros(image.shape) - - img_poly_on_img = np.copy(image) - - mask_poly = cv2.fillPoly(mask_poly, pts=[ind_poly], color=(1, 1, 1)) - - - - mask_poly = mask_poly[y:y+h, x:x+w, :] - img_crop = img_poly_on_img[y:y+h, x:x+w, :] - - img_crop[mask_poly==0] = 255 - - if w_scaled < 640:#1.5*image_width: - img_fin = preprocess_and_resize_image_for_ocrcnn_model(img_crop, image_height, image_width) - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(0) - else: - splited_images, splited_images_bin = return_textlines_split_if_needed(img_crop, None) - - if splited_images: - img_fin = preprocess_and_resize_image_for_ocrcnn_model(splited_images[0], - image_height, - image_width) - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(1) - - img_fin = preprocess_and_resize_image_for_ocrcnn_model(splited_images[1], - image_height, - image_width) - - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(-1) - - else: - img_fin = preprocess_and_resize_image_for_ocrcnn_model(img_crop, - image_height, - image_width) - cropped_lines.append(img_fin) - cropped_lines_meging_indexing.append(0) - - indexer_text_region+=1 - - extracted_texts = [] - - n_iterations = math.ceil(len(cropped_lines) / b_s_ocr) - - for i in range(n_iterations): - if i==(n_iterations-1): - n_start = i*b_s_ocr - imgs = cropped_lines[n_start:] - imgs = np.array(imgs) - imgs = imgs.reshape(imgs.shape[0], image_height, image_width, 3) - - - else: - n_start = i*b_s_ocr - n_end = (i+1)*b_s_ocr - imgs = cropped_lines[n_start:n_end] - imgs = np.array(imgs).reshape(b_s_ocr, image_height, image_width, 3) - - - preds = prediction_model.predict(imgs, verbose=0) - - pred_texts = decode_batch_predictions(preds, num_to_char) - - for ib in range(imgs.shape[0]): - pred_texts_ib = pred_texts[ib].replace("[UNK]", "") - extracted_texts.append(pred_texts_ib) - - extracted_texts_merged = [extracted_texts[ind] - if cropped_lines_meging_indexing[ind]==0 - else extracted_texts[ind]+" "+extracted_texts[ind+1] - if cropped_lines_meging_indexing[ind]==1 - else None - for ind in range(len(cropped_lines_meging_indexing))] - - extracted_texts_merged = [ind for ind in extracted_texts_merged if ind is not None] - unique_cropped_lines_region_indexer = np.unique(cropped_lines_region_indexer) - - ocr_all_textlines = [] - for ind in unique_cropped_lines_region_indexer: - ocr_textline_in_textregion = [] - extracted_texts_merged_un = np.array(extracted_texts_merged)[np.array(cropped_lines_region_indexer)==ind] - for it_ind, text_textline in enumerate(extracted_texts_merged_un): - ocr_textline_in_textregion.append(text_textline) - ocr_all_textlines.append(ocr_textline_in_textregion) - return ocr_all_textlines + return img_rgb_curved, img_bin_curved diff --git a/tests/cli_tests/conftest.py b/tests/cli_tests/conftest.py index 601d76b..2e1501c 100644 --- a/tests/cli_tests/conftest.py +++ b/tests/cli_tests/conftest.py @@ -1,4 +1,5 @@ from typing import List +import os import pytest import logging @@ -31,6 +32,8 @@ def run_eynollah_ok_and_check_logs( subcommand, *args ] + if 'EYNOLLAH_OPTIONS' in os.environ: + args = os.environ['EYNOLLAH_OPTIONS'].split() + args if pytestconfig.getoption('verbose') > 0: args = ['-l', 'DEBUG'] + args caplog.set_level(logging.INFO) diff --git a/tests/cli_tests/test_layout.py b/tests/cli_tests/test_layout.py index 7cbe013..503aeac 100644 --- a/tests/cli_tests/test_layout.py +++ b/tests/cli_tests/test_layout.py @@ -6,11 +6,12 @@ from ocrd_models.constants import NAMESPACES as NS "options", [ [], # defaults - #["--allow_scaling", "--curved-line"], - ["--allow_scaling", "--curved-line", "--full-layout"], - ["--allow_scaling", "--curved-line", "--full-layout", "--reading_order_machine_based"], + #["--curved-line"], + ["--curved-line", "--full-layout"], + ["--curved-line", "--full-layout", "--reading_order_machine_based"], # -ep ... - # -eoi ... + # --input_binary + # --ignore_page_extraction # --skip_layout_and_reading_order ], ids=str) def test_run_eynollah_layout_filename( diff --git a/tests/cli_tests/test_ocr.py b/tests/cli_tests/test_ocr.py index 6bf3080..c4fba04 100644 --- a/tests/cli_tests/test_ocr.py +++ b/tests/cli_tests/test_ocr.py @@ -30,7 +30,7 @@ def test_run_eynollah_ocr_filename( '-o', str(outfile.parent), ] + options, [ - # FIXME: ocr has no logging! + str(infile) ] ) assert outfile.exists() @@ -57,7 +57,8 @@ def test_run_eynollah_ocr_directory( '-o', str(outdir), ], [ - # FIXME: ocr has no logging! + 'Job done in', + 'All jobs done in', ] ) assert len(list(outdir.iterdir())) == 2 diff --git a/tests/test_model_zoo.py b/tests/test_model_zoo.py index 9d37431..082f2a7 100644 --- a/tests/test_model_zoo.py +++ b/tests/test_model_zoo.py @@ -1,16 +1,29 @@ from eynollah.model_zoo import EynollahModelZoo +from eynollah.predictor import Predictor def test_trocr1( model_dir, ): model_zoo = EynollahModelZoo(model_dir) try: - from transformers import TrOCRProcessor, VisionEncoderDecoderModel - model_zoo.load_models('trocr_processor') - proc = model_zoo.get('trocr_processor') - assert isinstance(proc, TrOCRProcessor) - model_zoo.load_models(['ocr', 'tr']) + model_zoo.load_models(('ocr', 'tr')) model = model_zoo.get('ocr') - assert isinstance(model, VisionEncoderDecoderModel) + assert isinstance(model, Predictor) + shape = model.input_shape + assert len(shape) == 3 + except ImportError: + pass + +def test_cnnrnnocr1( + model_dir, +): + model_zoo = EynollahModelZoo(model_dir) + try: + model_zoo.load_models('ocr') + model = model_zoo.get('ocr') + assert isinstance(model, Predictor) + shape = model.input_shape + assert len(shape) == 2 + assert len(shape[0]) == 4 except ImportError: pass diff --git a/train/config_params_trocr.json b/train/config_params_trocr.json new file mode 100644 index 0000000..34c6376 --- /dev/null +++ b/train/config_params_trocr.json @@ -0,0 +1,82 @@ +{ + "backbone_type" : "transformer", + "task": "transformer-ocr", + "n_classes" : 2, + "max_len": 192, + "n_epochs" : 1, + "input_height" : 32, + "input_width" : 512, + "weight_decay" : 1e-6, + "n_batch" : 1, + "learning_rate": 1e-5, + "save_interval": 1500, + "patches" : false, + "pretraining" : false, + "augmentation" : true, + "flip_aug" : false, + "blur_aug" : true, + "scaling" : false, + "adding_rgb_background": true, + "adding_rgb_foreground": true, + "add_red_textlines": true, + "white_noise_strap": true, + "textline_right_in_depth": true, + "textline_left_in_depth": true, + "textline_up_in_depth": true, + "textline_down_in_depth": true, + "textline_right_in_depth_bin": true, + "textline_left_in_depth_bin": true, + "textline_up_in_depth_bin": true, + "textline_down_in_depth_bin": true, + "bin_deg": true, + "textline_skewing": true, + "textline_skewing_bin": true, + "channels_shuffling": true, + "degrading": true, + "brightening": true, + "binarization" : true, + "pepper_aug": true, + "pepper_bin_aug": true, + "image_inversion": true, + "scaling_bluring" : false, + "scaling_binarization" : false, + "scaling_flip" : false, + "rotation": false, + "color_padding_rotation": true, + "padding_white": true, + "rotation_not_90": true, + "transformer_num_patches_xy": [56, 56], + "transformer_patchsize_x": 4, + "transformer_patchsize_y": 4, + "transformer_projection_dim": 64, + "transformer_mlp_head_units": [128, 64], + "transformer_layers": 1, + "transformer_num_heads": 1, + "transformer_cnn_first": false, + "blur_k" : ["blur","gauss","median"], + "padd_colors" : ["white", "black"], + "scales" : [0.6, 0.7, 0.8, 0.9], + "brightness" : [1.3, 1.5, 1.7, 2], + "degrade_scales" : [0.2, 0.4], + "pepper_indexes": [0.01, 0.005], + "skewing_amplitudes" : [5, 8], + "flip_index" : [0, 1, -1], + "shuffle_indexes" : [ [0,2,1], [1,2,0], [1,0,2] , [2,1,0]], + "thetha" : [0.1, 0.2, -0.1, -0.2], + "thetha_padd": [-0.6, -1, -1.4, -1.8, 0.6, 1, 1.4, 1.8], + "white_padds" : [0.1, 0.3, 0.5, 0.7, 0.9], + "number_of_backgrounds_per_image": 2, + "continue_training": false, + "index_start" : 0, + "dir_of_start_model" : " ", + "weighted_loss": false, + "is_loss_soft_dice": false, + "data_is_provided": false, + "dir_train": "/home/vahid/extracted_lines/1919_bin/train", + "dir_eval": "/home/vahid/Documents/test/sbb_pixelwise_segmentation/test_label/pageextractor_test/eval_new", + "dir_output": "/home/vahid/extracted_lines/1919_bin/output", + "dir_rgb_backgrounds": "/home/vahid/Documents/1_2_test_eynollah/set_rgb_background", + "dir_rgb_foregrounds": "/home/vahid/Documents/1_2_test_eynollah/out_set_rgb_foreground", + "dir_img_bin": "/home/vahid/extracted_lines/1919_bin/images_bin" + +} diff --git a/train/requirements.txt b/train/requirements.txt index 090bc50..cd4d74b 100644 --- a/train/requirements.txt +++ b/train/requirements.txt @@ -6,5 +6,10 @@ imutils scipy tensorflow-addons # for connected_components, depublished and only compatible with tensorflow < 2.16 tensorflow < 2.16 # for tensorflow-addons, so only needed in training -tf_data < 2.16 # for tensorflow-addons, so only needed in training +tf-keras < 2.16 # avoid keras 3 (also needs TF_USE_LEGACY_KERAS=1) protobuf < 5 # for tensorflow-addons, so only needed in training +torch +transformers <= 4.30.2 ; python_version < '3.10' +transformers >= 5 ; python_version >= '3.10' +eynollah-fork-tf2onnx == 1.17.0.post2 +ml_dtypes >= 0.5