172 lines
5 KiB
Bash
Executable file
172 lines
5 KiB
Bash
Executable file
#!/usr/bin/env bash
|
|
# Download OCR/object models used by the service.
|
|
|
|
set -euo pipefail
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
MODELS_DIR="${MODELS_DIR:-$ROOT_DIR/models}"
|
|
PADDLE_DIR="${PADDLE_DIR:-$MODELS_DIR/paddle}"
|
|
TESSDATA_LOCAL_DIR="${TESSDATA_LOCAL_DIR:-$MODELS_DIR/tessdata_best}"
|
|
|
|
DOWNLOAD_CORAL=1
|
|
DOWNLOAD_EAST=1
|
|
DOWNLOAD_PADDLE=1
|
|
DOWNLOAD_TESSDATA=1
|
|
|
|
usage() {
|
|
cat <<'EOF'
|
|
Usage:
|
|
./scripts/download_models.sh [options]
|
|
|
|
Options:
|
|
--only-coral Download only Coral model + labels
|
|
--only-east Download only EAST text detector model
|
|
--only-paddle Download only PaddleOCR models
|
|
--only-tessdata Download only local tessdata_best (eng+rus)
|
|
--no-coral Skip Coral model download
|
|
--no-east Skip EAST model download
|
|
--no-paddle Skip PaddleOCR model download
|
|
--no-tessdata Skip local tessdata_best download
|
|
-h, --help Show this help
|
|
|
|
Environment overrides:
|
|
MODELS_DIR=/abs/path/to/models
|
|
PADDLE_DIR=/abs/path/to/models/paddle
|
|
TESSDATA_LOCAL_DIR=/abs/path/to/tessdata_best
|
|
EOF
|
|
}
|
|
|
|
while [[ $# -gt 0 ]]; do
|
|
case "$1" in
|
|
--only-coral)
|
|
DOWNLOAD_CORAL=1; DOWNLOAD_EAST=0; DOWNLOAD_PADDLE=0; DOWNLOAD_TESSDATA=0
|
|
shift
|
|
;;
|
|
--only-east)
|
|
DOWNLOAD_CORAL=0; DOWNLOAD_EAST=1; DOWNLOAD_PADDLE=0; DOWNLOAD_TESSDATA=0
|
|
shift
|
|
;;
|
|
--only-paddle)
|
|
DOWNLOAD_CORAL=0; DOWNLOAD_EAST=0; DOWNLOAD_PADDLE=1; DOWNLOAD_TESSDATA=0
|
|
shift
|
|
;;
|
|
--only-tessdata)
|
|
DOWNLOAD_CORAL=0; DOWNLOAD_EAST=0; DOWNLOAD_PADDLE=0; DOWNLOAD_TESSDATA=1
|
|
shift
|
|
;;
|
|
--no-coral)
|
|
DOWNLOAD_CORAL=0
|
|
shift
|
|
;;
|
|
--no-east)
|
|
DOWNLOAD_EAST=0
|
|
shift
|
|
;;
|
|
--no-paddle)
|
|
DOWNLOAD_PADDLE=0
|
|
shift
|
|
;;
|
|
--no-tessdata)
|
|
DOWNLOAD_TESSDATA=0
|
|
shift
|
|
;;
|
|
-h|--help)
|
|
usage
|
|
exit 0
|
|
;;
|
|
*)
|
|
echo "[ERROR] Unknown option: $1" >&2
|
|
usage
|
|
exit 1
|
|
;;
|
|
esac
|
|
done
|
|
|
|
mkdir -p "$MODELS_DIR" "$PADDLE_DIR" "$TESSDATA_LOCAL_DIR"
|
|
|
|
if command -v curl >/dev/null 2>&1; then
|
|
DOWNLOADER="curl"
|
|
elif command -v wget >/dev/null 2>&1; then
|
|
DOWNLOADER="wget"
|
|
else
|
|
echo "[ERROR] curl or wget is required" >&2
|
|
exit 1
|
|
fi
|
|
|
|
download_file() {
|
|
local url="$1"
|
|
local dst="$2"
|
|
mkdir -p "$(dirname "$dst")"
|
|
if [[ -s "$dst" ]]; then
|
|
echo "[SKIP] $dst already exists"
|
|
return 0
|
|
fi
|
|
local tmp="${dst}.part"
|
|
echo "[GET] $url"
|
|
if [[ "$DOWNLOADER" == "curl" ]]; then
|
|
curl -L --fail --retry 3 --connect-timeout 20 --max-time 600 "$url" -o "$tmp"
|
|
else
|
|
wget -O "$tmp" "$url"
|
|
fi
|
|
mv "$tmp" "$dst"
|
|
}
|
|
|
|
download_and_extract_tar() {
|
|
local url="$1"
|
|
local tag="$2"
|
|
local out_dir="$3"
|
|
local cache_dir="$MODELS_DIR/.downloads"
|
|
local archive="$cache_dir/${tag}.tar"
|
|
mkdir -p "$cache_dir" "$out_dir"
|
|
download_file "$url" "$archive"
|
|
echo "[EXTRACT] $archive -> $out_dir"
|
|
tar -xf "$archive" -C "$out_dir"
|
|
}
|
|
|
|
# ---- URLs ----
|
|
CORAL_MODEL_URL="https://github.com/google-coral/test_data/raw/master/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite"
|
|
CORAL_LABELS_URL="https://github.com/google-coral/test_data/raw/master/coco_labels.txt"
|
|
EAST_URL="https://github.com/argman/EAST/releases/download/model/frozen_east_text_detection.pb"
|
|
TESS_ENG_URL="https://github.com/tesseract-ocr/tessdata_best/raw/main/eng.traineddata"
|
|
TESS_RUS_URL="https://github.com/tesseract-ocr/tessdata_best/raw/main/rus.traineddata"
|
|
|
|
# PaddleOCR v5 mobile models (detection + RU/EN recognition)
|
|
PADDLE_DET_URL="https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/PP-OCRv5_mobile_det_infer.tar"
|
|
PADDLE_REC_URL="https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/eslav_PP-OCRv5_mobile_rec_infer.tar"
|
|
|
|
if [[ "$DOWNLOAD_CORAL" == "1" ]]; then
|
|
download_file "$CORAL_MODEL_URL" "$MODELS_DIR/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite"
|
|
download_file "$CORAL_LABELS_URL" "$MODELS_DIR/coco_labels.txt"
|
|
fi
|
|
|
|
if [[ "$DOWNLOAD_EAST" == "1" ]]; then
|
|
download_file "$EAST_URL" "$MODELS_DIR/frozen_east_text_detection.pb"
|
|
fi
|
|
|
|
if [[ "$DOWNLOAD_PADDLE" == "1" ]]; then
|
|
download_and_extract_tar "$PADDLE_DET_URL" "PP-OCRv5_mobile_det_infer" "$PADDLE_DIR"
|
|
download_and_extract_tar "$PADDLE_REC_URL" "eslav_PP-OCRv5_mobile_rec_infer" "$PADDLE_DIR"
|
|
fi
|
|
|
|
if [[ "$DOWNLOAD_TESSDATA" == "1" ]]; then
|
|
download_file "$TESS_ENG_URL" "$TESSDATA_LOCAL_DIR/eng.traineddata"
|
|
download_file "$TESS_RUS_URL" "$TESSDATA_LOCAL_DIR/rus.traineddata"
|
|
fi
|
|
|
|
cat <<EOF
|
|
|
|
[DONE] Model download complete.
|
|
|
|
Recommended env:
|
|
export MODEL_EDGETPU_PATH="$MODELS_DIR/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite"
|
|
export MODEL_LABELS_PATH="$MODELS_DIR/coco_labels.txt"
|
|
export MODEL_EAST_PATH="$MODELS_DIR/frozen_east_text_detection.pb"
|
|
export TESSDATA_PREFIX="$TESSDATA_LOCAL_DIR"
|
|
export OCR_BACKEND="tesseract"
|
|
|
|
For PaddleOCR backend:
|
|
export OCR_BACKEND="paddle"
|
|
export PADDLE_LANG="ru"
|
|
export PADDLE_DET_MODEL_DIR="$PADDLE_DIR/PP-OCRv5_mobile_det_infer"
|
|
export PADDLE_REC_MODEL_DIR="$PADDLE_DIR/eslav_PP-OCRv5_mobile_rec_infer"
|
|
EOF
|