#!/usr/bin/env bash
# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Download and convert PaddleOCR PP-OCRv4 detection and recognition models
# to OpenVINO IR for the ocr-text-recognition use case.
# Usage: ./export_and_quantize.sh

set -euo pipefail

echo "--- Installing dependencies ---"
pip install -qU openvino

# Ask for approval before downloading models and sample files
echo ""
echo "This script will download:"
echo "  - Model weights and/or sample files"
echo ""
read -p "Continue with downloads? (yes/no): " APPROVAL
if [[ "${APPROVAL}" != "yes" ]]; then
    echo "Download cancelled by user."
    exit 0
fi
echo ""

DET_URL="https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_det_infer.tar"
# The larger "server" recognition model is noticeably more accurate than the
# mobile variant on stylized/decorative fonts (e.g. the sample video's plate
# text), at the cost of a bigger download and slightly slower inference.
REC_URL="https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_rec_server_infer.tar"

echo "--- Downloading PP-OCRv4 detection model ---"
if [[ ! -d "ch_PP-OCRv4_det_infer" ]]; then
    wget -q -O det.tar "${DET_URL}"
    tar xf det.tar
    rm -f det.tar
    echo "Downloaded and extracted: ch_PP-OCRv4_det_infer/"
else
    echo "Already present: ch_PP-OCRv4_det_infer/"
fi

echo "--- Converting detection model to OpenVINO IR ---"
if [[ ! -f "ch_PP-OCRv4_det_infer/inference.xml" ]]; then
    ovc ch_PP-OCRv4_det_infer/inference.pdmodel \
        --output_model ch_PP-OCRv4_det_infer/inference.xml
    echo "Converted detection model to OpenVINO IR"
else
    echo "Already converted: ch_PP-OCRv4_det_infer/inference.xml"
fi

echo "--- Downloading PP-OCRv4 recognition model (server variant) ---"
if [[ ! -d "ch_PP-OCRv4_rec_server_infer" ]]; then
    wget -q -O rec.tar "${REC_URL}"
    tar xf rec.tar
    rm -f rec.tar
    echo "Downloaded and extracted: ch_PP-OCRv4_rec_server_infer/"
else
    echo "Already present: ch_PP-OCRv4_rec_server_infer/"
fi

echo "--- Converting recognition model to OpenVINO IR ---"
if [[ ! -f "ch_PP-OCRv4_rec_server_infer/inference.xml" ]]; then
    ovc ch_PP-OCRv4_rec_server_infer/inference.pdmodel \
        --output_model ch_PP-OCRv4_rec_server_infer/inference.xml
    echo "Converted recognition model to OpenVINO IR"
else
    echo "Already converted: ch_PP-OCRv4_rec_server_infer/inference.xml"
fi

echo "--- Downloading sample test image ---"
if [[ ! -f test_ocr.jpg ]]; then
    wget -q -O test_ocr.jpg \
        "https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.7/doc/imgs_en/img_12.jpg"
    echo "Downloaded: test_ocr.jpg"
else
    echo "Already present: test_ocr.jpg"
fi

echo "--- Downloading PP-OCRv4 character dictionary ---"
if [[ ! -f ppocr_keys_v1.txt ]]; then
    wget -q -O ppocr_keys_v1.txt \
        "https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.7/ppocr/utils/ppocr_keys_v1.txt"
    echo "Downloaded: ppocr_keys_v1.txt"
else
    echo "Already present: ppocr_keys_v1.txt"
fi

echo "--- Downloading sample test video ---"
if [[ ! -f test_video.mp4 ]]; then
    wget -q -O test_video.mp4 \
        "https://www.pexels.com/download/video/5286217/"
    echo "Downloaded: test_video.mp4"
else
    echo "Already present: test_video.mp4"
fi

echo "--- Done ---"
echo "Detection : ch_PP-OCRv4_det_infer/inference.xml"
echo "Recognition: ch_PP-OCRv4_rec_server_infer/inference.xml"
echo "Dictionary : ppocr_keys_v1.txt"
echo "Sample     : test_ocr.jpg"
echo "Video      : test_video.mp4"