dhSegment: A generic deep-learning approach for document segmentation

In plain words

Old scans come with many recurring jobs: cut out the page, find the lines of text, mark the decorations, separate a photo from the card it is glued to. The usual approach has been one hand-built tool per job. dhSegment instead uses one image-to-image network for all of them. You give it scans together with masks that paint the wanted regions, and it learns to label every pixel. Simple follow-up steps then turn that label map into the final output, such as a rectangle, a set of lines or a set of boxes.

The network reuses an image classifier whose early layers were already trained on ordinary photos, so it needs comparatively little labelled material. Across the tasks, the authors change little beyond the training data and the final post-processing. The one setting that has to be chosen for each task is how large the input image is resized. The authors report results that are close to or better than specialised systems.

TL;DR

dhSegment is a single encoder-decoder network for pixel-wise segmentation of historical documents, trained with one shared recipe for several tasks. Its contracting path is a ResNet-50 initialised from ImageNet weights, and its expanding path upsamples back to the input resolution §III-B. Only thresholding, morphology, component filtering and polygon fitting follow the network §III-C. Author-reported results are close to state of the art on page extraction, baseline detection and layout analysis, and are compared with an ornament baseline at a different IoU tab. I tab. II tab. III tab. IV. Several details (learning rate, augmentation and resize settings) differ between the paper and its code, as listed below.

Problem

Historical document work needs many separate segmentation tools, such as page extraction, text-line detection, layout analysis and illustration extraction. Each tool has usually been built for one document type, which makes it hard for non-specialists to reuse. The authors argue that a generic approach can now match dedicated systems §I.

Method

1. Two stages. A fully convolutional network maps the input image to a probability for each pixel and each class. Task-specific post-processing then turns those probabilities into the required output §III-A fig. 1.
2. Training masks. Training labels are rasterised into per-pixel masks, and these masks are the network's training targets §III-A.
3. Contracting path. A ResNet-50 encoder (the yellow blocks in the architecture figure) with five levels; each level halves the spatial size §III-B fig. 2. The encoder starts from pretrained weights, which the authors report shortens training and makes the model less sensitive to outliers §III-B.
4. Expanding path. Five stages. Each stage upsamples the deeper map bilinearly, concatenates it with the encoder map of the same size, and applies a 3x3 convolution followed by a ReLU §III-B fig. 2. The channels of stages 4 and 5 are first reduced to 512 with a 1x1 convolution §III-B.
5. Output. A final convolution gives one score per class at input resolution, followed by a softmax §III-B estimator_fn.py.
6. Parameter count. The paper states 32.8M parameters, of which 9.36M are trained from scratch; a footnote says the 1.57M parameters of the reduction blocks could be counted as fine-tuned too, giving 7.79M §III-B. The count was not recomputed here.
7. Post-processing, per task §III-C:
   • thresholding, either at a fixed value or with Otsu's method;
   • opening and closing on the binary mask;
   • removal of small connected components;
   • vectorisation into quadrilaterals, boxes or polylines.
8. Training recipe §III-D: L2 weight decay, an initial learning rate between 1e-5 and 1e-4 with exponential decay 0.95, Adam with Xavier initialisation, and batch renormalisation (r_min 0.1, r_max 100, d_max 1). Images are resized to a total of 6e5 to 1e6 pixels, cropped to 300 x 300 patches with a margin, and augmented by rotation, scaling and mirroring. The resize size is the one parameter that has to be chosen per task.

# dhSegment training step  [§III-A] [§III-D]
# Input:  image (H, W, 3), mask (H, W) with one class id per pixel
# Output: scalar loss; network weights updated in place

def training_step(image, mask, network, optimizer):
    # Resize image and mask to the task's pixel budget, keeping aspect ratio
    # [§III-D] [§IV-A] [§IV-B]
    image, mask = resize_to_pixel_budget(image, mask, budget)

    # Optional augmentation: rotation of about +/-0.2 rad, scaling of
    # 0.8 to 1.2, mirroring  [§III-D]
    image, mask = augment(image, mask)

    # Cut 300 x 300 crops, each carrying a context margin  [§III-D]
    # [code: dh_segment/io/input.py#L85-L90]
    crops_image, crops_mask = random_crops(image, mask, size=300)

    # Per-pixel class scores, shape (B, h, w, classes)  [§III-B]
    logits = network(crops_image, training=True)

    # Per-pixel cross-entropy averaged per image; the margin is excluded
    # [code: dh_segment/estimator_fn.py#L129-L137]
    loss = masked_mean_cross_entropy(logits, crops_mask)

    # Adam step; learning rate starts at 1e-5 and decays by 0.95 every
    # 200 steps  [code: dh_segment/estimator_fn.py#L151-L158]
    optimizer.apply(loss, learning_rate=lr_schedule(step))
    return loss

# Page extraction at inference  [§IV-A]
# Input:  probs (H, W), page probability per pixel
# Output: four corner points of the page quadrilateral

def extract_page(probs):
    # Otsu threshold after a 5 x 5 Gaussian blur  [§IV-A]
    # [code: dh_segment/post_processing/binarization.py#L15-L20]
    mask = otsu_threshold(probs)

    # Morphological opening, then closing; the paper gives no kernel size
    # [§IV-A]
    # [code: dh_segment/post_processing/binarization.py#L28]
    mask = close(open(mask, kernel=5), kernel=5)

    # Convex hull of the largest blob; keep the hull points nearest to
    # the four image corners  [§IV-A]
    # [code: dh_segment/post_processing/boxes_detection.py#L123-L153]
    return quadrilateral_from_hull(mask)

# Baseline detection  [§IV-B]
# Input:  probs (H, W), baseline probability per pixel
# Output: list of polylines, one per kept component

def detect_baselines(probs):
    # Gaussian smoothing with sigma 1.5  [§IV-B]
    # [code: exps/cbad/demo_processing.py#L54]
    smooth = gaussian_blur(probs, sigma=1.5)

    # Hysteresis: low threshold 0.2 selects candidate pixels; keep only
    # components that also contain a pixel above 0.4  [§IV-B]
    mask = hysteresis(smooth, low=0.2, high=0.4)

    # Skeleton of each component, then a path between its two farthest
    # endpoints  [§IV-B]
    # [code: dh_segment/post_processing/line_vectorization.py#L237-L287]
    return [longest_path(component) for component in components(mask)]

Hyperparameters

Results

Author-reported. Each task is averaged over five runs, and post-processing parameters are chosen on each task's evaluation set §IV.

The ornament comparison is not like for like: the baseline is reported at IoU 0.5 only, while dhSegment is reported at 0.7, 0.8 and 0.9 tab. IV. The ornament and photo datasets are private §IV-D §IV-E. The discussion states that the approach needs little training data and sometimes under an hour of training §V.

Paper vs code

The official repository is at dhlab-epfl/dhSegment and belongs to the authors' lab (EPFL Digital Humanities Laboratory). Its licence is GPL-3.0; per the licensing rules, this card cites it by path and line and paraphrases the logic, quoting no code.

• Scope of code. The repository has page extraction and cBAD baseline experiments only process.py process.py. It has no code for layout (DIVA-HisDB), ornaments, photo extraction or the 5-pixel baseline label masks.
• Learning rate. Paper: initial value 1e-5 to 1e-4 §III-D. Code default: 1e-5 params_config.py. Shipped general config also uses 1e-5 general_config.json. The demo configs use 5e-5 demo_cbad_config.json demo_page_config.json.
• Decay schedule. Paper gives only the rate 0.95. Code decays every 200 steps, not staircased estimator_fn.py.
• Scaling augmentation. Paper: scale 0.8 to 1.2 §III-D. Code default: max scaling 0.05 params_config.py. General config: 0.2 general_config.json.
• Mirroring. Paper mentions mirroring §III-D. The general config turns left-right flips off general_config.json; the demos turn it on demo_cbad_config.json.
• Page input size. Paper: 6e5 pixels §IV-A. The page demo config does not set a size, so the code default of 7.2e5 applies params_config.py demo_page_config.json.
• Baseline input size. Paper: 1e6 pixels §IV-B. Matches the cBAD demo demo_cbad_config.json.
• Baseline post-processing. Paper: sigma 1.5, thresholds 0.4 and 0.2 §IV-B. The cBAD demo matches demo_processing.py. The library defaults differ: sigma 2.5, thresholds 0.8 and 0.9 process.py.
• Page binarisation. Paper: Otsu §IV-A. Code blurs with a 5 x 5 Gaussian before Otsu, which the paper does not mention, and carries a TODO about it binarization.py. The demo uses Otsu demo_processing.py.
• Page shape. Paper: quadrilateral §IV-A. The page demo instead fits a minimum rectangle demo_processing.py. The library extract_page fits a quadrilateral process.py.
• Baseline polylines. Paper: each component becomes a polygonal line §IV-B. Code: a skeleton with a path between the two farthest endpoints, sampled every fifth point line_vectorization.py.
• Encoder batch norm. The encoder is built with is_training=False, so its batch-norm layers run in inference mode during training model.py. The decoder follows the training flag model.py. The paper does not discuss this.
• Loss. The paper does not name a loss. Code uses per-pixel softmax cross-entropy estimator_fn.py. Focal loss, class weights and local-entropy weighting exist but default to off params_config.py.
• Margin. The paper adds a margin without giving its value §III-D. Code default is 16 px, and the demos set 0 params_config.py.
• Batch renormalisation. Matches the paper model.py.
• Weight decay. Matches the paper params_config.py.
• Patch size. 300 x 300 matches params_config.py.
• Framework. The code imports tensorflow.contrib model.py, so it targets TensorFlow 1.x. The paper says only TensorFlow.

Pitfalls

• Vague training recipe. The paper gives a learning-rate range, not a value per task, and no decay step. Start from the demo values and tune.
• Resize is the one per-task choice, but it is not uniform. The paper says the resize size sets the receptive field §III-D. Its own values differ per task: 6e5 pages, 1e6 baselines, 8e5 ornaments, no resize for layout §IV-A §IV-B §IV-C §IV-D.
• Encoder batch norm runs in inference mode during training. The encoder is built with is_training=False model.py. From this reading (not tested), its normalisation statistics are not updated during training. Verify this before changing the training setup.
• Loss and margin are implicit. Margin pixels are excluded from the loss and cropped at prediction estimator_fn.py estimator_fn.py. The paper says neither.
• Morphology kernel. The paper names opening and closing but not their size. The code default is 5 x 5 binarization.py.
• Quadrilateral fallback is dead code. Both branches of the quadrilateral fit return four corners, so the fallback to minimum rectangles cannot trigger boxes_detection.py.
• Selection bias. Post-processing parameters are picked on the evaluation set §IV. Reported numbers are therefore optimistic relative to a held-out choice.
• Mixed comparisons. The baseline result is not the best on the simple track tab. II. The layout result is below the best entry tab. III. The ornament comparison uses different IoU thresholds tab. IV.
• Unverified claims. The parameter counts and the photo and ornament numbers are not reproducible from the repository, because that code is absent.
• Trained weights and labels. The 5-pixel baseline mask generator and the ornament and layout datasets are not in the repository, so their results cannot be reproduced from the release.

Lineage

• Builds on: ResNet-50 as the encoder, and the contracting and expanding path terminology of U-Net §III-B. The paper names both without arXiv links in the text.
• Compared against: PageNet and GrabCut for page extraction, the cBAD ICDAR 2017 systems for baselines, and the ICDAR 2017 DIVA-HisDB systems for layout §IV-A §IV-B §IV-C.
• Superseded by: not stated in the wiki's content.