Pipeline steps

Pipeline step base class.

Pipeline steps are implemented by subclassing this class and overriding the run() method.

45
46
47
48
49
50
51
52
53
54

def run(self, collection: Collection) -> Collection:
    """
    Run the pipeline step.

    Arguments:
        collection: Input collection

    Returns:
        A new collection, updated with the results of the pipeline step.
    """

Bases: PipelineStep

Base for image preprocessing steps.

This is a base class for all image preprocessing steps. Subclasses define their image processing operation by overriding the op() method. This step does not alter the original image. Instead, a new copy of the image is saved in the directory specified by ProcessImages.output_directory. The PageNode's image path is then updated to point to the new processed image.

Attributes:

488
489
490
491
492
493
494
495
496
497
498

def op(self, image: NumpyImage) -> NumpyImage:
    """
    Perform the image processing operation on `image`.

    Arguments:
        image: Input image.

    Returns:
        A processed version of `image`.
    """
    pass

Bases: ProcessImages

Binarize images.

Runs image binarization on the collection's images. Saves the resulting images in a directory named binarized. All subsequent pipeline steps will use the binarized images.

Example YAML:

- step: Binarization

Bases: PipelineStep

Run model inference.

This is a generic pipeline step for any type of model inference. This step always runs the model on the images of the collection's leaf nodes.

Example YAML:

- step: Inference
  settings:
    model: DiT
    model_settings:
      model: ...

78
79
80
81
82

def __init__(self, model_class, model_kwargs, generation_kwargs):
    self.model_class = model_class
    self.model_kwargs = model_kwargs
    self.generation_kwargs = generation_kwargs
    self.model = None

Bases: Inference

Run a segmentation model.

See Segmentation models for available models.

Example YAML:

- step: Segmentation
  settings:
    model: yolo
    model_settings:
      model: Riksarkivet/yolov9-regions-1

78
79
80
81
82

def __init__(self, model_class, model_kwargs, generation_kwargs):
    self.model_class = model_class
    self.model_kwargs = model_kwargs
    self.generation_kwargs = generation_kwargs
    self.model = None

Bases: Inference

Run a text recognition model.

See Text recognition models for available models.

Example YAML:

- step: TextRecognition
  settings:
    model: TrOCR
    model_settings:
      model: Riksarkivet/trocr-base-handwritten-hist-swe-2

78
79
80
81
82

def __init__(self, model_class, model_kwargs, generation_kwargs):
    self.model_class = model_class
    self.model_kwargs = model_kwargs
    self.generation_kwargs = generation_kwargs
    self.model = None

Bases: PipelineStep

Remove nodes based on a given condition.

This is a generic pruning (filtering) step which removes nodes (segments, lines, words) based on the given condition. The condition is a function f such that f(node) == True if node should be removed from the tree. This step runs f on all nodes, at all segmentation levels. See the RemoveLowTextConfidence[Lines|Regions|Pages] steps for examples of how to formulate condition.

Parameters:

383
384
385
386
387
388
389

def __init__(self, condition: Callable[[Node], bool]):
    """
    Arguments:
        condition: A function `f` such that `f(node) == True` if
            `node` should be removed from the document tree.
    """
    self.condition = condition

Bases: Prune

Remove all pages where the average text confidence score is below threshold.

Example YAML:

- step: RemoveLowTextConfidencePages
  settings:
    threshold: 0.8

Parameters:

452
453
454
455
456
457
458
459

def __init__(self, threshold: float):
    """
    Arguments:
        threshold: Confidence score threshold.
    """
    super().__init__(
        lambda node: node.parent and node.parent.is_root() and metrics.average_text_confidence(node) < threshold
    )

Bases: Prune

Remove all regions where the average text confidence score is below threshold.

Example YAML:

- step: RemoveLowTextConfidenceRegions
  settings:
    threshold: 0.8

Parameters:

430
431
432
433
434
435
436
437

def __init__(self, threshold: float):
    """
    Arguments:
        threshold: Confidence score threshold.
    """
    super().__init__(
        lambda node: all(child.is_line() for child in node) and metrics.average_text_confidence(node) < threshold
    )

Bases: Prune

Remove all lines with text confidence score below threshold.

Example YAML:

- step: RemoveLowTextConfidenceLines
  settings:
    threshold: 0.8

Parameters:

410
411
412
413
414
415

def __init__(self, threshold: float):
    """
    Arguments:
        threshold: Confidence score threshold.
    """
    super().__init__(lambda node: node.is_line() and metrics.line_text_confidence(node) < threshold)

Bases: PipelineStep

Order regions and lines by reading order.

This step orders the pages' first- and second-level segments (corresponding to regions and lines). Both the regions and their lines are ordered using reading_order.order_regions.

Parameters:

268
269
270
271
272
273
274
275
276
277

def __init__(self, two_page: Literal["auto"] | bool = False):
    """
    Arguments:
        two_page: Whether the page is a two-page spread. Three modes:
            - 'auto': determine heuristically for each page using
                `layout.is_twopage`
            - True: assume all pages are spreads
            - False: assume all pages are single pages
    """
    self.two_page = two_page

Bases: PipelineStep

Order lines top-down.

This step orders the lines within each region top-down.

Example YAML:

- step: OrderLines

Bases: PipelineStep

Segment lines into words.

This step segments lines of text into words. It estimates the word boundaries from the recognized text, which means that this step must be run after a line-based text recognition model.

See also <models.huggingface.trocr.WordLevelTrOCR>, which is a version of TrOCR that outputs word-level text directly using a more sophisticated method.

Example YAML:

- step: WordSegmentation

Bases: PipelineStep

Export results.

Exports the current state of the collection in the given format. This step is typically the last step of a pipeline, however, it can be inserted at any pipeline stage. For example, you could put an Export step before a post processing step in order to save a copy without post processing. A pipeline can include as many Export steps as you like.

See Export formats or the <serialization.serialization> module for more details about each export format.

Example:

- step: Export
  settings:
    format: Alto
    dest: alto-outputs

Parameters:

239
240
241
242
243
244
245
246
247
248
249
250
251

def __init__(
    self,
    dest: str,
    format: Literal["alto", "page", "txt", "json"],
    **serializer_kwargs,
):
    """
    Arguments:
        dest: Output directory.
        format: Output format as a string.
    """
    self.serializer = get_serializer(format, **serializer_kwargs)
    self.dest = dest

Bases: PipelineStep

Export the collection's images.

This step writes all existing images (regions, lines, etc.) in the collection to disk. The exported images are the images that have been passed to previous Inference steps and the images that would be passed to a following Inference step.

Example YAML:

- step: ExportImages
  settings:
    dest: exported_images

Parameters:

337
338
339
340
341
342
343

def __init__(self, dest: str):
    """
    Arguments:
        dest: Destination directory.
    """
    self.dest = dest
    os.makedirs(self.dest, exist_ok=True)

Bases: PipelineStep

Break the pipeline! Used for testing.

Example YAML:

- step: Break

Bases: PipelineStep

Import segmentation from PageXML files.

This step replicates the line segmentation from PageXML files. It can be used to import ground truth segmentation for evaluation purposes.

Example YAML:

- step: ImportSegmentation
  settings:
    source: /path/to/pageXMLs

Parameters:

125
126
127
128
129
130
131
132

def __init__(self, source: str):
    """
    Arguments:
        source: Path to a directory with PageXML files. The XML files
            must have the same names as the input image files (ignoring
            the file extension).
    """
    self.source = source