test.py

"""
Awiros-ANPR-OCR single-image / directory inference script.

Usage:
    pip install -r requirements.txt
    python test.py --image_path plate.jpg
    python test.py --image_path plates_dir/ --output_json results.json

PaddleOCR repo is needed for model construction. On first run the script
auto-clones it into a PaddleOCR/ subfolder next to this file.
Pass --paddleocr_dir to point to an existing clone instead.
"""

import argparse
import copy
import json
import os
import subprocess
import sys
from pathlib import Path

import cv2
import numpy as np

_SCRIPT_DIR = Path(__file__).resolve().parent

# ---------------------------------------------------------------------------
# Model architecture config (PP-OCRv5 server rec, SVTR_HGNet)
# CTC head output: 64 classes (63 dict chars + blank)
# NRTR head output: 68 classes (64 + bos/eos/pad/unk)
# ---------------------------------------------------------------------------
CTC_NUM_CLASSES = 64
NRTR_NUM_CLASSES = 67  # NRTRHead internally adds +1, so 67 -> 68 to match weights

MODEL_CONFIG = {
    "Architecture": {
        "model_type": "rec",
        "algorithm": "SVTR_HGNet",
        "Transform": None,
        "Backbone": {"name": "PPHGNetV2_B4", "text_rec": True},
        "Head": {
            "name": "MultiHead",
            "out_channels_list": {
                "CTCLabelDecode": CTC_NUM_CLASSES,
                "NRTRLabelDecode": NRTR_NUM_CLASSES,
            },
            "head_list": [
                {
                    "CTCHead": {
                        "Neck": {
                            "name": "svtr",
                            "dims": 120,
                            "depth": 2,
                            "hidden_dims": 120,
                            "kernel_size": [1, 3],
                            "use_guide": True,
                        },
                        "Head": {"fc_decay": 1e-05},
                    }
                },
                {"NRTRHead": {"nrtr_dim": 384, "max_text_length": 25}},
            ],
        },
    },
}

IMAGE_SHAPE = [3, 48, 320]
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp"}


# ---------------------------------------------------------------------------
# PaddleOCR path setup
# ---------------------------------------------------------------------------
def _find_paddleocr(explicit_path=None):
    """Find a directory containing the ppocr package."""
    candidates = []
    if explicit_path:
        candidates.append(Path(explicit_path))
    candidates += [
        _SCRIPT_DIR / "PaddleOCR",
        _SCRIPT_DIR,
        Path.cwd(),
        Path.cwd() / "PaddleOCR",
    ]
    for c in candidates:
        if (c / "ppocr" / "__init__.py").is_file():
            return c
    return None


def _ensure_paddleocr(explicit_path=None):
    """Make ppocr importable. Auto-clones PaddleOCR if not found."""
    root = _find_paddleocr(explicit_path)
    if root is None:
        clone_target = _SCRIPT_DIR / "PaddleOCR"
        print(f"ppocr not found. Cloning PaddleOCR into {clone_target} ...")
        subprocess.check_call([
            "git", "clone", "--depth", "1",
            "https://github.com/PaddlePaddle/PaddleOCR.git",
            str(clone_target),
        ])
        root = clone_target
    root_str = str(root)
    if root_str not in sys.path:
        sys.path.insert(0, root_str)
    return root


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def parse_args():
    p = argparse.ArgumentParser("Awiros-ANPR-OCR inference")
    p.add_argument("--image_path", required=True,
                    help="Path to a single image or a directory of images.")
    p.add_argument("--weights", default="",
                    help="Path to model.safetensors (default: next to this script).")
    p.add_argument("--dict_path", default="",
                    help="Path to en_dict.txt (default: next to this script).")
    p.add_argument("--device", default="gpu", choices=["gpu", "cpu"],
                    help="Device for inference.")
    p.add_argument("--output_json", default="",
                    help="Optional output JSON path for results.")
    p.add_argument("--paddleocr_dir", default="",
                    help="Path to PaddleOCR repo root (auto-cloned if omitted).")
    return p.parse_args()


def resolve_path(user_path: str, filename: str) -> str:
    """Use user-supplied path if it exists, else fall back to script dir."""
    if user_path and os.path.exists(user_path):
        return user_path
    alt = _SCRIPT_DIR / filename
    if alt.exists():
        return str(alt)
    raise FileNotFoundError(
        f"Could not find {filename}. Place it next to this script or pass its path."
    )


def load_safetensors_to_paddle(paddle_mod, weight_path: str):
    from safetensors.numpy import load_file
    np_state = load_file(weight_path)
    return {k: paddle_mod.to_tensor(v) for k, v in np_state.items()}


def resize_for_rec(img_bgr, target_shape):
    _, h, w = target_shape
    img_h, img_w = img_bgr.shape[:2]
    ratio = h / img_h
    new_w = min(int(img_w * ratio), w)
    resized = cv2.resize(img_bgr, (new_w, h))
    if new_w < w:
        padded = np.zeros((h, w, 3), dtype=np.uint8)
        padded[:, :new_w, :] = resized
        resized = padded
    return resized


def preprocess(img_bgr, target_shape):
    img = resize_for_rec(img_bgr, target_shape)
    img = img.astype(np.float32) / 255.0
    img = (img - 0.5) / 0.5
    return img.transpose((2, 0, 1))


def collect_images(path: str):
    p = Path(path)
    if p.is_file():
        return [p]
    if p.is_dir():
        return sorted(f for f in p.iterdir()
                       if f.is_file() and f.suffix.lower() in IMAGE_EXTENSIONS)
    raise FileNotFoundError(f"Path not found: {path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    args = parse_args()

    # 1. Ensure ppocr is importable, then import paddle + ppocr
    _ensure_paddleocr(args.paddleocr_dir or None)

    import paddle
    from ppocr.modeling.architectures import build_model as ppocr_build_model
    from ppocr.postprocess import build_post_process

    # 2. Device
    if args.device == "gpu" and not paddle.is_compiled_with_cuda():
        print("CUDA not available, falling back to CPU.")
        paddle.set_device("cpu")
    else:
        paddle.set_device(args.device)

    # 3. Resolve file paths
    weights_path = resolve_path(args.weights, "model.safetensors")
    dict_path = resolve_path(args.dict_path, "en_dict.txt")

    # 4. Build CTC post-processor
    post_process = build_post_process({
        "name": "CTCLabelDecode",
        "character_dict_path": dict_path,
        "use_space_char": True,
    })

    # 5. Build model and load weights
    config = copy.deepcopy(MODEL_CONFIG)
    model = ppocr_build_model(config["Architecture"])
    model.eval()

    state_dict = load_safetensors_to_paddle(paddle, weights_path)
    model.set_state_dict(state_dict)
    print(f"Loaded weights from {weights_path}")

    # 6. Run inference
    image_paths = collect_images(args.image_path)
    print(f"Found {len(image_paths)} image(s)\n")

    results = []
    for img_path in image_paths:
        img_bgr = cv2.imread(str(img_path))
        if img_bgr is None:
            print(f"WARNING: Could not read {img_path}, skipping.")
            continue

        tensor = paddle.to_tensor(
            np.expand_dims(preprocess(img_bgr, IMAGE_SHAPE), axis=0)
        )

        with paddle.no_grad():
            preds = model(tensor)

        if isinstance(preds, dict):
            pred_tensor = preds.get("ctc", next(iter(preds.values())))
        elif isinstance(preds, (list, tuple)):
            pred_tensor = preds[0]
        else:
            pred_tensor = preds

        post_result = post_process(pred_tensor.numpy())
        if isinstance(post_result, (list, tuple)) and len(post_result) > 0:
            text, confidence = post_result[0]
        else:
            text, confidence = "", 0.0

        text = text.strip().upper()
        result = {
            "image": str(img_path.name),
            "prediction": text,
            "confidence": round(float(confidence), 4),
        }
        results.append(result)
        print(f"  {img_path.name}: {text}  (conf: {confidence:.4f})")

    # 7. Save JSON
    if args.output_json:
        out_path = Path(args.output_json)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json.dumps(results, indent=2))
        print(f"\nResults saved to {out_path}")


if __name__ == "__main__":
    main()