modified-inference-pipeline.sh

#!/bin/bash
VERSION=2.1.2
set -e
STOPONERROR=1

## TOGGLES: set		1 = on, 0 = off 	to add/ remove recognition steps from the pipeline

# Baseline Recognition:
BASELINELAYPA=1

# Region (Field) Recognition (only pick one of these two)

# AI-powered (trained) Region Recognition:
REGIONLAYPA=0
# Rule-based Region Recognition (needs Baselines to work):
REGIONRULEBASED=1

# Text Recognition
HTRLOGHI=0

# Recalculate Reading Order (custom region reading order recalculation script)
NEWREADINGORDER=1

## MODEL SELECTION
# specify model paths (only needed if corresponding recognition step is set to on)
# Baseline
LAYPABASELINEMODEL=INSERT_FULL_PATH_TO_YAML_HERE
LAYPABASELINEMODELWEIGHTS=INSERT_FULLPATH_TO_PTH_HERE
# Region
LAYPAREGIONMODEL=INSERT_FULL_PATH_TO_YAML_HERE
LAYPAREGIONMODELWEIGHTS=INSERT_FULLPATH_TO_PTH_HERE
# Text Recognition
HTRLOGHIMODEL=INSERT_FULL_PATH_TO_LOGHI_HTR_MODEL_HERE


# Further settings for REGIONRULEBASED
# If the edge of baseline is closer than x pixels...
REGIONRULEBASEDBORDERMARGIN=0
# Clean borders if 1
REGIONRULEBASEDCLEANBORDERS=0
# How many threads to use for recalculating reading order
REGIONRULEBASEDTHREADS=4

# Detect language of Text Lines
DETECTLANGUAGE=1
# Interpolate word locations
SPLITWORDS=1
# BEAMWIDTH: higher makes results slightly better at the expense of lot of computation time. In general don't set higher than 10
BEAMWIDTH=4
# Used gpu ids, set to "-1" to use CPU, "0" for first, "1" for second, etc
GPU=0

# Use 2013 PageXML namespace
USE2013NAMESPACE=1


# DO NOT MODIFY BELOW THIS LINE
# ------------------------------

function check_error_and_exit {
    # $1 - The error message to display
    # $2 - The status of the last command executed before the function was called

    if [[ $STOPONERROR -eq 1 && $2 -ne 0 ]]; then
        echo "$1 has failed"
        exit 1
    fi
}

# Check for proper command-line arguments
if [ "$#" -ne 1 ]; then
    echo "Usage: $0 <path_to_images>"
    exit 1
fi

tmpdir=$(mktemp -d)
echo "Temporary directory created at: $tmpdir"

mkdir -p $tmpdir/imagesnippets/
mkdir -p $tmpdir/output

DOCKERLOGHITOOLING=loghi/docker.loghi-tooling:$VERSION
DOCKERLAYPA=loghi/docker.laypa:$VERSION
DOCKERLOGHIHTR=loghi/docker.htr:$VERSION

DOCKERGPUPARAMS=""
if [[ $GPU -gt -1 ]]; then
    DOCKERGPUPARAMS="--gpus device=${GPU}"
    echo "Using GPU ${GPU}"
fi

IMAGES_PATH=`realpath $1`

# Housekeeping: remove any existing *.done files
find $IMAGES_PATH -name '*.done' -exec rm -f "{}" \;

if [[ $USE2013NAMESPACE -eq 1 ]]; then
    namespace=" -use_2013_namespace "
fi

# First step: Laypa
if [[ $BASELINELAYPA -eq 1 ]]; then
    echo "Running Laypa baseline detection"

    LAYPA_IN=$IMAGES_PATH
    LAYPA_OUT=$IMAGES_PATH
    LAYPADIR="$(dirname "${LAYPABASELINEMODEL}")"

    docker run $DOCKERGPUPARAMS --rm -it -u $(id -u ${USER}):$(id -g ${USER}) -m 32000m --shm-size 10240m \
        -v $LAYPADIR:$LAYPADIR \
        -v $LAYPA_IN:$LAYPA_IN \
        -v $LAYPA_OUT:$LAYPA_OUT \
        $DOCKERLAYPA \
            python run.py \
            -c $LAYPABASELINEMODEL \
            -i $LAYPA_IN \
            -o $LAYPA_OUT \
            --opts MODEL.WEIGHTS "" TEST.WEIGHTS $LAYPABASELINEMODELWEIGHTS | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "Laypa baseline detection" $status

    # Set as_single_region flag by default
    as_single_region=" -as_single_region "
    echo "Laypa baseline detection done"

    if [[ $REGIONLAYPA -eq 1 ]]; then
        echo "Running Laypa region detection"

        docker run $DOCKERGPUPARAMS --rm -it -u $(id -u ${USER}):$(id -g ${USER}) -m 32000m --shm-size 10240m \
        -v $LAYPADIR:$LAYPADIR \
        -v $LAYPA_IN:$LAYPA_IN \
        -v $LAYPA_OUT:$LAYPA_OUT \
        $DOCKERLAYPA \
            python run.py \
            -c $LAYPAREGIONMODEL \
            -i $LAYPA_IN \
            -o $LAYPA_OUT \
            --opts MODEL.WEIGHTS "" TEST.WEIGHTS $LAYPAREGIONMODELWEIGHTS | tee -a $tmpdir/log.txt

        # Check if failed
        status=$?
        check_error_and_exit "Laypa region detection" $status

        as_single_region=""
        echo "Laypa region detection done"
    fi

    # Second step: extract baselines and regions
    echo "Extracting baselines and regions"

    docker run --rm -u $(id -u ${USER}):$(id -g ${USER}) \
        -v $LAYPA_IN:$LAYPA_IN \
        -v $LAYPA_OUT:$LAYPA_OUT \
        $DOCKERLOGHITOOLING \
            /src/loghi-tooling/minions/target/appassembler/bin/MinionExtractBaselines \
            -input_path_image $LAYPA_IN \
            -input_path_png $LAYPA_OUT/page/ \
            -input_path_page $LAYPA_OUT/page/ \
            -output_path_page $LAYPA_OUT/page/ \
            -recalculate_textline_contours_from_baselines \
            $as_single_region \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionExtractBaselines" $status

    echo "Extracting baselines and regions done"
fi

# Third step: cut out snippets
if [[ $HTRLOGHI -eq 1 ]]; then
    echo "Cutting out snippets"

    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
       -v $IMAGES_PATH/:$IMAGES_PATH/ \
       -v $tmpdir:$tmpdir \
       $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionCutFromImageBasedOnPageXMLNew \
           -input_path $IMAGES_PATH \
           -outputbase $tmpdir/imagesnippets/ \
           -output_type png \
           -channels 4 \
           -threads 4 $namespace| tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionCutFromImageBasedOnPageXMLNew" $status

    # Collect all the snippets in a file
    find $tmpdir/imagesnippets/ -type f -name '*.png' > $tmpdir/lines.txt

    echo "Running HTR"
    LOGHIDIR="$(dirname "${HTRLOGHIMODEL}")"

    echo docker run $DOCKERGPUPARAMS -u $(id -u ${USER}):$(id -g ${USER}) --rm -m 32000m --shm-size 10240m -ti \
        -v /tmp:/tmp \
        -v $tmpdir:$tmpdir \
        -v $LOGHIDIR:$LOGHIDIR \
        $DOCKERLOGHIHTR \
            bash -c "LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 python3 /src/loghi-htr/src/main.py \
            --model $HTRLOGHIMODEL  \
            --batch_size 64 \
            --use_mask \
            --inference_list $tmpdir/lines.txt \
            --results_file $tmpdir/results.txt \
            --gpu $GPU \
            --output $tmpdir/output/ \
            --beam_width $BEAMWIDTH " | tee -a $tmpdir/log.txt


    docker run $DOCKERGPUPARAMS -u $(id -u ${USER}):$(id -g ${USER}) --rm -m 32000m --shm-size 10240m -ti \
        -v /tmp:/tmp \
        -v $tmpdir:$tmpdir \
        -v $LOGHIDIR:$LOGHIDIR \
        $DOCKERLOGHIHTR \
            bash -c "LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 python3 /src/loghi-htr/src/main.py \
            --model $HTRLOGHIMODEL  \
            --batch_size 64 \
            --use_mask \
            --inference_list $tmpdir/lines.txt \
            --results_file $tmpdir/results.txt \
            --gpu $GPU \
            --output $tmpdir/output/ \
            --beam_width $BEAMWIDTH " | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "Loghi HTR" $status
    
    echo "Loghi HTR done"

    # Fourth step: merge results back into PageXML
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $LOGHIDIR:$LOGHIDIR \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionLoghiHTRMergePageXML \
            -input_path $IMAGES_PATH/page \
            -results_file $tmpdir/results.txt \
            -config_file $HTRLOGHIMODEL/config.json \
            -htr_code_config_file $tmpdir/output/config.json \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionLoghiHTRMergePageXML" $status
    
    echo "Merging HTR results into PageXML done"
fi

# Fifth step: Recalculate reading order
if [[ $REGIONRULEBASED -eq 1 ]]
then
    echo "Recalculating reading order"
    clean_borders=""

    if [[ $REGIONRULEBASEDCLEANBORDERS -eq 1 ]]; then
        echo "and cleaning borders"
        clean_borders=" -clean_borders "
    fi
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionRecalculateReadingOrderNew \
            -input_dir $IMAGES_PATH/page/ \
            -border_margin $REGIONRULEBASEDBORDERMARGIN \
            -threads $REGIONRULEBASEDTHREADS \
            $clean_borders \
            $namespace| tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionRecalculateReadingOrder" $status
fi


if [[ $DETECTLANGUAGE -eq 1 ]]
then
    echo "Detecting language..."
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionDetectLanguageOfPageXml \
            -page $IMAGES_PATH/page/ \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionDetectLanguageOfPageXml" $status
fi


if [[ $SPLITWORDS -eq 1 ]]
then
    echo "MinionSplitPageXMLTextLineIntoWords..."
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionSplitPageXMLTextLineIntoWords \
            -input_path $IMAGES_PATH/page/ \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionSplitPageXMLTextLineIntoWords" $status
fi

if [[ $NEWREADINGORDER -eq 1 ]]
then
    echo "Recalculating the Reading Order of the Regions..."
    docker run $DOCKERGPUPARAMS -u $(id -u ${USER}):$(id -g ${USER}) --rm -it \
        -v $IMAGES_PATH:$IMAGES_PATH/ \
        -v $tmpdir:/tmp \
        -v $(pwd):/workspace \
        reorder-image \
    	bash -c "ls /workspace && python /workspace/reorder/reorder.py $IMAGES_PATH/page --overwrite"
    # Check if failed
    status=$?
    check_error_and_exit "Recalculating the Reading Order of the regions" $status
    echo "Finished Recalculating the Reading Order of the regions."
fi


# cleanup results
rm -rf $tmpdir
    # Check if failed
    status=$?
    check_error_and_exit "MinionCutFromImageBasedOnPageXMLNew" $status

    # Collect all the snippets in a file
    find $tmpdir/imagesnippets/ -type f -name '*.png' > $tmpdir/lines.txt

    echo "Running HTR"
    LOGHIDIR="$(dirname "${HTRLOGHIMODEL}")"

    echo docker run $DOCKERGPUPARAMS -u $(id -u ${USER}):$(id -g ${USER}) --rm -m 32000m --shm-size 10240m -ti \
        -v /tmp:/tmp \
        -v $tmpdir:$tmpdir \
        -v $LOGHIDIR:$LOGHIDIR \
        $DOCKERLOGHIHTR \
            bash -c "LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 python3 /src/loghi-htr/src/main.py \
            --model $HTRLOGHIMODEL  \
            --batch_size 64 \
            --use_mask \
            --inference_list $tmpdir/lines.txt \
            --results_file $tmpdir/results.txt \
            --gpu $GPU \
            --output $tmpdir/output/ \
            --beam_width $BEAMWIDTH " | tee -a $tmpdir/log.txt


    docker run $DOCKERGPUPARAMS -u $(id -u ${USER}):$(id -g ${USER}) --rm -m 32000m --shm-size 10240m -ti \
        -v /tmp:/tmp \
        -v $tmpdir:$tmpdir \
        -v $LOGHIDIR:$LOGHIDIR \
        $DOCKERLOGHIHTR \
            bash -c "LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 python3 /src/loghi-htr/src/main.py \
            --model $HTRLOGHIMODEL  \
            --batch_size 64 \
            --use_mask \
            --inference_list $tmpdir/lines.txt \
            --results_file $tmpdir/results.txt \
            --gpu $GPU \
            --output $tmpdir/output/ \
            --beam_width $BEAMWIDTH " | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "Loghi HTR" $status
    
    echo "Loghi HTR done"

    # Fourth step: merge results back into PageXML
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $LOGHIDIR:$LOGHIDIR \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionLoghiHTRMergePageXML \
            -input_path $IMAGES_PATH/page \
            -results_file $tmpdir/results.txt \
            -config_file $HTRLOGHIMODEL/config.json \
            -htr_code_config_file $tmpdir/output/config.json \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionLoghiHTRMergePageXML" $status
    
    echo "Merging HTR results into PageXML done"
fi

# Fifth step: Recalculate reading order
if [[ $REGIONRULEBASED -eq 1 ]]
then
    echo "Recalculating reading order"
    clean_borders=""

    if [[ $REGIONRULEBASEDCLEANBORDERS -eq 1 ]]; then
        echo "and cleaning borders"
        clean_borders=" -clean_borders "
    fi
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionRecalculateReadingOrderNew \
            -input_dir $IMAGES_PATH/page/ \
            -border_margin $REGIONRULEBASEDBORDERMARGIN \
            -threads $REGIONRULEBASEDTHREADS \
            $clean_borders \
            $namespace| tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionRecalculateReadingOrder" $status
fi


if [[ $DETECTLANGUAGE -eq 1 ]]
then
    echo "Detecting language..."
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionDetectLanguageOfPageXml \
            -page $IMAGES_PATH/page/ \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionDetectLanguageOfPageXml" $status
fi


if [[ $SPLITWORDS -eq 1 ]]
then
    echo "MinionSplitPageXMLTextLineIntoWords..."
    docker run -u $(id -u ${USER}):$(id -g ${USER}) --rm \
        -v $IMAGES_PATH/:$IMAGES_PATH/ \
        -v $tmpdir:$tmpdir \
        $DOCKERLOGHITOOLING /src/loghi-tooling/minions/target/appassembler/bin/MinionSplitPageXMLTextLineIntoWords \
            -input_path $IMAGES_PATH/page/ \
            $namespace | tee -a $tmpdir/log.txt

    # Check if failed
    status=$?
    check_error_and_exit "MinionSplitPageXMLTextLineIntoWords" $status
fi

if [[ $NEWREADINGORDER -eq 1 ]]
then
    echo "Recalculating the Reading Order of the Regions..."
    docker run $DOCKERGPUPARAMS -u $(id -u ${USER}):$(id -g ${USER}) --rm -it \
        -v $IMAGES_PATH:$IMAGES_PATH/ \
        -v $tmpdir:/tmp \
        -v $(pwd):/workspace \
        reorder-image \
    	bash -c "python /workspace/reorder/reorder.py $IMAGES_PATH/page --overwrite"
    # Check if failed
    status=$?
    check_error_and_exit "Recalculating the Reading Order of the regions" $status
    echo "Finished Recalculating the Reading Order of the regions."
fi


# cleanup results
rm -rf $tmpdir