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 d544a67..82209be 100644 --- a/src/eynollah/cli/cli_binarize.py +++ b/src/eynollah/cli/cli_binarize.py @@ -33,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, @@ -46,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 42b1d41..bcb8263 100644 --- a/src/eynollah/cli/cli_enhance.py +++ b/src/eynollah/cli/cli_enhance.py @@ -48,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 """ @@ -62,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_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 99e03c5..daeccbe 100644 --- a/src/eynollah/cli/cli_ocr.py +++ b/src/eynollah/cli/cli_ocr.py @@ -73,11 +73,6 @@ import click 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.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 ocr_cli( ctx, @@ -92,7 +87,6 @@ def ocr_cli( do_not_mask_with_textline_contour, batch_size, min_conf_value_of_textline_text, - device, ): """ Recognize text with a CNN/RNN or transformer ML model. @@ -101,11 +95,12 @@ def ocr_cli( 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, - device=device) + ) eynollah_ocr.run(overwrite=overwrite, dir_in=dir_in, dir_in_bin=dir_in_bin, diff --git a/src/eynollah/cli/cli_readingorder.py b/src/eynollah/cli/cli_readingorder.py index 9bb7092..ac52e38 100644 --- a/src/eynollah/cli/cli_readingorder.py +++ b/src/eynollah/cli/cli_readingorder.py @@ -22,19 +22,15 @@ import click type=click.Path(exists=True, file_okay=False), required=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 readingorder_cli(ctx, input, dir_in, out, device): +def readingorder_cli(ctx, input, dir_in, out): """ Generate ReadingOrder with a ML model """ 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 = Reorder(model_zoo=ctx.obj.model_zoo, device=device) + orderer = Reorder(model_zoo=ctx.obj.model_zoo, + device=ctx.obj.device) orderer.run(xml_filename=input, dir_in=dir_in, dir_out=out,