text_detect/scripts/download_models.sh
2026-03-08 14:27:38 +03:00

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 v3 mobile models (good compatibility with paddleocr 2.7.x / Python 3.6)
PADDLE_DET_URL="https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/PP-OCRv3_mobile_det_infer.tar"
PADDLE_REC_URL="https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/cyrillic_PP-OCRv3_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-OCRv3_mobile_det_infer" "$PADDLE_DIR"
download_and_extract_tar "$PADDLE_REC_URL" "cyrillic_PP-OCRv3_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-OCRv3_mobile_det_infer"
export PADDLE_REC_MODEL_DIR="$PADDLE_DIR/cyrillic_PP-OCRv3_mobile_rec_infer"
EOF