diff --git a/README.md b/README.md index df235dd..e423342 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,8 @@ Python 3.6-compatible Flask service that accepts an image and returns detected t - Accepts JPEG/PNG/WebP/BMP/TIFF/GIF (first frame) via multipart or base64 JSON. - Optional Google Coral EdgeTPU detection for scene context (loaded once at start). - EAST text detector (OpenCV DNN) when model exists; otherwise OpenCV contour-based text-region fallback. -- OCR via `pytesseract` with per-block confidences; configurable languages. +- OCR via `pytesseract` (default) with per-block confidences; configurable languages. +- Optional `PaddleOCR` backend (`OCR_BACKEND=paddle`) for better text-on-photo quality. - Auto-selects `rus` or `eng` (via Tesseract OSD script detection) before full OCR when `OCR_LANG=eng+rus`. - If OCR result is empty, retries with fallback language/full-frame pass to improve recall. - Hard payload cap (10MB default) and graceful degradation if Coral/EAST unavailable. @@ -32,12 +33,28 @@ sudo apt-get install -y python3-opencv pip3 install -r requirements.txt ``` +Optional PaddleOCR backend (Python 3.8+ only): +```bash +pip3 install -r requirements-paddle.txt +``` + ## Models Place your models in `./models` (paths can be overridden via env): - EdgeTPU SSD: `ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite` - Labels: `coco_labels.txt` - EAST text detector (optional): `frozen_east_text_detection.pb` +Auto-download all models (Coral + EAST + Paddle + local `tessdata_best`): +```bash +./scripts/download_models.sh +``` + +Examples: +```bash +./scripts/download_models.sh --only-paddle +./scripts/download_models.sh --no-paddle +``` + ### Better OCR models (tessdata_best) ```bash sudo apt-get install -y wget @@ -56,9 +73,17 @@ OpenCV from JetPack repos (4.x) usually ships without CUDA dnn. To use CUDA, bui - `MODEL_EDGETPU_PATH` (default `./models/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite`) - `MODEL_LABELS_PATH` (default `./models/coco_labels.txt`) - `MODEL_EAST_PATH` (default `./models/frozen_east_text_detection.pb`) +- `OCR_BACKEND` (default `tesseract`; `tesseract|paddle|auto`) - `OCR_LANG` (default `eng+rus`; request language set for OCR) - `OCR_AUTO_LANG` (default `true`; auto-pick `eng` or `rus` by script when `OCR_LANG` contains both) - `OCR_CONFIG` (default `--oem 1 --psm 6`, forwarded to Tesseract) +- `PADDLE_LANG` (default `ru`; used when `OCR_BACKEND=paddle`) +- `PADDLE_USE_GPU` (default `false`) +- `PADDLE_USE_ANGLE_CLS` (default `false`) +- `PADDLE_DET_MODEL_DIR` (default `./models/paddle/PP-OCRv5_mobile_det_infer`) +- `PADDLE_REC_MODEL_DIR` (default `./models/paddle/eslav_PP-OCRv5_mobile_rec_infer`) +- `PADDLE_CLS_MODEL_DIR` (default empty) +- `PADDLE_CPU_THREADS` (default `4`) - `TESSDATA_PREFIX` (path to tessdata / tessdata_best if you install better models) - `OCR_UPSCALE` (default `1.5`, upscale factor before OCR; max 3.0) - `OCR_CLAHE` (default `true`, contrast-limited adaptive histogram) @@ -80,7 +105,7 @@ python3 app.py --host 0.0.0.0 --port 8000 Waitress or gunicorn can also serve the app if desired (not required). ## API -- `GET /health` -> `{ "status": "ok", "coral": true/false }` +- `GET /health` -> `{ "status": "ok", "coral": true/false, "paddle": true/false, "ocr_backend": "..." }` - `POST /v1/image2text` - Multipart: field `image` - JSON: `{ "image_base64": "" }` @@ -108,6 +133,7 @@ Example JSON response (fields may vary): "meta": { "latency_ms": 123, "coral_used": true, + "ocr_backend": "tesseract", "east_used": true, "text_locator": "east", "ocr_lang": "eng", @@ -121,6 +147,7 @@ Example JSON response (fields may vary): ## Notes - If Coral is absent or fails, the service still returns OCR with `coral_used=false`. +- If `OCR_BACKEND=paddle` but PaddleOCR package/models are unavailable, service falls back to Tesseract. - If EAST model is missing, service first tries OpenCV contour text regions and only then full-image OCR. - Keep images under `MAX_IMAGE_MB` (413 returned otherwise). - Optimize Tesseract languages to those you need to speed up OCR. diff --git a/__pycache__/app.cpython-313.pyc b/__pycache__/app.cpython-313.pyc index bf2035c..7969f80 100644 Binary files a/__pycache__/app.cpython-313.pyc and b/__pycache__/app.cpython-313.pyc differ diff --git a/app.py b/app.py index e160a8f..78c49b0 100644 --- a/app.py +++ b/app.py @@ -14,7 +14,7 @@ import uuid import threading from flask import Flask, request, jsonify -from core import image_io, coral_vision, text_detect, ocr, schema +from core import image_io, coral_vision, text_detect, ocr, schema, paddle_ocr_backend try: import cv2 @@ -48,6 +48,7 @@ def env_int(name, default): MODEL_EDGETPU_PATH = os.getenv("MODEL_EDGETPU_PATH", "./models/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite") MODEL_LABELS_PATH = os.getenv("MODEL_LABELS_PATH", "./models/coco_labels.txt") MODEL_EAST_PATH = os.getenv("MODEL_EAST_PATH", "./models/frozen_east_text_detection.pb") +OCR_BACKEND = (os.getenv("OCR_BACKEND", "tesseract") or "tesseract").strip().lower() OCR_LANG = os.getenv("OCR_LANG", "eng+rus") MAX_IMAGE_MB = env_int("MAX_IMAGE_MB", 10) CONF_THRESHOLD = env_float("CONF_THRESHOLD", 0.4) @@ -65,8 +66,19 @@ OCR_BLUR = env_int("OCR_BLUR", 0) # 0 disables gaussian blur OCR_INVERT = env_bool("OCR_INVERT", False) OCR_USE_THRESHOLD = env_bool("OCR_USE_THRESHOLD", True) OCR_AUTO_LANG = env_bool("OCR_AUTO_LANG", True) +PADDLE_LANG = os.getenv("PADDLE_LANG", "ru") +PADDLE_USE_GPU = env_bool("PADDLE_USE_GPU", False) +PADDLE_USE_ANGLE_CLS = env_bool("PADDLE_USE_ANGLE_CLS", False) +PADDLE_DET_MODEL_DIR = os.getenv("PADDLE_DET_MODEL_DIR", "") +PADDLE_REC_MODEL_DIR = os.getenv("PADDLE_REC_MODEL_DIR", "") +PADDLE_CLS_MODEL_DIR = os.getenv("PADDLE_CLS_MODEL_DIR", "") +PADDLE_CPU_THREADS = env_int("PADDLE_CPU_THREADS", 0) MAX_CONTENT_LENGTH = MAX_IMAGE_MB * 1024 * 1024 +if OCR_BACKEND not in ("tesseract", "paddle", "auto"): + print('[WARN] Unsupported OCR_BACKEND "{}"; fallback to "tesseract"'.format(OCR_BACKEND), file=sys.stderr) + OCR_BACKEND = "tesseract" + app = Flask(__name__) app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH @@ -145,6 +157,17 @@ txt_detector = text_detect.TextDetector( use_cuda=USE_CUDA_DNN, ) +paddle_backend = paddle_ocr_backend.PaddleOCRBackend( + enable=(OCR_BACKEND in ("paddle", "auto")), + lang=PADDLE_LANG, + use_angle_cls=PADDLE_USE_ANGLE_CLS, + use_gpu=PADDLE_USE_GPU, + det_model_dir=PADDLE_DET_MODEL_DIR, + rec_model_dir=PADDLE_REC_MODEL_DIR, + cls_model_dir=PADDLE_CLS_MODEL_DIR, + cpu_threads=PADDLE_CPU_THREADS, +) + # ---- Helpers ---- def parse_image_from_request(req): @@ -323,7 +346,12 @@ def build_text_gui_image(img_rgb, boxes, blocks, target_width=1600): @app.route('/health', methods=['GET']) def health(): - return jsonify({"status": "ok", "coral": coral.available}) + return jsonify({ + "status": "ok", + "coral": coral.available, + "paddle": paddle_backend.available, + "ocr_backend": OCR_BACKEND, + }) @app.route('/v1/image2text', methods=['POST']) @@ -362,82 +390,90 @@ def image2text(): objects = [] coral_used = False - boxes, east_used = txt_detector.detect(img_rgb) h, w = img_rgb.shape[:2] - text_locator = "east" if east_used else "opencv_fallback" - if not east_used and len(boxes) == 1: - try: - if tuple([int(v) for v in boxes[0]]) == (0, 0, w, h): - text_locator = "full_image" - except Exception: - pass - - prep_options = { - "upscale": OCR_UPSCALE, - "clahe": OCR_CLAHE, - "bilateral": OCR_BILATERAL, - "blur": OCR_BLUR, - "invert": OCR_INVERT, - "use_threshold": OCR_USE_THRESHOLD, - } - + ocr_backend_used = "tesseract" ocr_lang_selected = OCR_LANG ocr_lang_strategy = "fixed" - if OCR_AUTO_LANG: - try: - ocr_lang_selected, ocr_lang_strategy = ocr.choose_ocr_lang( - img_rgb, - boxes, - requested_lang=OCR_LANG, - prep_options=prep_options, - ) - except Exception as exc: - print("[WARN] OCR lang auto-detect failed {}: {}".format(request_id, exc), file=sys.stderr) - ocr_lang_selected = OCR_LANG - ocr_lang_strategy = "auto_error_fallback" - - text, blocks = ocr.run_ocr( - img_rgb, - boxes, - lang=ocr_lang_selected, - config=OCR_CONFIG, - return_boxes=RETURN_BOXES, - prep_options=prep_options, - ) ocr_passes = 1 + text = "" + blocks = [] + boxes = [] + east_used = False + text_locator = "full_image" - # Retry with requested language set when auto-picked language produced no text. - if (not (text or '').strip()) and ocr_lang_selected != OCR_LANG: - retry_text, retry_blocks = ocr.run_ocr( + use_paddle = (OCR_BACKEND in ("paddle", "auto")) and paddle_backend.available + if use_paddle: + ocr_backend_used = "paddle" + text, blocks = paddle_backend.run_ocr(img_rgb, return_boxes=RETURN_BOXES) + ocr_lang_selected = PADDLE_LANG + ocr_lang_strategy = "paddle" + if blocks: + paddle_boxes = [] + for blk in blocks: + bbox = blk.get("bbox", [0, 0, 0, 0]) + if len(bbox) != 4: + continue + paddle_boxes.append((int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))) + boxes = paddle_boxes or [(0, 0, w, h)] + text_locator = "paddle_det" + else: + boxes = [(0, 0, w, h)] + text_locator = "paddle_full_image" + else: + fallback_prefix = "" + if OCR_BACKEND == "paddle": + fallback_prefix = "paddle_unavailable_fallback" + ocr_lang_strategy = fallback_prefix + boxes, east_used = txt_detector.detect(img_rgb) + text_locator = "east" if east_used else "opencv_fallback" + if not east_used and len(boxes) == 1: + try: + if tuple([int(v) for v in boxes[0]]) == (0, 0, w, h): + text_locator = "full_image" + except Exception: + pass + + prep_options = { + "upscale": OCR_UPSCALE, + "clahe": OCR_CLAHE, + "bilateral": OCR_BILATERAL, + "blur": OCR_BLUR, + "invert": OCR_INVERT, + "use_threshold": OCR_USE_THRESHOLD, + } + + if OCR_AUTO_LANG: + try: + ocr_lang_selected, ocr_lang_strategy = ocr.choose_ocr_lang( + img_rgb, + boxes, + requested_lang=OCR_LANG, + prep_options=prep_options, + ) + except Exception as exc: + print("[WARN] OCR lang auto-detect failed {}: {}".format(request_id, exc), file=sys.stderr) + ocr_lang_selected = OCR_LANG + ocr_lang_strategy = "auto_error_fallback" + + if fallback_prefix and ocr_lang_strategy != fallback_prefix: + ocr_lang_strategy = "{}+{}".format(fallback_prefix, ocr_lang_strategy) + + text, blocks = ocr.run_ocr( img_rgb, boxes, - lang=OCR_LANG, + lang=ocr_lang_selected, config=OCR_CONFIG, return_boxes=RETURN_BOXES, prep_options=prep_options, ) - ocr_passes += 1 - if (retry_text or '').strip(): - text, blocks = retry_text, retry_blocks - ocr_lang_selected = OCR_LANG - ocr_lang_strategy = "{}+lang_retry".format(ocr_lang_strategy) + ocr_passes = 1 - # Retry with full-frame OCR if EAST boxes were too aggressive and OCR is empty. - if not (text or '').strip(): - h, w = img_rgb.shape[:2] - full_box = (0, 0, w, h) - needs_full_pass = True - if len(boxes) == 1: - try: - cur_box = tuple([int(v) for v in boxes[0]]) - needs_full_pass = cur_box != full_box - except Exception: - needs_full_pass = True - if needs_full_pass: + # Retry with requested language set when auto-picked language produced no text. + if (not (text or '').strip()) and ocr_lang_selected != OCR_LANG: retry_text, retry_blocks = ocr.run_ocr( img_rgb, - [full_box], - lang=ocr_lang_selected, + boxes, + lang=OCR_LANG, config=OCR_CONFIG, return_boxes=RETURN_BOXES, prep_options=prep_options, @@ -445,7 +481,32 @@ def image2text(): ocr_passes += 1 if (retry_text or '').strip(): text, blocks = retry_text, retry_blocks - ocr_lang_strategy = "{}+full_frame_retry".format(ocr_lang_strategy) + ocr_lang_selected = OCR_LANG + ocr_lang_strategy = "{}+lang_retry".format(ocr_lang_strategy) + + # Retry with full-frame OCR if EAST boxes were too aggressive and OCR is empty. + if not (text or '').strip(): + full_box = (0, 0, w, h) + needs_full_pass = True + if len(boxes) == 1: + try: + cur_box = tuple([int(v) for v in boxes[0]]) + needs_full_pass = cur_box != full_box + except Exception: + needs_full_pass = True + if needs_full_pass: + retry_text, retry_blocks = ocr.run_ocr( + img_rgb, + [full_box], + lang=ocr_lang_selected, + config=OCR_CONFIG, + return_boxes=RETURN_BOXES, + prep_options=prep_options, + ) + ocr_passes += 1 + if (retry_text or '').strip(): + text, blocks = retry_text, retry_blocks + ocr_lang_strategy = "{}+full_frame_retry".format(ocr_lang_strategy) scene_summary = schema.summarize_scene(objects) @@ -489,6 +550,7 @@ def image2text(): metadata={ "latency_ms": latency_ms, "coral_used": coral_used, + "ocr_backend": ocr_backend_used, "east_used": east_used, "text_locator": text_locator, "ocr_lang": ocr_lang_selected, @@ -504,10 +566,12 @@ def image2text(): if debug_image_b64: resp["debug_image_base64"] = debug_image_b64 - print("[INFO] req={} size={}KB coral={} east={} latency={}ms".format( + print("[INFO] req={} size={}KB coral={} ocr_backend={} locator={} east={} latency={}ms".format( request_id, len(data) // 1024, coral_used, + ocr_backend_used, + text_locator, east_used, latency_ms )) diff --git a/core/__pycache__/paddle_ocr_backend.cpython-313.pyc b/core/__pycache__/paddle_ocr_backend.cpython-313.pyc new file mode 100644 index 0000000..bed4e5f Binary files /dev/null and b/core/__pycache__/paddle_ocr_backend.cpython-313.pyc differ diff --git a/core/paddle_ocr_backend.py b/core/paddle_ocr_backend.py new file mode 100644 index 0000000..998edb2 --- /dev/null +++ b/core/paddle_ocr_backend.py @@ -0,0 +1,171 @@ +# -*- coding: utf-8 -*- +"""Optional PaddleOCR backend wrapper.""" +import os +import sys + +try: + from paddleocr import PaddleOCR +except ImportError: # pragma: no cover + PaddleOCR = None + + +class PaddleOCRBackend(object): + def __init__(self, + enable=False, + lang='ru', + use_angle_cls=False, + use_gpu=False, + det_model_dir='', + rec_model_dir='', + cls_model_dir='', + cpu_threads=0): + self.enable = bool(enable) + self.lang = lang or 'ru' + self.use_angle_cls = bool(use_angle_cls) + self.use_gpu = bool(use_gpu) + self.det_model_dir = det_model_dir or '' + self.rec_model_dir = rec_model_dir or '' + self.cls_model_dir = cls_model_dir or '' + self.cpu_threads = int(cpu_threads) if cpu_threads else 0 + + self.available = False + self.engine = None + self.last_error = '' + self._load() + + def _load(self): + if not self.enable: + return + if PaddleOCR is None: + self.last_error = 'paddleocr_not_installed' + print('[WARN] PaddleOCR backend requested but package is not installed', file=sys.stderr) + return + kwargs = { + 'lang': self.lang, + 'use_angle_cls': self.use_angle_cls, + 'use_gpu': self.use_gpu, + 'show_log': False, + } + det_dir = self._resolve_dir(self.det_model_dir, 'det') + rec_dir = self._resolve_dir(self.rec_model_dir, 'rec') + cls_dir = self._resolve_dir(self.cls_model_dir, 'cls') + if det_dir: + kwargs['det_model_dir'] = det_dir + if rec_dir: + kwargs['rec_model_dir'] = rec_dir + if cls_dir: + kwargs['cls_model_dir'] = cls_dir + if self.cpu_threads > 0: + kwargs['cpu_threads'] = self.cpu_threads + try: + self.engine = PaddleOCR(**kwargs) + self.available = True + print('[INFO] PaddleOCR backend ready (lang={})'.format(self.lang)) + except Exception as exc: + self.last_error = str(exc) + self.available = False + print('[WARN] Failed to initialize PaddleOCR backend: {}'.format(exc), file=sys.stderr) + + def run_ocr(self, img_rgb, return_boxes=True): + if not self.available or img_rgb is None: + return '', [] + try: + if img_rgb.ndim == 3 and img_rgb.shape[2] == 3: + img_bgr = img_rgb[:, :, ::-1].copy() + else: + img_bgr = img_rgb.copy() + raw = self.engine.ocr(img_bgr, cls=self.use_angle_cls) + except Exception as exc: + print('[WARN] PaddleOCR inference failed: {}'.format(exc), file=sys.stderr) + return '', [] + return self._parse_result(raw, return_boxes=return_boxes) + + def _parse_result(self, raw, return_boxes=True): + lines = [] + self._walk_lines(raw, lines) + blocks = [] + texts = [] + for line in lines: + poly = line[0] + rec = line[1] + text = self._safe_text(rec[0] if len(rec) > 0 else '') + if not text: + continue + conf = self._safe_float(rec[1] if len(rec) > 1 else -1.0, default=-1.0) + bbox = self._poly_to_bbox(poly) + texts.append(text) + if return_boxes: + blocks.append({ + 'bbox': [bbox[0], bbox[1], bbox[2], bbox[3]], + 'text': text, + 'conf': conf, + }) + return '\n'.join(texts), blocks if return_boxes else [] + + def _walk_lines(self, node, out): + if self._is_line_entry(node): + out.append(node) + return + if isinstance(node, (list, tuple)): + for item in node: + self._walk_lines(item, out) + + def _is_line_entry(self, node): + if not isinstance(node, (list, tuple)) or len(node) < 2: + return False + poly = node[0] + rec = node[1] + return self._is_poly(poly) and self._is_rec(rec) + + def _is_poly(self, poly): + if not isinstance(poly, (list, tuple)) or len(poly) < 3: + return False + first = poly[0] + return isinstance(first, (list, tuple)) and len(first) >= 2 + + def _is_rec(self, rec): + if not isinstance(rec, (list, tuple)) or len(rec) < 1: + return False + return isinstance(rec[0], (str, bytes)) + + def _poly_to_bbox(self, poly): + xs = [] + ys = [] + for pt in poly: + if not isinstance(pt, (list, tuple)) or len(pt) < 2: + continue + xs.append(self._safe_int(pt[0], 0)) + ys.append(self._safe_int(pt[1], 0)) + if not xs or not ys: + return (0, 0, 0, 0) + return (max(0, min(xs)), max(0, min(ys)), max(xs), max(ys)) + + def _resolve_dir(self, path, name): + if not path: + return '' + if os.path.isdir(path): + return path + print('[WARN] PaddleOCR {} model dir not found: {}'.format(name, path), file=sys.stderr) + return '' + + def _safe_text(self, value): + if value is None: + return '' + if isinstance(value, bytes): + try: + value = value.decode('utf-8', 'ignore') + except Exception: + return '' + return str(value).strip() + + def _safe_float(self, value, default=0.0): + try: + return float(value) + except Exception: + return float(default) + + def _safe_int(self, value, default=0): + try: + return int(round(float(value))) + except Exception: + return int(default) diff --git a/requirements-paddle.txt b/requirements-paddle.txt new file mode 100644 index 0000000..d885cd6 --- /dev/null +++ b/requirements-paddle.txt @@ -0,0 +1,3 @@ +# Optional OCR backend dependencies. +# PaddleOCR 3.x supports Python 3.8+. +paddleocr>=3.0.0 diff --git a/run_service.sh b/run_service.sh index af308cc..6af84c3 100755 --- a/run_service.sh +++ b/run_service.sh @@ -8,9 +8,17 @@ set -euo pipefail export MODEL_EDGETPU_PATH=${MODEL_EDGETPU_PATH:-./models/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite} export MODEL_LABELS_PATH=${MODEL_LABELS_PATH:-./models/coco_labels.txt} export MODEL_EAST_PATH=${MODEL_EAST_PATH:-./models/frozen_east_text_detection.pb} +export OCR_BACKEND=${OCR_BACKEND:-tesseract} # tesseract|paddle|auto export OCR_LANG=${OCR_LANG:-eng+rus} export OCR_AUTO_LANG=${OCR_AUTO_LANG:-true} export OCR_CONFIG=${OCR_CONFIG:---oem 1 --psm 6} +export PADDLE_LANG=${PADDLE_LANG:-ru} +export PADDLE_USE_GPU=${PADDLE_USE_GPU:-false} +export PADDLE_USE_ANGLE_CLS=${PADDLE_USE_ANGLE_CLS:-false} +export PADDLE_DET_MODEL_DIR=${PADDLE_DET_MODEL_DIR:-./models/paddle/PP-OCRv5_mobile_det_infer} +export PADDLE_REC_MODEL_DIR=${PADDLE_REC_MODEL_DIR:-./models/paddle/eslav_PP-OCRv5_mobile_rec_infer} +export PADDLE_CLS_MODEL_DIR=${PADDLE_CLS_MODEL_DIR:-} +export PADDLE_CPU_THREADS=${PADDLE_CPU_THREADS:-4} # Point to tessdata_best if installed export TESSDATA_PREFIX=${TESSDATA_PREFIX:-} export MAX_IMAGE_MB=${MAX_IMAGE_MB:-10} diff --git a/scripts/download_models.sh b/scripts/download_models.sh new file mode 100755 index 0000000..02ca6a8 --- /dev/null +++ b/scripts/download_models.sh @@ -0,0 +1,172 @@ +#!/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 <