#!/usr/bin/env python3
"""
PP-OCRv6 Detection ONNX Inference & Evaluation (standalone, zero Paddle dependency)

Dependencies:
    numpy, opencv-python, onnxruntime, pyyaml, shapely, pyclipper

Data format (same as PaddleOCR official):
    Label file: image_path<TAB>json_label
    json_label: [{"transcription": "text_or_###", "points": [[x,y]*4]}, ...]

Usage:
    # Single image inference
    python ppocrv6_det_onnx.py --det_onnx det.onnx --image test.jpg --visualize

    # Batch evaluation
    python ppocrv6_det_onnx.py --det_onnx det.onnx \\
        --label_file val.txt --dataset_root ./ocr_det_dataset_examples \\
        --visualize --output_json result.json --verbose
"""

import argparse
import json
import math
import os
from collections import namedtuple
from typing import List, Optional, Tuple

import cv2
import numpy as np
import axengine as ort
import yaml
from shapely.geometry import Polygon
import pyclipper


def _get_dim_value(dim):
    """Extract integer value from an ONNX Runtime dimension, returning 0 for dynamic dims."""
    if dim is None:
        return 0
    if isinstance(dim, str):
        return 0
    if hasattr(dim, 'dim_value'):
        return int(dim.dim_value) if dim.dim_value else 0
    if hasattr(dim, 'dim_param'):
        # named dim like "batch_size" → dynamic
        return 0
    try:
        v = int(dim)
        return v
    except (TypeError, ValueError):
        return 0


# ============================================================================
# 1. Detection Preprocessing
# ============================================================================

class _DetResizeForTest:
    def __init__(self, limit_side_len=960, limit_type="max", max_side_limit=4000,
                 image_shape=None, keep_ratio=False):
        self.max_side_limit = max_side_limit
        # resize_type=0: limit_side_len (dynamic input)
        # resize_type=1: image_shape (fixed ONNX input)
        if image_shape is not None:
            self.resize_type = 1
            self.image_shape = image_shape
            self.keep_ratio = keep_ratio
        else:
            self.resize_type = 0
            self.limit_side_len = limit_side_len
            self.limit_type = limit_type

    def _image_padding(self, im, value=0):
        h, w, c = im.shape
        im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
        im_pad[:h, :w, :] = im
        return im_pad

    def _resize_image_type0(self, img):
        h, w, _ = img.shape
        limit_side_len = self.limit_side_len

        if self.limit_type == "max":
            if max(h, w) > limit_side_len:
                ratio = float(limit_side_len) / max(h, w)
            else:
                ratio = 1.0
        elif self.limit_type == "min":
            if min(h, w) < limit_side_len:
                ratio = float(limit_side_len) / min(h, w)
            else:
                ratio = 1.0
        elif self.limit_type == "resize_long":
            ratio = float(limit_side_len) / max(h, w)
        else:
            raise ValueError(f"not support limit_type: {self.limit_type}")

        resize_h = int(h * ratio)
        resize_w = int(w * ratio)
        if max(resize_h, resize_w) > self.max_side_limit:
            ratio = float(self.max_side_limit) / max(resize_h, resize_w)
            resize_h, resize_w = int(resize_h * ratio), int(resize_w * ratio)

        resize_h = max(int(round(resize_h / 32) * 32), 32)
        resize_w = max(int(round(resize_w / 32) * 32), 32)

        if int(resize_w) <= 0 or int(resize_h) <= 0:
            return None, (None, None)
        img = cv2.resize(img, (int(resize_w), int(resize_h)))
        ratio_h = resize_h / float(h)
        ratio_w = resize_w / float(w)
        return img, [ratio_h, ratio_w]

    def _resize_image_type1(self, img):
        """Direct resize to fixed [H, W]. Used when ONNX has fixed input dimensions."""
        resize_h, resize_w = self.image_shape
        ori_h, ori_w = img.shape[:2]
        if self.keep_ratio:
            resize_w = ori_w * resize_h / ori_h
            N = math.ceil(resize_w / 32)
            resize_w = N * 32
        ratio_h = float(resize_h) / ori_h
        ratio_w = float(resize_w) / ori_w
        img = cv2.resize(img, (int(resize_w), int(resize_h)))
        return img, [ratio_h, ratio_w]

    def __call__(self, img):
        src_h, src_w = img.shape[:2]
        if sum([src_h, src_w]) < 64:
            img = self._image_padding(img)
        if self.resize_type == 1:
            img, [ratio_h, ratio_w] = self._resize_image_type1(img)
        else:
            img, [ratio_h, ratio_w] = self._resize_image_type0(img)
        shape = np.array([src_h, src_w, ratio_h, ratio_w])
        return img, shape


class _NormalizeImage:
    def __init__(self, mean, std, scale=1.0 / 255.0, order="hwc"):
        self.scale = np.float32(scale)
        shape = (1, 1, 3) if order == "hwc" else (3, 1, 1)
        self.mean = np.array(mean, dtype=np.float32).reshape(shape)
        self.std = np.array(std, dtype=np.float32).reshape(shape)

    def __call__(self, img):
        return (img.astype("float32") * self.scale - self.mean) / self.std


class _ToCHWImage:
    def __call__(self, img):
        return img.transpose((2, 0, 1))


# ============================================================================
# 2. Detection Postprocessing (DB)
# ============================================================================

class _DBPostProcess:
    def __init__(
        self,
        thresh=0.3,
        box_thresh=0.7,
        max_candidates=1000,
        unclip_ratio=2.0,
        use_dilation=False,
        score_mode="fast",
    ):
        self.thresh = thresh
        self.box_thresh = box_thresh
        self.max_candidates = max_candidates
        self.unclip_ratio = unclip_ratio
        self.min_size = 3
        self.score_mode = score_mode
        assert score_mode in ("slow", "fast")
        self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]])

    def _unclip(self, box, unclip_ratio):
        poly = Polygon(box)
        distance = poly.area * unclip_ratio / poly.length
        offset = pyclipper.PyclipperOffset()
        offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
        expanded = offset.Execute(distance)
        return expanded

    def _get_mini_boxes(self, contour):
        bounding_box = cv2.minAreaRect(contour)
        points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
        i1, i2, i3, i4 = 0, 1, 2, 3
        if points[1][1] > points[0][1]:
            i1, i4 = 0, 1
        else:
            i1, i4 = 1, 0
        if points[3][1] > points[2][1]:
            i2, i3 = 2, 3
        else:
            i2, i3 = 3, 2
        box = [points[i1], points[i2], points[i3], points[i4]]
        return box, min(bounding_box[1])

    def _box_score_fast(self, bitmap, _box):
        h, w = bitmap.shape[:2]
        box = _box.copy()
        xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
        xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
        ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
        ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
        box[:, 0] = box[:, 0] - xmin
        box[:, 1] = box[:, 1] - ymin
        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
        return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]

    def _boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
        bitmap = _bitmap
        height, width = bitmap.shape
        outs = cv2.findContours(
            (bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
        )
        if len(outs) == 3:
            _, contours, _ = outs
        else:
            contours, _ = outs
        num_contours = min(len(contours), self.max_candidates)
        boxes, scores = [], []
        for index in range(num_contours):
            contour = contours[index]
            points, sside = self._get_mini_boxes(contour)
            if sside < self.min_size:
                continue
            points = np.array(points)
            if self.score_mode == "fast":
                score = self._box_score_fast(pred, points.reshape(-1, 2))
            else:
                score = self._box_score_slow(pred, contour)
            if self.box_thresh > score:
                continue
            box = self._unclip(points, self.unclip_ratio)
            if len(box) > 1:
                continue
            box = np.array(box).reshape(-1, 1, 2)
            box, sside = self._get_mini_boxes(box)
            if sside < self.min_size + 2:
                continue
            box = np.array(box)
            box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
            box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
            boxes.append(box.astype("int32"))
            scores.append(score)
        return np.array(boxes, dtype="int32"), scores

    def _box_score_slow(self, bitmap, contour):
        h, w = bitmap.shape[:2]
        contour = contour.copy().reshape((-1, 2))
        xmin = np.clip(np.min(contour[:, 0]), 0, w - 1)
        xmax = np.clip(np.max(contour[:, 0]), 0, w - 1)
        ymin = np.clip(np.min(contour[:, 1]), 0, h - 1)
        ymax = np.clip(np.max(contour[:, 1]), 0, h - 1)
        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
        contour[:, 0] = contour[:, 0] - xmin
        contour[:, 1] = contour[:, 1] - ymin
        cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype("int32"), 1)
        return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]

    def __call__(self, pred, shape_list):
        pred = pred[:, 0, :, :]
        segmentation = pred > self.thresh
        boxes_batch = []
        for batch_index in range(pred.shape[0]):
            src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
            if self.dilation_kernel is not None:
                mask = cv2.dilate(
                    np.array(segmentation[batch_index]).astype(np.uint8),
                    self.dilation_kernel,
                )
            else:
                mask = segmentation[batch_index]
            boxes, _ = self._boxes_from_bitmap(pred[batch_index], mask, src_w, src_h)
            boxes_batch.append(boxes)
        return boxes_batch


# ============================================================================
# 3. Detection IoU Evaluator (inline from eval_det_iou.py)
# ============================================================================

_Rectangle = namedtuple("Rectangle", "xmin ymin xmax ymax")


def _get_intersection(pD, pG):
    return Polygon(pD).intersection(Polygon(pG)).area


def _get_union(pD, pG):
    return Polygon(pD).union(Polygon(pG)).area


def _get_iou(pD, pG):
    return _get_intersection(pD, pG) / _get_union(pD, pG)


class _DetectionIoUEvaluator:
    def __init__(self, iou_constraint=0.5, area_precision_constraint=0.5):
        self.iou_constraint = iou_constraint
        self.area_precision_constraint = area_precision_constraint

    def evaluate_image(self, gt: List[dict], det: List[dict]) -> dict:
        gt_pols = []
        gt_dont_care = []

        for n, g in enumerate(gt):
            points = g.get("points", [])
            if not points:
                continue
            try:
                if not Polygon(points).is_valid:
                    continue
            except Exception:
                continue
            gt_pols.append(points)
            if g.get("ignore", False):
                gt_dont_care.append(len(gt_pols) - 1)

        det_pols = []
        det_dont_care = []

        for n, d in enumerate(det):
            points = d.get("points", [])
            if not points:
                continue
            try:
                if not Polygon(points).is_valid:
                    continue
            except Exception:
                continue
            det_pols.append(points)
            if gt_dont_care:
                for dc_idx in gt_dont_care:
                    dc_pol = gt_pols[dc_idx]
                    inter = _get_intersection(dc_pol, points)
                    pd_area = Polygon(points).area
                    prec = 0 if pd_area == 0 else inter / pd_area
                    if prec > self.area_precision_constraint:
                        det_dont_care.append(len(det_pols) - 1)
                        break

        det_matched = 0

        if gt_pols and det_pols:
            iou_mat = np.empty([len(gt_pols), len(det_pols)])
            for g_i, g_pts in enumerate(gt_pols):
                for d_i, d_pts in enumerate(det_pols):
                    iou_mat[g_i, d_i] = _get_iou(d_pts, g_pts)

            gt_matched = np.zeros(len(gt_pols), dtype=np.uint8)
            det_matched_arr = np.zeros(len(det_pols), dtype=np.uint8)
            for g_i in range(len(gt_pols)):
                for d_i in range(len(det_pols)):
                    if (
                        gt_matched[g_i] == 0
                        and det_matched_arr[d_i] == 0
                        and g_i not in gt_dont_care
                        and d_i not in det_dont_care
                    ):
                        if iou_mat[g_i, d_i] > self.iou_constraint:
                            gt_matched[g_i] = 1
                            det_matched_arr[d_i] = 1
                            det_matched += 1

        num_gt_care = len(gt_pols) - len(gt_dont_care)
        num_det_care = len(det_pols) - len(det_dont_care)

        return {
            "gt_care": num_gt_care,
            "det_care": num_det_care,
            "det_matched": det_matched,
        }

    def combine_results(self, results: List[dict]) -> dict:
        num_gt = sum(r["gt_care"] for r in results)
        num_det = sum(r["det_care"] for r in results)
        matched = sum(r["det_matched"] for r in results)

        recall = 0 if num_gt == 0 else float(matched) / num_gt
        precision = 0 if num_det == 0 else float(matched) / num_det
        hmean = (
            0
            if (precision + recall) == 0
            else 2.0 * precision * recall / (precision + recall)
        )

        return {
            "precision": precision,
            "recall": recall,
            "hmean": hmean,
        }


# ============================================================================
# 4. Visualization
# ============================================================================

def draw_det_result(
    img: np.ndarray,
    det_boxes: np.ndarray,
    gt_boxes: Optional[List[dict]] = None,
    matched_pairs: Optional[List[Tuple[int, int]]] = None,
    thickness: int = 2,
) -> np.ndarray:
    """Draw detection boxes on image with optional ground-truth and matching info.

    Colors:
        green  = matched detection
        blue   = unmatched detection
        red    = unmatched ground truth (missed)
    """
    vis = img.copy()

    if gt_boxes is not None and matched_pairs is not None:
        gt_matched = set()
        det_matched = set()
        for p in matched_pairs:
            gt_matched.add(p["gt"])
            det_matched.add(p["det"])

        for i, d in enumerate(det_boxes):
            box = np.array(d, dtype=np.int32).reshape((-1, 1, 2))
            color = (0, 255, 0) if i in det_matched else (255, 0, 0)
            cv2.polylines(vis, [box], True, color, thickness)

        for i, g in enumerate(gt_boxes):
            if g.get("ignore", False):
                continue
            if i not in gt_matched:
                pts = np.array(g["points"], dtype=np.int32).reshape((-1, 1, 2))
                cv2.polylines(vis, [pts], True, (0, 0, 255), max(thickness, 3))
                # Draw dashed effect by alternating segments
                cx, cy = int(np.mean(pts[:, 0, 0])), int(np.mean(pts[:, 0, 1]))
                cv2.putText(vis, "MISS", (cx, cy),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
    else:
        for box in det_boxes:
            box = np.array(box, dtype=np.int32).reshape((-1, 1, 2))
            cv2.polylines(vis, [box], True, (0, 255, 0), thickness)

    return vis


# ============================================================================
# 5. Detection Engine
# ============================================================================

class PPOCRv6DetOnnx:

    def __init__(
        self,
        det_onnx: str,
        det_limit_side_len: int = 960,
        det_db_thresh: float = 0.2,
        det_db_box_thresh: float = 0.4,
        det_db_unclip_ratio: float = 1.4,
        det_max_candidates: int = 3000,
        use_gpu: bool = False,
        onnx_providers: Optional[List[str]] = None,
        resize_mode: str = "letterbox",
    ):
        assert resize_mode in ("letterbox", "stretch"), f"invalid resize_mode: {resize_mode}"

        self.session = ort.InferenceSession(det_onnx)
        self.input_name = self.session.get_inputs()[0].name

        # Detect fixed vs dynamic input dimensions
        det_input = self.session.get_inputs()[0]
        img_h = _get_dim_value(det_input.shape[2])
        img_w = _get_dim_value(det_input.shape[3])

        self._fixed_h = img_h if img_h > 0 else 0
        self._fixed_w = img_w if img_w > 0 else 0

        print(f"[PPOCRv6Det] ONNX input shape: {det_input.shape}, fixed_h={self._fixed_h}, fixed_w={self._fixed_w}, resize_mode={resize_mode}")

        self._resize_mode = resize_mode

        # Preprocessing: choose resize strategy
        if self._fixed_h > 0 and self._fixed_w > 0:
            # Fully fixed ONNX input — handled in _preprocess
            self._resize_style = "fixed"
        else:
            # Dynamic or partially-fixed — ratio-preserving resize
            self._resize = _DetResizeForTest(
                limit_side_len=det_limit_side_len, limit_type="max"
            )
            self._resize_style = "dynamic"

        self._normalize = _NormalizeImage(
            mean=[0., 0., 0.],
            std=[1.0, 1.0, 1.0],
            scale=1.0,
        )

        self._to_chw = _ToCHWImage()
        self._post = _DBPostProcess(
            thresh=det_db_thresh,
            box_thresh=det_db_box_thresh,
            unclip_ratio=det_db_unclip_ratio,
            max_candidates=det_max_candidates,
        )

    def _preprocess(self, img: np.ndarray):
        src_h, src_w = img.shape[:2]
        fixed_w = self._fixed_w
        fixed_h = self._fixed_h

        # Stretch mode: direct resize to fixed size (official PaddleOCR behavior)
        if self._resize_mode == "stretch" and fixed_h > 0 and fixed_w > 0:
            img_resized = cv2.resize(img, (fixed_w, fixed_h))
            ratio_h = float(fixed_h) / src_h
            ratio_w = float(fixed_w) / src_w
            # Post-processing maps: origin = fm_coord / fm_dim * dest_dim
            # For direct stretch, fm_dim corresponds uniformly to src_dim.
            shape = np.array([src_h, src_w, ratio_h, ratio_w])

        elif fixed_w > 0 or fixed_h > 0:
            # Letterbox mode (default): ratio-preserving + pad to fixed size
            ratios = []
            if fixed_w > 0:
                ratios.append(fixed_w / src_w)
            if fixed_h > 0:
                ratios.append(fixed_h / src_h)
            ratio = min(ratios)
            new_w = max(int(round(src_w * ratio / 32) * 32), 32)
            new_h = max(int(round(src_h * ratio / 32) * 32), 32)
            new_w = min(new_w, fixed_w) if fixed_w > 0 else new_w
            new_h = min(new_h, fixed_h) if fixed_h > 0 else new_h

            img_resized = cv2.resize(img, (new_w, new_h))
            ratio_h = new_h / float(src_h)
            ratio_w = new_w / float(src_w)

            pad_h = max(0, fixed_h - new_h)
            pad_w = max(0, fixed_w - new_w)
            if pad_h > 0 or pad_w > 0:
                img_resized = cv2.copyMakeBorder(
                    img_resized, 0, pad_h, 0, pad_w,
                    cv2.BORDER_CONSTANT, value=(0, 0, 0),
                )

            # Adjust shape for correct coordinate mapping after padding
            adj_h = src_h * fixed_h / new_h if fixed_h > 0 else src_h
            adj_w = src_w * fixed_w / new_w if fixed_w > 0 else src_w
            shape = np.array([adj_h, adj_w, ratio_h, ratio_w])

        else:
            img_resized, shape = self._resize(img)

        img_norm = self._normalize(img_resized)
        img_chw = self._to_chw(img_norm)
        tensor = np.expand_dims(img_chw.astype(np.float32), axis=0)
        return tensor, shape

    def _postprocess(self, output: np.ndarray, shape: np.ndarray):
        shape_list = np.expand_dims(shape, axis=0)
        boxes_batch = self._post(output, shape_list)
        return boxes_batch[0]

    def __call__(self, img: np.ndarray) -> np.ndarray:
        """Detect text boxes. Returns (N, 4, 2) int32 array."""
        tensor, shape = self._preprocess(img)
        onnx_out = self.session.run(None, {self.input_name: tensor})
        boxes = self._postprocess(onnx_out[0], shape)
        return boxes

    def predict_image(self, path: str) -> np.ndarray:
        im = cv2.imread(path)
        if im is None:
            raise FileNotFoundError(f"Cannot read: {path}")
        return self.__call__(im)


# ============================================================================
# 6. Evaluation
# ============================================================================

def evaluate(
    det: PPOCRv6DetOnnx,
    label_file: str,
    dataset_root: str = "",
    iou_constraint: float = 0.5,
    verbose: bool = False,
) -> dict:
    """Evaluate detection against a PaddleOCR format label file.

    Label format (one per line, tab-separated)::

        rel/path/to/img.jpg<TAB>[{"transcription":"text_or_###","points":[[x,y]*4]}, ...]

    "###" means ignored / don't-care region.

    Returns:
        dict: precision, recall, hmean, total_images, total_gt, total_det, det_matched, per_sample
    """
    samples = []
    with open(label_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) < 2:
                continue
            img_path = os.path.join(dataset_root, parts[0].strip())
            try:
                gt_label = json.loads(parts[1])
            except json.JSONDecodeError:
                continue
            if not isinstance(gt_label, list):
                continue
            samples.append((img_path, gt_label))

    if not samples:
        print("[WARN] No samples found in label file.")
        return {
            "precision": 0, "recall": 0, "hmean": 0,
            "total_images": 0, "total_gt": 0, "total_det": 0, "det_matched": 0,
            "per_sample": [],
        }

    evaluator = _DetectionIoUEvaluator(iou_constraint=iou_constraint)
    per_sample = []
    total_gt = 0
    total_det = 0
    total_matched = 0

    for idx, (img_path, gt_label) in enumerate(samples):
        img = cv2.imread(img_path)
        if img is None:
            print(f"[WARN] Cannot read {img_path}, skipping.")
            per_sample.append({
                "image": img_path, "error": "cannot read",
                "gt_care": 0, "det_care": 0, "det_matched": 0,
            })
            continue

        # Run detection
        det_boxes = det(img)

        # Prepare GT
        gt_info = []
        for g in gt_label:
            pts = g.get("points", [])
            if not pts or len(pts) < 4:
                continue
            is_ignored = g.get("transcription", "") == "###"
            gt_info.append({"points": pts, "ignore": is_ignored})

        # Prepare DET
        det_info = [{"points": d.tolist()} for d in det_boxes]

        # Evaluate
        result = evaluator.evaluate_image(gt_info, det_info)
        total_gt += result["gt_care"]
        total_det += result["det_care"]
        total_matched += result["det_matched"]

        per_sample.append({
            "image": img_path,
            "gt_care": result["gt_care"],
            "det_care": result["det_care"],
            "det_matched": result["det_matched"],
            "det_boxes": [d.tolist() for d in det_boxes],
        })

        if verbose:
            miss = result["gt_care"] - result["det_matched"]
            extra = result["det_care"] - result["det_matched"]
            parts = [
                f"gt={result['gt_care']}",
                f"det={result['det_care']}",
                f"match={result['det_matched']}",
            ]
            if miss > 0:
                parts.append(f"MISS={miss}")
            if extra > 0:
                parts.append(f"EXTRA={extra}")
            print(f"[{os.path.basename(img_path)}] " + " ".join(parts))

    metrics = evaluator.combine_results(
        [r for r in per_sample if "error" not in r]
    )

    return {
        "precision": round(metrics["precision"], 6),
        "recall": round(metrics["recall"], 6),
        "hmean": round(metrics["hmean"], 6),
        "total_images": len(samples),
        "total_gt": total_gt,
        "total_det": total_det,
        "det_matched": total_matched,
        "per_sample": per_sample,
    }


# ============================================================================
# 7. CLI
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="PP-OCRv6 Detection ONNX – inference & evaluation"
    )
    # Model
    parser.add_argument("--det_onnx", type=str,
                        default="./axmodel/ax650/det_npu2.axmodel",
                        help="Path to detection ONNX model")
    parser.add_argument("--limit_side_len", type=int, default=960)
    parser.add_argument("--det_db_thresh", type=float, default=0.2)
    parser.add_argument("--det_db_box_thresh", type=float, default=0.45)
    parser.add_argument("--det_db_unclip_ratio", type=float, default=1.4)
    parser.add_argument("--use_gpu", action="store_true", help="Enable GPU inference")
    parser.add_argument("--resize_mode", type=str, default="letterbox",
                        choices=["letterbox", "stretch"],
                        help="Resize strategy for fixed-size ONNX: letterbox (keep ratio+pad) or stretch (direct resize)")

    # Single image mode
    parser.add_argument("--image", type=str, default=None, help="Single image path")

    # Evaluation mode
    parser.add_argument("--label_file", type=str, 
                        default='dataset/ocr_det_dataset_examples/val.txt',
                        help="Label file (image_path<TAB>json_label per line)")
    parser.add_argument("--dataset_root", type=str, 
                        default="dataset/ocr_det_dataset_examples",
                        help="Prefix directory for image paths in label file")

    # Common
    parser.add_argument("--visualize", action="store_true", help="Draw boxes on image")
    parser.add_argument("--output", type=str, default=None,
                        help="Save visualized image (implies --visualize)")
    parser.add_argument("--verbose", action="store_true", help="Print per-image metrics")
    parser.add_argument("--output_json", type=str, default=None,
                        help="Save results to JSON file")

    args = parser.parse_args()

    det = PPOCRv6DetOnnx(
        det_onnx=args.det_onnx,
        det_limit_side_len=args.limit_side_len,
        det_db_thresh=args.det_db_thresh,
        det_db_box_thresh=args.det_db_box_thresh,
        det_db_unclip_ratio=args.det_db_unclip_ratio,
        use_gpu=args.use_gpu,
        resize_mode=args.resize_mode,
    )

    # --- Single image mode ---
    if args.image and not args.label_file:
        img = cv2.imread(args.image)
        if img is None:
            raise FileNotFoundError(f"Cannot read: {args.image}")
        boxes = det(img)
        print(f"Detected {len(boxes)} text boxes:")
        for i, box in enumerate(boxes):
            print(f"  [{i}] {box.tolist()}")

        do_viz = args.visualize or args.output
        if do_viz:
            vis = draw_det_result(img, boxes)
            out_path = args.output or "det_result.jpg"
            cv2.imwrite(out_path, vis)
            print(f"Visualization saved to: {out_path}")

        if args.output_json:
            with open(args.output_json, "w") as f:
                json.dump(
                    {"image": args.image, "boxes": [b.tolist() for b in boxes]},
                    f, indent=2,
                )
            print(f"Results saved to: {args.output_json}")
        return

    # --- Evaluation mode ---
    if args.label_file:
        metrics = evaluate(
            det, args.label_file,
            dataset_root=args.dataset_root,
            verbose=args.verbose,
        )

        print()
        print("=" * 60)
        print("Evaluation Results")
        print("=" * 60)
        print(f"  Images:       {metrics['total_images']}")
        print(f"  GT boxes:     {metrics['total_gt']}")
        print(f"  DET boxes:    {metrics['total_det']}")
        print(f"  Matched:      {metrics['det_matched']}")
        print(f"  Precision:    {metrics['precision']:.4f} ({metrics['precision']*100:.2f}%)")
        print(f"  Recall:       {metrics['recall']:.4f} ({metrics['recall']*100:.2f}%)")
        print(f"  Hmean (F1):   {metrics['hmean']:.4f}")
        print("=" * 60)

        # Visualization for eval mode
        do_viz = args.visualize or args.output
        if do_viz:
            out_dir = args.output if args.output else "det_eval_vis"
            os.makedirs(out_dir, exist_ok=True)
            for i, smp in enumerate(metrics["per_sample"]):
                img = cv2.imread(smp["image"])
                if img is None:
                    continue
                # Load GT boxes with matching info
                with open(args.label_file, "r") as f:
                    lines = f.readlines()
                gt_label = []
                for line in lines:
                    line = line.strip()
                    if not line:
                        continue
                    parts = line.split("\t")
                    if len(parts) < 2:
                        continue
                    if os.path.join(args.dataset_root, parts[0].strip()) == smp["image"]:
                        gt_label = json.loads(parts[1])
                        break

                det_boxes = np.array(smp.get("det_boxes", []))
                # Simple matching for visualization (re-run evaluate_image)
                gt_info = []
                for g in gt_label:
                    pts = g.get("points", [])
                    if not pts or len(pts) < 4:
                        continue
                    gt_info.append({
                        "points": pts,
                        "ignore": g.get("transcription", "") == "###",
                    })
                det_info = [{"points": d} for d in det_boxes.tolist()]

                # Compute matches for coloring
                matched_pairs = _compute_matched_pairs(gt_info, det_info)
                vis = draw_det_result(img, det_boxes, gt_info, matched_pairs)
                fname = os.path.basename(smp["image"])
                cv2.imwrite(os.path.join(out_dir, fname), vis)

            print(f"Visualization saved to: {out_dir}/")

        if args.output_json:
            out = {k: v for k, v in metrics.items() if k != "per_sample"}
            out["per_sample"] = metrics["per_sample"]
            with open(args.output_json, "w") as f:
                json.dump(out, f, indent=2)
            print(f"Results saved to: {args.output_json}")
        return

    parser.error("Either --image or --label_file must be provided.")


def _compute_matched_pairs(gt_info, det_info, iou_thr=0.5):
    """Compute matched GT-det pairs for visualization coloring."""
    pairs = []
    gt_pols = [g["points"] for g in gt_info if not g.get("ignore")]
    det_pols = [d["points"] for d in det_info]
    if not gt_pols or not det_pols:
        return pairs

    iou_mat = np.empty([len(gt_pols), len(det_pols)])
    for g_i, g_pts in enumerate(gt_pols):
        for d_i, d_pts in enumerate(det_pols):
            try:
                int_area = _get_intersection(d_pts, g_pts)
                union_area = _get_union(d_pts, g_pts)
                iou_mat[g_i, d_i] = int_area / union_area if union_area > 0 else 0
            except Exception:
                iou_mat[g_i, d_i] = 0

    gt_used = set()
    det_used = set()
    # Greedy matching by descending IoU
    flat = []
    for g_i in range(len(gt_pols)):
        for d_i in range(len(det_pols)):
            flat.append((iou_mat[g_i, d_i], g_i, d_i))
    flat.sort(key=lambda x: x[0], reverse=True)
    for iou, g_i, d_i in flat:
        if iou > iou_thr and g_i not in gt_used and d_i not in det_used:
            pairs.append({"gt": g_i, "det": d_i})
            gt_used.add(g_i)
            det_used.add(d_i)
    return pairs


if __name__ == "__main__":
    main()