diff --git a/ocrd-tool.json b/ocrd-tool.json new file mode 120000 index 0000000..a1e5650 --- /dev/null +++ b/ocrd-tool.json @@ -0,0 +1 @@ +qurator/sbb_textline_detector/ocrd-tool.json \ No newline at end of file diff --git a/qurator/sbb_textline_detector/__init__.py b/qurator/sbb_textline_detector/__init__.py index 15b6a64..b7c0712 100644 --- a/qurator/sbb_textline_detector/__init__.py +++ b/qurator/sbb_textline_detector/__init__.py @@ -1 +1,2 @@ from .main import * +from .ocrd_cli import * diff --git a/qurator/sbb_textline_detector/main.py b/qurator/sbb_textline_detector/main.py index c7d47e2..d775aa5 100644 --- a/qurator/sbb_textline_detector/main.py +++ b/qurator/sbb_textline_detector/main.py @@ -35,16 +35,18 @@ __doc__ = \ class textlineerkenner: - def __init__(self, image_dir, dir_out, dir_models): + def __init__(self, image_dir, dir_out, f_name, dir_models): self.image_dir = image_dir self.dir_out = dir_out + self.f_name = f_name + if self.f_name is None: + try: + self.f_name = image_dir.split('/')[len(image_dir.split('/')) - 1] + self.f_name = self.f_name.split('.')[0] + print(self.f_name) + except: + self.f_name = self.f_name.split('.')[0] self.dir_models = dir_models - try: - self.f_name = image_dir.split('/')[len(image_dir.split('/')) - 1] - self.f_name = self.f_name.split('.')[0] - print(self.f_name) - except: - self.f_name = self.f_name.split('.')[0] self.kernel = np.ones((5, 5), np.uint8) self.model_page_dir = dir_models + '/model_page.h5' self.model_region_dir = dir_models + '/model_strukturerkennung.h5' @@ -365,6 +367,7 @@ class textlineerkenner: return img_r def get_image_and_scales(self): + print(self.image_dir) self.image = cv2.imread(self.image_dir) self.height_org = self.image.shape[0] self.width_org = self.image.shape[1] @@ -1298,7 +1301,7 @@ class textlineerkenner: unireg.text = ' ' tree = ET.ElementTree(data) - tree.write(dir_of_image + self.f_name + ".xml") + tree.write(os.path.join(dir_of_image, self.f_name) + ".xml") def run(self): self.get_image_and_scales() @@ -1307,7 +1310,7 @@ class textlineerkenner: boxes,contours=self.get_text_region_contours_and_boxes(text_regions) self.get_all_image_patches_based_on_text_regions(boxes,image_page) textline_mask_tot=self.textline_contours(image_page) - + self.get_textlines_for_each_textregions(textline_mask_tot,boxes) self.get_slopes_for_each_text_region(contours) self.deskew_textline_patches(contours, boxes) @@ -1321,7 +1324,7 @@ class textlineerkenner: def main(image, out, model): possibles = globals() # XXX unused? possibles.update(locals()) - x = textlineerkenner(image, out, model) + x = textlineerkenner(image, out, None, model) x.run() diff --git a/qurator/sbb_textline_detector/ocrd-tool.json b/qurator/sbb_textline_detector/ocrd-tool.json new file mode 100644 index 0000000..b76f439 --- /dev/null +++ b/qurator/sbb_textline_detector/ocrd-tool.json @@ -0,0 +1,19 @@ +{ + "version": "0.0.1", + "tools": { + "ocrd_sbb_textline_detector": { + "executable": "ocrd_sbb_textline_detector", + "description": "Detect lines", + "steps": ["layout/segmentation/line"], + "input_file_grp": [ + "OCR-D-IMG" + ], + "output_file_grp": [ + "OCR-D-SBB-SEG-LINE" + ], + "parameters": { + "model": {"type": "string", "format": "file", "cacheable": true} + } + } + } +} diff --git a/qurator/sbb_textline_detector/ocrd_cli.py b/qurator/sbb_textline_detector/ocrd_cli.py new file mode 100644 index 0000000..854a5dc --- /dev/null +++ b/qurator/sbb_textline_detector/ocrd_cli.py @@ -0,0 +1,61 @@ +import json +import os + +import click +from ocrd import Processor +from ocrd.decorators import ocrd_cli_options, ocrd_cli_wrap_processor +from ocrd_utils import concat_padded, getLogger +from pkg_resources import resource_string + +from qurator.sbb_textline_detector import textlineerkenner + +log = getLogger('processor.OcrdSbbTextlineDetectorRecognize') + +OCRD_TOOL = json.loads(resource_string(__name__, 'ocrd-tool.json').decode('utf8')) + + +@click.command() +@ocrd_cli_options +def ocrd_sbb_textline_detector(*args, **kwargs): + return ocrd_cli_wrap_processor(OcrdSbbTextlineDetectorRecognize, *args, **kwargs) + + +class OcrdSbbTextlineDetectorRecognize(Processor): + + def __init__(self, *args, **kwargs): + kwargs['ocrd_tool'] = OCRD_TOOL['tools']['ocrd_sbb_textline_detector'] + super(OcrdSbbTextlineDetectorRecognize, self).__init__(*args, **kwargs) + + def _make_file_id(self, input_file, input_file_grp, n): + file_id = input_file.ID.replace(input_file_grp, self.output_file_grp) + if file_id == input_file.ID: + file_id = concat_padded(self.output_file_grp, n) + return file_id + + def process(self): + for n, page_id in enumerate(self.workspace.mets.physical_pages): + image_file = self.workspace.mets.find_files(fileGrp=self.input_file_grp, pageId=page_id)[0] + log.info("INPUT FILE %i / %s", n, image_file) + + file_id = self._make_file_id(image_file, self.output_file_grp, n) + + # Process the files + try: + os.mkdir(self.output_file_grp) + except FileExistsError: + pass + + model = self.parameter['model'] + x = textlineerkenner(image_file.local_filename, self.output_file_grp, file_id, model) + x.run() + + self.workspace.add_file( + ID=file_id + '.xml', + file_grp=self.output_file_grp, + pageId=page_id, + mimetype='application/vnd.prima.page+xml', + local_filename=self.output_file_grp + '/' + file_id) + + +if __name__ == '__main__': + ocrd_sbb_textline_detector() diff --git a/requirements.txt b/requirements.txt index 58d61f7..acc0101 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ opencv-python -numpy<=1.14.5 +numpy matplotlib seaborn tqdm @@ -9,3 +9,4 @@ scikit-learn tensorflow-gpu < 2.0 scipy click +ocrd >= 1.0.0b19 diff --git a/setup.py b/setup.py index faff412..26eff52 100644 --- a/setup.py +++ b/setup.py @@ -21,6 +21,7 @@ setup( entry_points={ 'console_scripts': [ "sbb_textline_detector=qurator.sbb_textline_detector:main", + "ocrd_sbb_textline_detector=qurator.sbb_textline_detector:ocrd_sbb_textline_detector", ] }, python_requires='>=3.6.0',