From 283a4e32fede6b47d41432be1cfd2ff9d22f4eb4 Mon Sep 17 00:00:00 2001 From: Benjamin Kiessling Date: Wed, 20 Dec 2023 17:52:09 +0100 Subject: [PATCH] Linting and type hint imports --- kraken/align.py | 8 +- kraken/binarization.py | 11 ++- kraken/containers.py | 15 ++-- kraken/contrib/extract_lines.py | 1 + kraken/contrib/heatmap_overlay.py | 1 + .../hyperparameters/tune_pretraining.py | 2 + .../contrib/hyperparameters/tune_training.py | 1 - kraken/contrib/segmentation_overlay.py | 1 + kraken/ketos/transcription.py | 2 +- kraken/kraken.py | 10 +-- kraken/lib/arrow_dataset.py | 12 +-- kraken/lib/dataset/recognition.py | 8 +- kraken/lib/dataset/ro.py | 10 ++- kraken/lib/dataset/segmentation.py | 9 +- kraken/lib/functional_im_transforms.py | 8 +- kraken/lib/layers.py | 82 ++++++++++++----- kraken/lib/lineest.py | 13 +-- kraken/lib/models.py | 9 +- kraken/lib/pretrain/layers.py | 6 +- kraken/lib/pretrain/model.py | 11 +-- kraken/lib/progress.py | 26 +++--- kraken/lib/ro/layers.py | 7 +- kraken/lib/ro/model.py | 13 +-- kraken/lib/segmentation.py | 88 +++++++++---------- kraken/lib/train.py | 28 +++--- kraken/lib/util.py | 12 +-- kraken/lib/xml.py | 10 ++- kraken/repo.py | 9 +- kraken/rpred.py | 7 +- kraken/serialization.py | 24 ++--- 30 files changed, 265 insertions(+), 179 deletions(-) diff --git a/kraken/align.py b/kraken/align.py index db2948ce0..b496821ed 100644 --- a/kraken/align.py +++ b/kraken/align.py @@ -29,16 +29,18 @@ from bidi.algorithm import get_display from dataclasses import dataclass -from typing import Optional, Literal +from typing import Optional, Literal, TYPE_CHECKING from kraken import rpred from kraken.containers import Segmentation, BaselineOCRRecord -from kraken.lib.models import TorchSeqRecognizer + +if TYPE_CHECKING: + from kraken.lib.models import TorchSeqRecognizer logger = logging.getLogger('kraken') -def forced_align(doc: Segmentation, model: TorchSeqRecognizer, base_dir: Optional[Literal['L', 'R']] = None) -> Segmentation: +def forced_align(doc: Segmentation, model: 'TorchSeqRecognizer', base_dir: Optional[Literal['L', 'R']] = None) -> Segmentation: """ Performs a forced character alignment of text with recognition model output activations. diff --git a/kraken/binarization.py b/kraken/binarization.py index 703497632..db3adacc7 100644 --- a/kraken/binarization.py +++ b/kraken/binarization.py @@ -23,19 +23,24 @@ import logging import numpy as np -from PIL import Image from kraken.lib.util import pil2array, array2pil, is_bitonal, get_im_str from scipy.ndimage import affine_transform, percentile_filter, gaussian_filter, binary_dilation from scipy.ndimage import zoom as _zoom +from typing import TYPE_CHECKING + from kraken.lib.exceptions import KrakenInputException +if TYPE_CHECKING: + from PIL import Image + + __all__ = ['nlbin'] logger = logging.getLogger(__name__) -def nlbin(im: Image.Image, +def nlbin(im: 'Image.Image', threshold: float = 0.5, zoom: float = 0.5, escale: float = 1.0, @@ -43,7 +48,7 @@ def nlbin(im: Image.Image, perc: int = 80, range: int = 20, low: int = 5, - high: int = 90) -> Image.Image: + high: int = 90) -> 'Image.Image': """ Performs binarization using non-linear processing. diff --git a/kraken/containers.py b/kraken/containers.py index 54885b81d..188bb8e25 100644 --- a/kraken/containers.py +++ b/kraken/containers.py @@ -22,13 +22,16 @@ import numpy as np import bidi.algorithm as bd -from os import PathLike -from typing import Literal, List, Dict, Union, Optional, Tuple, Any +from typing import Literal, List, Dict, Union, Optional, Tuple, Any, TYPE_CHECKING from dataclasses import dataclass, asdict from abc import ABC, abstractmethod from kraken.lib.segmentation import compute_polygon_section +if TYPE_CHECKING: + from os import PathLike + + __all__ = ['BaselineLine', 'BBoxLine', 'Segmentation', @@ -85,7 +88,7 @@ class BaselineLine: text: Optional[str] = None base_dir: Optional[Literal['L', 'R']] = None type: str = 'baselines' - imagename: Optional[Union[str, PathLike]] = None + imagename: Optional[Union[str, 'PathLike']] = None tags: Optional[Dict[str, str]] = None split: Optional[Literal['train', 'validation', 'test']] = None regions: Optional[List[str]] = None @@ -124,7 +127,7 @@ class BBoxLine: text: Optional[str] = None base_dir: Optional[Literal['L', 'R']] = None type: str = 'bbox' - imagename: Optional[Union[str, PathLike]] = None + imagename: Optional[Union[str, 'PathLike']] = None tags: Optional[Dict[str, str]] = None split: Optional[Literal['train', 'validation', 'test']] = None regions: Optional[List[str]] = None @@ -145,7 +148,7 @@ class Region: """ id: str boundary: List[Tuple[int, int]] - imagename: Optional[Union[str, PathLike]] = None + imagename: Optional[Union[str, 'PathLike']] = None tags: Optional[Dict[str, str]] = None @@ -175,7 +178,7 @@ class Segmentation: Each reading order is a list of line indices. """ type: Literal['baselines', 'bbox'] - imagename: Union[str, PathLike] + imagename: Union[str, 'PathLike'] text_direction: Literal['horizontal-lr', 'horizontal-rl', 'vertical-lr', 'vertical-rl'] script_detection: bool lines: List[Union[BaselineLine, BBoxLine]] diff --git a/kraken/contrib/extract_lines.py b/kraken/contrib/extract_lines.py index e2a41b8ba..3dec3bc51 100755 --- a/kraken/contrib/extract_lines.py +++ b/kraken/contrib/extract_lines.py @@ -1,6 +1,7 @@ #! /usr/bin/env python import click + @click.command() @click.option('-f', '--format-type', type=click.Choice(['xml', 'alto', 'page', 'binary']), default='xml', help='Sets the input document format. In ALTO and PageXML mode all ' diff --git a/kraken/contrib/heatmap_overlay.py b/kraken/contrib/heatmap_overlay.py index 0b1c22197..8e2c7236b 100755 --- a/kraken/contrib/heatmap_overlay.py +++ b/kraken/contrib/heatmap_overlay.py @@ -4,6 +4,7 @@ """ import click + @click.command() @click.option('-i', '--model', default=None, show_default=True, type=click.Path(exists=True), help='Baseline detection model to use.') diff --git a/kraken/contrib/hyperparameters/tune_pretraining.py b/kraken/contrib/hyperparameters/tune_pretraining.py index 923f8463c..1348280c1 100644 --- a/kraken/contrib/hyperparameters/tune_pretraining.py +++ b/kraken/contrib/hyperparameters/tune_pretraining.py @@ -51,6 +51,7 @@ def train_tune(config, training_data=None, epochs=100, spec=RECOGNITION_SPEC): enable_progress_bar=False) trainer.fit(model, datamodule=data_module) + @click.command() @click.option('-v', '--verbose', default=0, count=True) @click.option('-s', '--seed', default=42, type=click.INT, @@ -83,5 +84,6 @@ def cli(verbose, seed, output, num_samples, epochs, spec, training_files, files) click.echo("Best hyperparameters found were: ", analysis.get_best_config(metric='accuracy', mode='max')) + if __name__ == '__main__': cli() diff --git a/kraken/contrib/hyperparameters/tune_training.py b/kraken/contrib/hyperparameters/tune_training.py index 3bfd1b73b..98f3781b4 100644 --- a/kraken/contrib/hyperparameters/tune_training.py +++ b/kraken/contrib/hyperparameters/tune_training.py @@ -11,7 +11,6 @@ from kraken.lib.default_spec import RECOGNITION_PRETRAIN_HYPER_PARAMS, RECOGNITION_SPEC from kraken.lib.pretrain.model import PretrainDataModule, RecognitionPretrainModel -from ray.tune.schedulers import ASHAScheduler import pytorch_lightning as pl diff --git a/kraken/contrib/segmentation_overlay.py b/kraken/contrib/segmentation_overlay.py index a9ff235ff..181a92ddc 100755 --- a/kraken/contrib/segmentation_overlay.py +++ b/kraken/contrib/segmentation_overlay.py @@ -11,6 +11,7 @@ from itertools import cycle from collections import defaultdict + cmap = cycle([(230, 25, 75, 127), (60, 180, 75, 127)]) diff --git a/kraken/ketos/transcription.py b/kraken/ketos/transcription.py index dd4402c66..2246c3f5c 100644 --- a/kraken/ketos/transcription.py +++ b/kraken/ketos/transcription.py @@ -202,7 +202,7 @@ def transcription(ctx, text_direction, scale, bw, maxcolseps, else: with click.open_file(lines, 'r') as fp: try: - fp = cast(IO[Any], fp) + fp = cast('IO[Any]', fp) res = json.load(fp) except ValueError as e: raise click.UsageError('{} invalid segmentation: {}'.format(lines, str(e))) diff --git a/kraken/kraken.py b/kraken/kraken.py index b06c38ce1..ba543a995 100644 --- a/kraken/kraken.py +++ b/kraken/kraken.py @@ -86,7 +86,7 @@ def binarizer(threshold, zoom, escale, border, perc, range, low, high, input, ou low, high) if ctx.meta['last_process'] and ctx.meta['output_mode'] != 'native': with click.open_file(output, 'w', encoding='utf-8') as fp: - fp = cast(IO[Any], fp) + fp = cast('IO[Any]', fp) logger.info('Serializing as {} into {}'.format(ctx.meta['output_mode'], output)) res.save(f'{output}.png') from kraken import serialization @@ -159,7 +159,7 @@ def segmenter(legacy, model, text_direction, scale, maxcolseps, black_colseps, ctx.exit(1) if ctx.meta['last_process'] and ctx.meta['output_mode'] != 'native': with click.open_file(output, 'w', encoding='utf-8') as fp: - fp = cast(IO[Any], fp) + fp = cast('IO[Any]', fp) logger.info('Serializing as {} into {}'.format(ctx.meta['output_mode'], output)) from kraken import serialization fp.write(serialization.serialize_segmentation(res, @@ -170,7 +170,7 @@ def segmenter(legacy, model, text_direction, scale, maxcolseps, black_colseps, processing_steps=ctx.meta['steps'])) else: with click.open_file(output, 'w') as fp: - fp = cast(IO[Any], fp) + fp = cast('IO[Any]', fp) json.dump(dataclasses.asdict(res), fp) message('\u2713', fg='green') @@ -208,7 +208,7 @@ def recognizer(model, pad, no_segmentation, bidi_reordering, tags_ignore, input, if not bounds and ctx.meta['base_image'] != input: with click.open_file(input, 'r') as fp: try: - fp = cast(IO[Any], fp) + fp = cast('IO[Any]', fp) bounds = Segmentation(**json.load(fp)) except ValueError as e: raise click.UsageError(f'{input} invalid segmentation: {str(e)}') @@ -244,7 +244,7 @@ def recognizer(model, pad, no_segmentation, bidi_reordering, tags_ignore, input, ctx = click.get_current_context() with click.open_file(output, 'w', encoding='utf-8') as fp: - fp = cast(IO[Any], fp) + fp = cast('IO[Any]', fp) message(f'Writing recognition results for {ctx.meta["orig_file"]}\t', nl=False) logger.info('Serializing as {} into {}'.format(ctx.meta['output_mode'], output)) if ctx.meta['output_mode'] != 'native': diff --git a/kraken/lib/arrow_dataset.py b/kraken/lib/arrow_dataset.py index 7f7e00962..87bc9f316 100755 --- a/kraken/lib/arrow_dataset.py +++ b/kraken/lib/arrow_dataset.py @@ -26,7 +26,7 @@ from PIL import Image, UnidentifiedImageError from functools import partial from collections import Counter -from typing import Optional, List, Union, Callable, Tuple, Dict +from typing import Optional, List, Union, Callable, Tuple, Dict, TYPE_CHECKING from multiprocessing import Pool from kraken.containers import Segmentation from kraken.lib import functional_im_transforms as F_t @@ -34,7 +34,9 @@ from kraken.lib.xml import XMLPage from kraken.lib.util import is_bitonal, make_printable from kraken.lib.exceptions import KrakenInputException -from os import PathLike + +if TYPE_CHECKING: + from os import PathLike import logging @@ -89,7 +91,7 @@ def _extract_path_line(xml_record, skip_empty_lines: bool = True): return [line], im.mode -def parse_path(path: Union[str, PathLike], +def parse_path(path: Union[str, 'PathLike'], suffix: str = '.gt.txt', split=F_t.default_split, skip_empty_lines: bool = True): @@ -101,8 +103,8 @@ def parse_path(path: Union[str, PathLike], return {'image': path, 'lines': [{'text': gt}]} -def build_binary_dataset(files: Optional[List[Union[str, PathLike, Dict]]] = None, - output_file: Union[str, PathLike] = None, +def build_binary_dataset(files: Optional[List[Union[str, 'PathLike', Dict]]] = None, + output_file: Union[str, 'PathLike'] = None, format_type: str = 'xml', num_workers: int = 0, ignore_splits: bool = False, diff --git a/kraken/lib/dataset/recognition.py b/kraken/lib/dataset/recognition.py index e7d6bffac..2fda9896c 100644 --- a/kraken/lib/dataset/recognition.py +++ b/kraken/lib/dataset/recognition.py @@ -24,12 +24,11 @@ import pyarrow as pa from PIL import Image -from os import PathLike from functools import partial from torchvision import transforms from collections import Counter from torch.utils.data import Dataset -from typing import List, Tuple, Callable, Optional, Any, Union, Literal +from typing import List, Tuple, Callable, Optional, Any, Union, Literal, TYPE_CHECKING from kraken.containers import BaselineLine, BBoxLine, Segmentation from kraken.lib.util import is_bitonal @@ -39,6 +38,9 @@ from kraken.lib import functional_im_transforms as F_t +if TYPE_CHECKING: + from os import PathLike + __all__ = ['DefaultAugmenter', 'ArrowIPCRecognitionDataset', 'PolygonGTDataset', @@ -137,7 +139,7 @@ def __init__(self, self.im_mode = self.transforms.mode - def add(self, file: Union[str, PathLike]) -> None: + def add(self, file: Union[str, 'PathLike']) -> None: """ Adds an Arrow IPC file to the dataset. diff --git a/kraken/lib/dataset/ro.py b/kraken/lib/dataset/ro.py index 46325d35b..65638da8f 100644 --- a/kraken/lib/dataset/ro.py +++ b/kraken/lib/dataset/ro.py @@ -19,14 +19,16 @@ import numpy as np from math import factorial -from os import PathLike from torch.utils.data import Dataset -from typing import Dict, Sequence, Union, Literal, Optional +from typing import Dict, Sequence, Union, Literal, Optional, TYPE_CHECKING from kraken.lib.xml import XMLPage from kraken.lib.exceptions import KrakenInputException +if TYPE_CHECKING: + from os import PathLike + __all__ = ['PairWiseROSet', 'PageWiseROSet'] import logging @@ -40,7 +42,7 @@ class PairWiseROSet(Dataset): Returns random pairs of lines from the same page. """ - def __init__(self, files: Sequence[Union[PathLike, str]] = None, + def __init__(self, files: Sequence[Union['PathLike', str]] = None, mode: Optional[Literal['alto', 'page', 'xml']] = 'xml', level: Literal['regions', 'baselines'] = 'baselines', ro_id: Optional[str] = None, @@ -142,7 +144,7 @@ class PageWiseROSet(Dataset): Returns all lines from the same page. """ - def __init__(self, files: Sequence[Union[PathLike, str]] = None, + def __init__(self, files: Sequence[Union['PathLike', str]] = None, mode: Optional[Literal['alto', 'page', 'xml']] = 'xml', level: Literal['regions', 'baselines'] = 'baselines', ro_id: Optional[str] = None, diff --git a/kraken/lib/dataset/segmentation.py b/kraken/lib/dataset/segmentation.py index 01b7abb90..1de44ac6e 100644 --- a/kraken/lib/dataset/segmentation.py +++ b/kraken/lib/dataset/segmentation.py @@ -27,12 +27,13 @@ from torchvision import transforms from collections import defaultdict from torch.utils.data import Dataset -from typing import Dict, Tuple, Sequence, Callable, Any, Union, Literal, Optional +from typing import Dict, Tuple, Sequence, Callable, Any, Union, Literal, Optional, TYPE_CHECKING from skimage.draw import polygon -from kraken.containers import Segmentation -from kraken.lib.xml import XMLPage +if TYPE_CHECKING: + from kraken.containers import Segmentation + from kraken.lib.xml import XMLPage __all__ = ['BaselineSet'] @@ -124,7 +125,7 @@ def __init__(self, self.transforms = im_transforms self.seg_type = None - def add(self, doc: Union[Segmentation, XMLPage]): + def add(self, doc: Union['Segmentation', 'XMLPage']): """ Adds a page to the dataset. diff --git a/kraken/lib/functional_im_transforms.py b/kraken/lib/functional_im_transforms.py index 2823b4e5e..804428cc2 100644 --- a/kraken/lib/functional_im_transforms.py +++ b/kraken/lib/functional_im_transforms.py @@ -21,16 +21,18 @@ import unicodedata import bidi.algorithm as bd -from os import PathLike from pathlib import Path -from PIL import Image from PIL.Image import Resampling -from typing import Tuple, Optional, Callable, Any, Union +from typing import Tuple, Optional, Callable, Any, Union, TYPE_CHECKING from kraken.binarization import nlbin from kraken.lib.lineest import dewarp, CenterNormalizer +if TYPE_CHECKING: + from os import PathLike + from PIL import Image + def pil_to_mode(im: Image.Image, mode: str) -> Image.Image: return im.convert(mode) diff --git a/kraken/lib/layers.py b/kraken/lib/layers.py index 0c4fee1ef..22ffdd6af 100644 --- a/kraken/lib/layers.py +++ b/kraken/lib/layers.py @@ -187,7 +187,10 @@ def __init__(self, dim: int, chunk_size: int) -> None: self.chunk_size = chunk_size super().__init__() - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[List[int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[List[int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: out = inputs.unfold(self.dim, self.chunk_size, self.chunk_size) out = out.sum(self.dim, keepdim=True) out = out.transpose(-1, self.dim).squeeze(-1) @@ -237,7 +240,10 @@ def __init__(self) -> None: """ super().__init__() - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: return inputs, seq_len def get_shape(self, input: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]: @@ -289,7 +295,10 @@ def __init__(self, src_dim: int, part_a: int, part_b: int, high: int, low: int) self.high = high self.low = low - def forward(self, input: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + input: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: initial_len = input.shape[3] # split dimension src_dim into part_a x part_b input = input.reshape(input.shape[:self.src_dim] + (self.part_a, self.part_b) + input.shape[self.src_dim + 1:]) @@ -354,7 +363,10 @@ def __init__(self, kernel_size: Tuple[int, int], stride: Tuple[int, int]) -> Non self.stride = stride self.layer = torch.nn.MaxPool2d(kernel_size, stride) - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: o = self.layer(inputs) if seq_len is not None: seq_len = torch.floor((seq_len-(self.kernel_size[1]-1)-1).float()/self.stride[1]+1).int() @@ -403,7 +415,10 @@ def __init__(self, p: float, dim: int) -> None: elif dim == 2: self.layer = torch.nn.Dropout2d(p) - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: return self.layer(inputs), seq_len def get_shape(self, input: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]: @@ -479,7 +494,10 @@ def __init__(self, batch_first=True, bias=False if legacy else True) - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: # NCHW -> HNWC inputs = inputs.permute(2, 0, 3, 1) if self.transpose: @@ -671,7 +689,10 @@ def __init__(self, input_size: int, output_size: int, augmentation: bool = False self.lin = torch.nn.Linear(self.input_size, output_size) - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: # move features (C) to last dimension for linear activation # NCHW -> NWHC inputs = inputs.transpose(1, 3) @@ -752,7 +773,14 @@ class ActConv2D(Module): dropped columns. """ - def __init__(self, in_channels: int, out_channels: int, kernel_size: Tuple[int, int], stride: Tuple[int, int], nl: str = 'l', dilation: Tuple[int, int] = (1, 1), transposed: bool = False) -> None: + def __init__(self, + in_channels: int, + out_channels: int, + kernel_size: Tuple[int, int], + stride: Tuple[int, int], + nl: str = 'l', + dilation: Tuple[int, int] = (1, 1), + transposed: bool = False) -> None: super().__init__() self.in_channels = in_channels self.kernel_size = kernel_size @@ -782,13 +810,24 @@ def __init__(self, in_channels: int, out_channels: int, kernel_size: Tuple[int, self.nl = lambda x: x if self.transposed: - self.co = torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, - stride=stride, padding=self.padding, dilation=self.dilation) + self.co = torch.nn.ConvTranspose2d(in_channels, + out_channels, + kernel_size, + stride=stride, + padding=self.padding, + dilation=self.dilation) else: - self.co = torch.nn.Conv2d(in_channels, out_channels, kernel_size, - stride=stride, padding=self.padding, dilation=self.dilation) - - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor] = None, output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + self.co = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size, + stride=stride, + padding=self.padding, + dilation=self.dilation) + + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor] = None, + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if self.transposed: o = self.co(inputs, output_size=output_shape) else: @@ -822,10 +861,10 @@ def get_shape(self, input: Tuple[int, int, int, int], target_shape: Optional[Tup min(min_x + self.stride[1] - 1, max(target_x, min_x))) else: self.output_shape = (input[0], - self.out_channels, - int(max(np.floor((input[2]+2*self.padding[0]-self.dilation[0]*(self.kernel_size[0]-1)-1) / - self.stride[0]+1), 1) if input[2] != 0 else 0), - int(max(np.floor((input[3]+2*self.padding[1]-self.dilation[1]*(self.kernel_size[1]-1)-1)/self.stride[1]+1), 1) if input[3] != 0 else 0)) + self.out_channels, + int(max(np.floor((input[2]+2*self.padding[0]-self.dilation[0]*(self.kernel_size[0]-1)-1) / + self.stride[0]+1), 1) if input[2] != 0 else 0), + int(max(np.floor((input[3]+2*self.padding[1]-self.dilation[1]*(self.kernel_size[1]-1)-1)/self.stride[1]+1), 1) if input[3] != 0 else 0)) return self.output_shape def deserialize(self, name: str, spec) -> None: @@ -897,7 +936,7 @@ def resize(self, output_size: int, del_indices: Optional[Iterable[int]] = None) bias = torch.cat([bias, torch.zeros(output_size - bias.size(0))]) if self.transposed: self.co = torch.nn.ConvTranspose2d(self.in_channels, self.out_channels, self.kernel_size, - stride=self.stride, padding=self.padding) + stride=self.stride, padding=self.padding) else: self.co = torch.nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size, stride=self.stride, padding=self.padding) @@ -917,7 +956,10 @@ def __init__(self, in_channels: int, num_groups: int) -> None: self.layer = torch.nn.GroupNorm(num_groups, in_channels) - def forward(self, inputs: torch.Tensor, seq_len: Optional[torch.Tensor], output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, + inputs: torch.Tensor, + seq_len: Optional[torch.Tensor], + output_shape: Optional[Tuple[int, int]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: t = inputs.dtype # XXX: verify that pytorch AMP casts the inputs to float32 correctly at # some point. diff --git a/kraken/lib/lineest.py b/kraken/lib/lineest.py index 2fc87c832..d7c5f3024 100644 --- a/kraken/lib/lineest.py +++ b/kraken/lib/lineest.py @@ -1,10 +1,14 @@ import warnings -from PIL import Image import numpy as np from kraken.lib.util import pil2array, array2pil from scipy.ndimage import affine_transform, gaussian_filter, uniform_filter +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from PIL import Image + __all__ = ['CenterNormalizer', 'dewarp'] @@ -63,15 +67,14 @@ def normalize(self, img, order=1, dtype=np.dtype('f'), cval=0): return scaled -def dewarp(normalizer: CenterNormalizer, im: Image.Image) -> Image.Image: +def dewarp(normalizer: CenterNormalizer, im: 'Image.Image') -> 'Image.Image': """ Dewarps an image of a line using a kraken.lib.lineest.CenterNormalizer instance. Args: - normalizer (kraken.lib.lineest.CenterNormalizer): A line normalizer - instance - im (PIL.Image.Image): Image to dewarp + normalizer: A line normalizer instance + im: Image to dewarp Returns: PIL.Image containing the dewarped image. diff --git a/kraken/lib/models.py b/kraken/lib/models.py index ac1219b23..ed0f9c173 100644 --- a/kraken/lib/models.py +++ b/kraken/lib/models.py @@ -5,7 +5,6 @@ Wrapper around TorchVGSLModel including a variety of forward pass helpers for sequence classification. """ -from os import PathLike from os.path import expandvars, expanduser, abspath import torch @@ -13,11 +12,15 @@ import kraken.lib.lineest import kraken.lib.ctc_decoder -from typing import List, Tuple, Optional, Union +from typing import List, Tuple, Optional, Union, TYPE_CHECKING from kraken.lib.vgsl import TorchVGSLModel from kraken.lib.exceptions import KrakenInvalidModelException, KrakenInputException +if TYPE_CHECKING: + from os import PathLike + + __all__ = ['TorchSeqRecognizer', 'load_any'] import logging @@ -171,7 +174,7 @@ def predict_labels(self, line: torch.tensor, lens: torch.Tensor = None) -> List[ return oseqs -def load_any(fname: Union[PathLike, str], +def load_any(fname: Union['PathLike', str], train: bool = False, device: str = 'cpu') -> TorchSeqRecognizer: """ diff --git a/kraken/lib/pretrain/layers.py b/kraken/lib/pretrain/layers.py index acbc7aa2f..009ce72fd 100644 --- a/kraken/lib/pretrain/layers.py +++ b/kraken/lib/pretrain/layers.py @@ -3,12 +3,14 @@ """ import torch -from typing import Tuple, Optional +from typing import Tuple, Optional, TYPE_CHECKING from torch.nn import Module, Embedding, Linear -from kraken.lib.vgsl import VGSLBlock from kraken.lib.pretrain.util import compute_mask_indices, sample_negatives +if TYPE_CHECKING: + from kraken.lib.vgsl import VGSLBlock + # all tensors are ordered NCHW, the "feature" dimension is C, so the output of # an LSTM will be put into C same as the filters of a CNN. diff --git a/kraken/lib/pretrain/model.py b/kraken/lib/pretrain/model.py index 63f53f199..92ab35684 100644 --- a/kraken/lib/pretrain/model.py +++ b/kraken/lib/pretrain/model.py @@ -36,10 +36,9 @@ import torch.nn.functional as F import pytorch_lightning as pl -from os import PathLike from itertools import chain from torch.optim import lr_scheduler -from typing import Dict, Optional, Sequence, Union, Any +from typing import Dict, Optional, Sequence, Union, Any, TYPE_CHECKING from pytorch_lightning.callbacks import EarlyStopping from pytorch_lightning.utilities.memory import is_oom_error, garbage_collection_cuda @@ -58,6 +57,8 @@ from torch.utils.data import DataLoader, random_split, Subset +if TYPE_CHECKING: + from os import PathLike logger = logging.getLogger(__name__) @@ -74,8 +75,8 @@ def _star_fun(fun, kwargs): class PretrainDataModule(pl.LightningDataModule): def __init__(self, - training_data: Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]] = None, - evaluation_data: Optional[Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]]] = None, + training_data: Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]] = None, + evaluation_data: Optional[Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]]] = None, partition: Optional[float] = 0.9, binary_dataset_split: bool = False, batch_size: int = 4, @@ -253,7 +254,7 @@ def __init__(self, hyper_params: Dict[str, Any] = None, output: str = 'model', spec: str = default_specs.RECOGNITION_SPEC, - model: Optional[Union[PathLike, str]] = None, + model: Optional[Union['PathLike', str]] = None, load_hyper_parameters: bool = False, len_train_set: int = -1): """ diff --git a/kraken/lib/progress.py b/kraken/lib/progress.py index e3dbbac8b..e9c8e6eb5 100644 --- a/kraken/lib/progress.py +++ b/kraken/lib/progress.py @@ -15,18 +15,20 @@ """ Handlers for rich-based progress bars. """ -from typing import Union +from typing import Union, TYPE_CHECKING from dataclasses import dataclass from pytorch_lightning.callbacks.progress.rich_progress import CustomProgress, RichProgressBar, MetricsTextColumn from rich import get_console, reconfigure -from rich.console import RenderableType from rich.progress import BarColumn, Progress, ProgressColumn, TextColumn, TimeRemainingColumn, TimeElapsedColumn, DownloadColumn from rich.text import Text -from rich.style import Style from rich.default_styles import DEFAULT_STYLES +if TYPE_CHECKING: + from rich.style import Style + from rich.console import RenderableType + __all__ = ['KrakenProgressBar', 'KrakenDownloadProgressBar', 'KrakenTrainProgressBar'] @@ -34,7 +36,7 @@ class BatchesProcessedColumn(ProgressColumn): def __init__(self): super().__init__() - def render(self, task) -> RenderableType: + def render(self, task) -> 'RenderableType': total = task.total if task.total != float("inf") else "--" return Text(f"{int(task.completed)}/{total}", style='magenta') @@ -147,11 +149,11 @@ class RichProgressBarTheme: https://rich.readthedocs.io/en/stable/style.html """ - description: Union[str, Style] = DEFAULT_STYLES['progress.description'] - progress_bar: Union[str, Style] = DEFAULT_STYLES['bar.complete'] - progress_bar_finished: Union[str, Style] = DEFAULT_STYLES['bar.finished'] - progress_bar_pulse: Union[str, Style] = DEFAULT_STYLES['bar.pulse'] - batch_progress: Union[str, Style] = DEFAULT_STYLES['progress.description'] - time: Union[str, Style] = DEFAULT_STYLES['progress.elapsed'] - processing_speed: Union[str, Style] = DEFAULT_STYLES['progress.data.speed'] - metrics: Union[str, Style] = DEFAULT_STYLES['progress.description'] + description: Union[str, 'Style'] = DEFAULT_STYLES['progress.description'] + progress_bar: Union[str, 'Style'] = DEFAULT_STYLES['bar.complete'] + progress_bar_finished: Union[str, 'Style'] = DEFAULT_STYLES['bar.finished'] + progress_bar_pulse: Union[str, 'Style'] = DEFAULT_STYLES['bar.pulse'] + batch_progress: Union[str, 'Style'] = DEFAULT_STYLES['progress.description'] + time: Union[str, 'Style'] = DEFAULT_STYLES['progress.elapsed'] + processing_speed: Union[str, 'Style'] = DEFAULT_STYLES['progress.data.speed'] + metrics: Union[str, 'Style'] = DEFAULT_STYLES['progress.description'] diff --git a/kraken/lib/ro/layers.py b/kraken/lib/ro/layers.py index 0ab931338..a90315b6e 100644 --- a/kraken/lib/ro/layers.py +++ b/kraken/lib/ro/layers.py @@ -4,9 +4,10 @@ import torch from torch import nn -from typing import Tuple +from typing import Tuple, TYPE_CHECKING -from kraken.lib.vgsl import VGSLBlock +if TYPE_CHECKING: + from kraken.lib.vgsl import VGSLBlock # all tensors are ordered NCHW, the "feature" dimension is C, so the output of # an LSTM will be put into C same as the filters of a CNN. @@ -38,7 +39,7 @@ def get_shape(self, input: Tuple[int, int, int, int]) -> Tuple[int, int, int, in """ return input - def get_spec(self, name) -> VGSLBlock: + def get_spec(self, name) -> 'VGSLBlock': """ Generates a VGSL spec block from the layer instance. """ diff --git a/kraken/lib/ro/model.py b/kraken/lib/ro/model.py index 6724500d2..4dbb762de 100644 --- a/kraken/lib/ro/model.py +++ b/kraken/lib/ro/model.py @@ -23,11 +23,9 @@ import torch.nn.functional as F import pytorch_lightning as pl -from os import PathLike from torch.optim import lr_scheduler from dataclasses import dataclass, field -from torch.nn import Module -from typing import Dict, Optional, Sequence, Union, Any, Literal, List +from typing import Dict, Optional, Sequence, Union, Any, Literal, List, TYPE_CHECKING from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor @@ -39,6 +37,9 @@ from torch.utils.data import DataLoader, Subset +if TYPE_CHECKING: + from os import PathLike + from torch.nn import Module logger = logging.getLogger(__name__) @@ -48,7 +49,7 @@ class DummyVGSLModel: hyper_params: Dict[str, int] = field(default_factory=dict) user_metadata: Dict[str, List] = field(default_factory=dict) one_channel_mode: Literal['1', 'L'] = '1' - ptl_module: Module = None + ptl_module: 'Module' = None model_type: str = 'unknown' def __post_init__(self): @@ -67,8 +68,8 @@ class ROModel(pl.LightningModule): def __init__(self, hyper_params: Dict[str, Any] = None, output: str = 'model', - training_data: Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]] = None, - evaluation_data: Optional[Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]]] = None, + training_data: Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]] = None, + evaluation_data: Optional[Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]]] = None, partition: Optional[float] = 0.9, num_workers: int = 1, format_type: Literal['alto', 'page', 'xml'] = 'xml', diff --git a/kraken/lib/segmentation.py b/kraken/lib/segmentation.py index 662c51eb4..efc35be8b 100644 --- a/kraken/lib/segmentation.py +++ b/kraken/lib/segmentation.py @@ -15,7 +15,6 @@ """ Processing for baseline segmenter output """ -import PIL import torch import logging import numpy as np @@ -40,14 +39,14 @@ from skimage.morphology import skeletonize from skimage.transform import PiecewiseAffineTransform, SimilarityTransform, AffineTransform, warp -from typing import List, Tuple, Union, Dict, Sequence, Optional, Literal - -from kraken.containers import Segmentation +from typing import List, Tuple, Union, Dict, Sequence, Optional, Literal, TYPE_CHECKING from kraken.lib import default_specs -from kraken.lib.vgsl import TorchVGSLModel from kraken.lib.exceptions import KrakenInputException +if TYPE_CHECKING: + from kraken.lib.vgsl import TorchVGSLModel + from kraken.containers import Segmentation logger = logging.getLogger('kraken') @@ -631,7 +630,7 @@ def _find_closest_point(pt, intersects): return env_up, env_bottom -def calculate_polygonal_environment(im: PIL.Image.Image = None, +def calculate_polygonal_environment(im: Image.Image = None, baselines: Sequence[Sequence[Tuple[int, int]]] = None, suppl_obj: Sequence[Sequence[Tuple[int, int]]] = None, im_feats: np.ndarray = None, @@ -643,27 +642,23 @@ def calculate_polygonal_environment(im: PIL.Image.Image = None, environment around each baseline. Args: - im (PIL.Image): grayscale input image (mode 'L') - baselines (sequence): List of lists containing a single baseline per - entry. - suppl_obj (sequence): List of lists containing additional polylines - that should be considered hard boundaries for - polygonizaton purposes. Can be used to prevent - polygonization into non-text areas such as - illustrations or to compute the polygonization of - a subset of the lines in an image. - im_feats (numpy.array): An optional precomputed seamcarve energy map. - Overrides data in `im`. The default map is - `gaussian_filter(sobel(im), 2)`. - scale (tuple): A 2-tuple (h, w) containing optional scale factors of - the input. Values of 0 are used for aspect-preserving - scaling. `None` skips input scaling. - topline (bool): Switch to change default baseline location for offset - calculation purposes. If set to False, baselines are - assumed to be on the bottom of the text line and will - be offset upwards, if set to True, baselines are on the - top and will be offset downwards. If set to None, no - offset will be applied. + im: grayscale input image (mode 'L') + baselines: List of lists containing a single baseline per entry. + suppl_obj: List of lists containing additional polylines that should be + considered hard boundaries for polygonizaton purposes. Can + be used to prevent polygonization into non-text areas such + as illustrations or to compute the polygonization of a + subset of the lines in an image. + im_feats: An optional precomputed seamcarve energy map. Overrides data + in `im`. The default map is `gaussian_filter(sobel(im), 2)`. + scale: A 2-tuple (h, w) containing optional scale factors of the input. + Values of 0 are used for aspect-preserving scaling. `None` skips + input scaling. + topline: Switch to change default baseline location for offset + calculation purposes. If set to False, baselines are assumed + to be on the bottom of the text line and will be offset + upwards, if set to True, baselines are on the top and will be + offset downwards. If set to None, no offset will be applied. raise_on_error: Raises error instead of logging them when they are not-blocking Returns: @@ -745,11 +740,10 @@ def polygonal_reading_order(lines: Sequence[Dict], and applies it to the input. Args: - lines (Sequence): List of tuples containing the baseline and its - polygonization. - regions (Sequence): List of region polygons. - text_direction (str): Set principal text direction for column ordering. - Can be 'lr' or 'rl' + lines: List of tuples containing the baseline and its polygonization. + regions: List of region polygons. + text_direction: Set principal text direction for column ordering. Can + be 'lr' or 'rl' Returns: The indices of the ordered input. @@ -798,14 +792,14 @@ def polygonal_reading_order(lines: Sequence[Dict], return ordered_idxs -def is_in_region(line, region) -> bool: +def is_in_region(line: geom.LineString, region: geom.Polygon) -> bool: """ Tests if a line is inside a region, i.e. if the mid point of the baseline is inside the region. Args: - line (geom.LineString): line to test - region (geom.Polygon): region to test against + line: line to test + region: region to test against Returns: False if line is not inside region, True otherwise. @@ -818,7 +812,7 @@ def neural_reading_order(lines: Sequence[Dict], text_direction: str = 'lr', regions: Optional[Sequence[geom.Polygon]] = None, im_size: Tuple[int, int] = None, - model: TorchVGSLModel = None, + model: 'TorchVGSLModel' = None, class_mapping: Dict[str, int] = None) -> Sequence[int]: """ Given a list of baselines and regions, calculates the correct reading order @@ -911,9 +905,8 @@ def scale_regions(regions: Sequence[Tuple[List[int], List[int]]], Scales baselines/polygon coordinates by a certain factor. Args: - lines (Sequence): List of tuples containing the baseline and it's - polygonization. - scale (float or tuple of floats): Scaling factor + lines: List of tuples containing the baseline and its polygonization. + scale: Scaling factor """ if isinstance(scale, float): scale = (scale, scale) @@ -928,9 +921,8 @@ def scale_polygonal_lines(lines: Sequence[Tuple[List, List]], scale: Union[float Scales baselines/polygon coordinates by a certain factor. Args: - lines (Sequence): List of tuples containing the baseline and it's - polygonization. - scale (float or tuple of floats): Scaling factor + lines: List of tuples containing the baseline and its polygonization. + scale: Scaling factor """ if isinstance(scale, float): scale = (scale, scale) @@ -974,11 +966,11 @@ def compute_polygon_section(baseline: Sequence[Tuple[int, int]], baseline will be extrapolated to the polygon edge. Args: - baseline (list): A polyline ((x1, y1), ..., (xn, yn)) - boundary (list): A bounding polygon around the baseline (same format as - baseline). - dist1 (int): Absolute distance along the baseline of the first point. - dist2 (int): Absolute distance along the baseline of the second point. + baseline: A polyline ((x1, y1), ..., (xn, yn)) + boundary: A bounding polygon around the baseline (same format as + baseline). Last and first point are automatically connected. + dist1: Absolute distance along the baseline of the first point. + dist2: Absolute distance along the baseline of the second point. Returns: A sequence of polygon points. @@ -1032,7 +1024,7 @@ def compute_polygon_section(baseline: Sequence[Tuple[int, int]], return tuple(o) -def extract_polygons(im: Image.Image, bounds: Segmentation) -> Image.Image: +def extract_polygons(im: Image.Image, bounds: 'Segmentation') -> Image.Image: """ Yields the subimages of image im defined in the list of bounding polygons with baselines preserving order. diff --git a/kraken/lib/train.py b/kraken/lib/train.py index a3136c928..a692a5376 100644 --- a/kraken/lib/train.py +++ b/kraken/lib/train.py @@ -23,11 +23,10 @@ import torch.nn.functional as F import pytorch_lightning as pl -from os import PathLike from torchmetrics.classification import MultilabelAccuracy, MultilabelJaccardIndex from torchmetrics.text import CharErrorRate, WordErrorRate from torch.optim import lr_scheduler -from typing import Callable, Dict, Optional, Sequence, Union, Any, Literal +from typing import Callable, Dict, Optional, Sequence, Union, Any, Literal, TYPE_CHECKING from pytorch_lightning.callbacks import Callback, EarlyStopping, BaseFinetuning, LearningRateMonitor from kraken.containers import Segmentation, XMLPage @@ -42,6 +41,8 @@ from torch.utils.data import DataLoader, random_split, Subset +if TYPE_CHECKING: + from os import PathLike logger = logging.getLogger(__name__) @@ -76,7 +77,7 @@ def __init__(self, max_epochs: int = 100, freeze_backbone=-1, pl_logger: Union[pl.loggers.logger.Logger, str, None] = None, - log_dir: Optional[PathLike] = None, + log_dir: Optional['PathLike'] = None, *args, **kwargs): kwargs['enable_checkpointing'] = False @@ -200,10 +201,10 @@ def __init__(self, output: str = 'model', spec: str = default_specs.RECOGNITION_SPEC, append: Optional[int] = None, - model: Optional[Union[PathLike, str]] = None, + model: Optional[Union['PathLike', str]] = None, reorder: Union[bool, str] = True, - training_data: Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]] = None, - evaluation_data: Optional[Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]]] = None, + training_data: Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]] = None, + evaluation_data: Optional[Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]]] = None, partition: Optional[float] = 0.9, binary_dataset_split: bool = False, num_workers: int = 1, @@ -494,8 +495,13 @@ def validation_step(self, batch, batch_idx): for i in range(self.hparams.batch_size): count = self.hparams.batch_size * batch_idx + i if count < 16: - self.logger.experiment.add_image(f'Validation #{count}, target: {decoded_targets[i]}', batch['image'][i], self.global_step, dataformats="CHW") - self.logger.experiment.add_text(f'Validation #{count}, target: {decoded_targets[i]}', pred[i], self.global_step) + self.logger.experiment.add_image(f'Validation #{count}, target: {decoded_targets[i]}', + batch['image'][i], + self.global_step, + dataformats="CHW") + self.logger.experiment.add_text(f'Validation #{count}, target: {decoded_targets[i]}', + pred[i], + self.global_step) def on_validation_epoch_end(self): accuracy = 1.0 - self.val_cer.compute() @@ -686,9 +692,9 @@ def __init__(self, message: Callable[[str], None] = lambda *args, **kwargs: None, output: str = 'model', spec: str = default_specs.SEGMENTATION_SPEC, - model: Optional[Union[PathLike, str]] = None, - training_data: Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]] = None, - evaluation_data: Optional[Union[Sequence[Union[PathLike, str]], Sequence[Dict[str, Any]]]] = None, + model: Optional[Union['PathLike', str]] = None, + training_data: Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]] = None, + evaluation_data: Optional[Union[Sequence[Union['PathLike', str]], Sequence[Dict[str, Any]]]] = None, partition: Optional[float] = 0.9, num_workers: int = 1, force_binarization: bool = False, diff --git a/kraken/lib/util.py b/kraken/lib/util.py index 66ac7faec..6eee8ab92 100644 --- a/kraken/lib/util.py +++ b/kraken/lib/util.py @@ -8,13 +8,15 @@ import numpy as np from PIL import Image -from os import PathLike -from typing import Union, Callable, Optional, Literal +from typing import Union, Callable, Optional, Literal, TYPE_CHECKING from kraken.lib import functional_im_transforms as F_t from kraken.containers import BBoxLine -from kraken.exceptions import KrakenInputException +from kraken.lib.exceptions import KrakenInputException + +if TYPE_CHECKING: + from os import PathLike __all__ = ['pil2array', 'array2pil', 'is_bitonal', 'make_printable', 'get_im_str', 'parse_gt_path'] @@ -101,9 +103,9 @@ def make_printable(char: str) -> str: return unicodedata.name(char) -def parse_gt_path(path: Union[str, PathLike], +def parse_gt_path(path: Union[str, 'PathLike'], suffix: str = '.gt.txt', - split: Callable[[Union[PathLike, str]], str] = F_t.default_split, + split: Callable[[Union['PathLike', str]], str] = F_t.default_split, skip_empty_lines: bool = True, base_dir: Optional[Literal['L', 'R']] = None, text_direction: Literal['horizontal-lr', 'horizontal-rl', 'vertical-lr', 'vertical-rl'] = 'horizontal-lr') -> BBoxLine: diff --git a/kraken/lib/xml.py b/kraken/lib/xml.py index d65535489..3806fb3f4 100644 --- a/kraken/lib/xml.py +++ b/kraken/lib/xml.py @@ -18,18 +18,20 @@ import re import logging -from os import PathLike from pathlib import Path from itertools import groupby from lxml import etree -from typing import Union, Dict, Any, Sequence, Tuple, Literal, Optional, List +from typing import Union, Dict, Any, Sequence, Tuple, Literal, Optional, List, TYPE_CHECKING from collections import defaultdict from kraken.containers import Segmentation, BaselineLine, Region logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from os import PathLike + __all__ = ['XMLPage'] @@ -61,7 +63,7 @@ class XMLPage(object): type: Literal['baselines', 'bbox'] = 'baselines' base_dir: Optional[Literal['L', 'R']] = None - imagename: PathLike = None + imagename: 'PathLike' = None image_size: Tuple[int, int] = None _orders: Dict[str, Dict[str, Any]] = None has_tags: bool = False @@ -70,7 +72,7 @@ class XMLPage(object): _split_set: Optional[List] = None def __init__(self, - filename: Union[str, PathLike], + filename: Union[str, 'PathLike'], filetype: Literal['xml', 'alto', 'page'] = 'xml'): super().__init__() self.filename = Path(filename) diff --git a/kraken/repo.py b/kraken/repo.py index 6f484db4a..af22f1ac9 100644 --- a/kraken/repo.py +++ b/kraken/repo.py @@ -20,13 +20,16 @@ import logging import requests -from os import PathLike from pathlib import Path from contextlib import closing -from typing import Callable, Any +from typing import Callable, Any, TYPE_CHECKING from kraken.lib.exceptions import KrakenRepoException +if TYPE_CHECKING: + from os import PathLike + + __all__ = ['get_model', 'get_description', 'get_listing', 'publish_model'] logger = logging.getLogger(__name__) @@ -35,7 +38,7 @@ SUPPORTED_MODELS = set(['kraken_pytorch']) -def publish_model(model_file: [str, PathLike] = None, +def publish_model(model_file: [str, 'PathLike'] = None, metadata: dict = None, access_token: str = None, callback: Callable[[int, int], Any] = lambda: None, diff --git a/kraken/rpred.py b/kraken/rpred.py index da7e6afcf..c2c31ebef 100644 --- a/kraken/rpred.py +++ b/kraken/rpred.py @@ -21,18 +21,19 @@ import logging import dataclasses -from PIL import Image from functools import partial from collections import defaultdict -from typing import List, Tuple, Optional, Generator, Union, Dict, Sequence +from typing import List, Tuple, Optional, Generator, Union, Dict, Sequence, TYPE_CHECKING from kraken.containers import BaselineOCRRecord, BBoxOCRRecord, ocr_record, Segmentation from kraken.lib.util import get_im_str, is_bitonal -from kraken.lib.models import TorchSeqRecognizer from kraken.lib.segmentation import extract_polygons from kraken.lib.exceptions import KrakenInputException from kraken.lib.dataset import ImageInputTransforms +if TYPE_CHECKING: + from PIL import Image + from kraken.lib.models import TorchSeqRecognizer __all__ = ['mm_rpred', 'rpred'] diff --git a/kraken/serialization.py b/kraken/serialization.py index 930f3eeb3..682635041 100644 --- a/kraken/serialization.py +++ b/kraken/serialization.py @@ -18,14 +18,16 @@ import logging import datetime -from os import PathLike from pkg_resources import get_distribution -from collections import Counter -from kraken.containers import Segmentation, ProcessingStep from kraken.lib.util import make_printable -from typing import List, Tuple, Iterable, Optional, Sequence, Literal +from typing import List, Tuple, Iterable, Optional, Sequence, Literal, TYPE_CHECKING + +if TYPE_CHECKING: + from os import PathLike + from collections import Counter + from kraken.containers import Segmentation, ProcessingStep logger = logging.getLogger(__name__) @@ -68,13 +70,13 @@ def max_bbox(boxes: Iterable[Sequence[int]]) -> Tuple[int, int, int, int]: return o -def serialize(results: Segmentation, +def serialize(results: 'Segmentation', image_size: Tuple[int, int] = (0, 0), writing_mode: Literal['horizontal-tb', 'vertical-lr', 'vertical-rl'] = 'horizontal-tb', scripts: Optional[Iterable[str]] = None, - template: [PathLike, str] = 'alto', + template: ['PathLike', str] = 'alto', template_source: Literal['native', 'custom'] = 'native', - processing_steps: Optional[List[ProcessingStep]] = None) -> str: + processing_steps: Optional[List['ProcessingStep']] = None) -> str: """ Serializes recognition and segmentation results into an output document. @@ -243,11 +245,11 @@ def _load_template(name): def render_report(model: str, chars: int, errors: int, - char_confusions: Counter, - scripts: Counter, - insertions: Counter, + char_confusions: 'Counter', + scripts: 'Counter', + insertions: 'Counter', deletions: int, - substitutions: Counter) -> str: + substitutions: 'Counter') -> str: """ Renders an accuracy report.