#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
合成训练数据：把「靶子」贴到随机背景上，并自动生成标注（无需手工标注）。

前置条件（推荐）：
  - 靶子用带透明通道的 PNG（抠图后），脚本按非透明像素算紧贴 bbox；
  - 若只有矩形靶图无 alpha，可用整张图作为矩形框贴入（略松）。

输出（默认 Pascal VOC，适配 MaixCam 等平台）：
  - images/xxx.jpg
  - xml/xxx.xml（与图片同名；单目标或多目标时可扩展）
  - 生成张数不超过 --max-images（默认 3000）
可选 YOLO：
  - labels/xxx.txt（class cx cy w h，相对 0~1）

多三角形检测（Pascal VOC 多 <object>，适配 YOLOv5 转 VOC 训练）：
  - 提供 --triangles-json，顶点在与 --fg 一致的原始靶图像素坐标系下；
  - 脚本先按 alpha 外接框裁切靶图，顶点会自动减去裁切偏移；
  - 透视变换时同步变换顶点，每张图输出多个三角形框；
  - 默认标注为顶点轴对齐最小外接矩形；可选 --triangle-bbox-pad-frac 四周加比例余量（与推理 margin 对齐）。

Stage2 ROI（对齐「先检整靶再裁小图」的第二步输入）：
  - --stage2-crop：在合成+增强后，按靶子外接框四周随机 padding 裁剪，标注改到裁剪图坐标系；
  - 有 --triangles-json 时默认要求裁剪后三角形数与 JSON 一致，否则丢弃重采样（可用 --stage2-allow-partial）。

运动模糊（模拟手持/快门，默认约一半样本会施加；标注仍为几何真值，与真机域更接近）：
  - --motion-prob：施加概率；--motion-kernel-min/max：模糊 streak 长度（奇数核，越大越糊）。
  - 可与 --blur-max 高斯模糊叠加；Stage2 建议：--motion-prob 0.5~0.7 --motion-kernel-max 35 --blur-max 1.2

依赖：OpenCV + NumPy（PC 上跑即可；Maix 上若内存够也可试）。

示例：
  python test/synth_compose_yolo.py --bg-dir ./bg --fg ./target_cutout.png --out ./synth_out --num 3000
  python test/synth_compose_yolo.py ... --triangles-json test/archery_triangles_default.json --class-name triangle --stage2-crop
  python test/synth_compose_yolo.py ... --zip ./dataset_voc.zip
  python test/synth_compose_yolo.py ... --format yolo --out ./synth_yolo
"""

from __future__ import annotations

import argparse
import json
import os
import random
import sys
import zipfile
import xml.etree.ElementTree as ET

import numpy as np


def _collect_images(folder: str, exts=(".jpg", ".jpeg", ".png", ".bmp")):
    out = []
    for name in sorted(os.listdir(folder)):
        low = name.lower()
        if low.endswith(exts):
            out.append(os.path.join(folder, name))
    return out


def _load_triangles_json(path: str) -> list[list[tuple[float, float]]]:
    with open(path, encoding="utf-8") as f:
        data = json.load(f)
    tris = data.get("triangles")
    if not isinstance(tris, list) or not tris:
        raise ValueError(f'JSON 需包含非空 "triangles" 数组: {path}')
    out: list[list[tuple[float, float]]] = []
    for t in tris:
        if not isinstance(t, list) or len(t) != 3:
            raise ValueError(f"每个三角形需 3 个顶点: {t!r}")
        pts = []
        for p in t:
            if not isinstance(p, (list, tuple)) or len(p) != 2:
                raise ValueError(f"顶点需为 [x,y]: {p!r}")
            pts.append((float(p[0]), float(p[1])))
        out.append(pts)
    return out


def _warp_triangle_points(
    corners_fg_orig: list[tuple[float, float]],
    fx0: float,
    fy0: float,
    fw0: float,
    fh0: float,
    new_w: int,
    new_h: int,
    persp_M,
    px: int,
    py: int,
    np,
    cv2,
) -> np.ndarray:
    """原始靶图像素坐标下的三角形顶点 -> 合成图上的 (3,2) float32。"""
    pts = np.array(corners_fg_orig, dtype=np.float32)
    pts[:, 0] -= fx0
    pts[:, 1] -= fy0
    pts[:, 0] *= new_w / max(fw0, 1e-6)
    pts[:, 1] *= new_h / max(fh0, 1e-6)
    if persp_M is not None:
        pts = cv2.perspectiveTransform(pts.reshape(1, -1, 2), persp_M).reshape(-1, 2)
    pts[:, 0] += px
    pts[:, 1] += py
    return pts


def _triangle_xyxy_exclusive(
    pts_xy: np.ndarray, img_w: int, img_h: int
) -> tuple[int, int, int, int] | None:
    xs = pts_xy[:, 0]
    ys = pts_xy[:, 1]
    bx0 = max(0, min(img_w - 1, int(np.floor(float(xs.min())))))
    by0 = max(0, min(img_h - 1, int(np.floor(float(ys.min())))))
    bx1 = max(bx0 + 1, min(img_w, int(np.ceil(float(xs.max())))))
    by1 = max(by0 + 1, min(img_h, int(np.ceil(float(ys.max())))))
    if bx1 <= bx0 or by1 <= by0:
        return None
    return bx0, by0, bx1, by1


def _expand_xyxy_half_open(
    bx0: int,
    by0: int,
    bx1: int,
    by1: int,
    img_w: int,
    img_h: int,
    pad_frac: float,
) -> tuple[int, int, int, int] | None:
    """在半开框 [bx0,bx1)×[by0,by1) 四周按 max(宽,高)×pad_frac 对称扩展，并裁入图像。"""
    if pad_frac <= 1e-9:
        return bx0, by0, bx1, by1
    bw = max(1, bx1 - bx0)
    bh = max(1, by1 - by0)
    base = float(max(bw, bh))
    p = float(pad_frac) * base
    x0 = int(np.floor(float(bx0) - p))
    y0 = int(np.floor(float(by0) - p))
    x1 = int(np.ceil(float(bx1) + p))
    y1 = int(np.ceil(float(by1) + p))
    iw, ih = max(1, img_w), max(1, img_h)
    x0 = max(0, min(x0, iw - 1))
    y0 = max(0, min(y0, ih - 1))
    x1 = max(x0 + 1, min(x1, iw))
    y1 = max(y0 + 1, min(y1, ih))
    if x1 <= x0 or y1 <= y0:
        return None
    return x0, y0, x1, y1


def _stage2_crop_window(
    tx0: int,
    ty0: int,
    tx1: int,
    ty1: int,
    img_w: int,
    img_h: int,
    pad_min_frac: float,
    pad_max_frac: float,
    rng: random.Random,
) -> tuple[int, int, int, int] | None:
    """
    以靶子轴对齐框 [tx0,tx1)×[ty0,ty1)（半开）为中心，四周加随机 padding（相对 max(宽,高) 的比例），
    再限制在图像内。返回 (cx0, cy0, cw, ch) 用于 comp[cy0:cy0+ch, cx0:cx0+cw]。
    """
    iw, ih = max(1, img_w), max(1, img_h)
    tw = max(1, tx1 - tx0)
    th = max(1, ty1 - ty0)
    base = float(max(tw, th))
    p0 = max(0.0, float(pad_min_frac))
    p1 = max(p0, float(pad_max_frac))
    pad = rng.uniform(p0, p1) * base
    cx0 = int(np.floor(float(tx0) - pad))
    cy0 = int(np.floor(float(ty0) - pad))
    cx1 = int(np.ceil(float(tx1) + pad))
    cy1 = int(np.ceil(float(ty1) + pad))
    cx0 = max(0, min(cx0, iw - 1))
    cy0 = max(0, min(cy0, ih - 1))
    cx1 = max(cx0 + 1, min(cx1, iw))
    cy1 = max(cy0 + 1, min(cy1, ih))
    cw, ch = cx1 - cx0, cy1 - cy0
    if cw < 4 or ch < 4:
        return None
    return cx0, cy0, cw, ch


def _triangle_to_voc_tuple(
    pts_xy: np.ndarray,
    img_w: int,
    img_h: int,
    class_name: str,
    bbox_pad_frac: float = 0.0,
) -> tuple | None:
    """
    返回 (VOC 元组, 半开 xyxy)；半开框与 VOC 一致地经 pad 扩展，供 YOLO 行写入。
    bbox_pad_frac>0 时在紧三角形 AABB 四周加 max(宽,高)×frac 余量（truncated 仍按顶点是否贴边）。
    """
    xyxy = _triangle_xyxy_exclusive(pts_xy, img_w, img_h)
    if xyxy is None:
        return None
    bx0, by0, bx1, by1 = xyxy
    if bbox_pad_frac > 1e-9:
        exp = _expand_xyxy_half_open(
            bx0, by0, bx1, by1, img_w, img_h, bbox_pad_frac
        )
        if exp is None:
            return None
        bx0, by0, bx1, by1 = exp
    xs = pts_xy[:, 0]
    ys = pts_xy[:, 1]
    truncated = (
        "1"
        if (
            xs.min() < -1e-3
            or xs.max() >= img_w - 1e-3
            or ys.min() < -1e-3
            or ys.max() >= img_h - 1e-3
        )
        else "0"
    )
    vx0, vy0, vx1, vy1 = _xyxy_exclusive_to_voc_inclusive(
        bx0, by0, bx1, by1, img_w, img_h
    )
    if vx1 < vx0 or vy1 < vy0:
        return None
    voc = (class_name, vx0, vy0, vx1, vy1, truncated)
    return voc, (bx0, by0, bx1, by1)


def _fg_bbox_from_alpha(fg_bgra):
    """非透明区域的外接矩形 (x,y,w,h)，BGRA。"""
    import numpy as np

    if fg_bgra.shape[2] < 4:
        h, w = fg_bgra.shape[:2]
        return 0, 0, w, h
    a = fg_bgra[:, :, 3]
    ys, xs = np.where(a > 10)
    if len(xs) == 0:
        h, w = fg_bgra.shape[:2]
        return 0, 0, w, h
    x0, x1 = int(xs.min()), int(xs.max())
    y0, y1 = int(ys.min()), int(ys.max())
    return x0, y0, x1 - x0 + 1, y1 - y0 + 1


def _paste_fg_on_bg(bg_bgr, x, y, fg_scaled_bgra):
    """左上角 (x,y) 将 fg_scaled_bgra（BGRA）贴到 bg_bgr，就地改 bg。"""
    import numpy as np

    fh, fw = fg_scaled_bgra.shape[:2]
    bh, bw = bg_bgr.shape[:2]
    x0, y0 = max(0, x), max(0, y)
    x1, y1 = min(bw, x + fw), min(bh, y + fh)
    if x0 >= x1 or y0 >= y1:
        return
    fx0, fy0 = x0 - x, y0 - y
    fx1, fy1 = fx0 + (x1 - x0), fy0 + (y1 - y0)
    roi_bg = bg_bgr[y0:y1, x0:x1]
    roi_fg = fg_scaled_bgra[fy0:fy1, fx0:fx1]
    a = roi_fg[:, :, 3:4].astype(np.float32) / 255.0
    fg_rgb = roi_fg[:, :, :3].astype(np.float32)
    bg_rgb = roi_bg.astype(np.float32)
    blended = fg_rgb * a + bg_rgb * (1.0 - a)
    roi_bg[:] = blended.astype(np.uint8)


def _perspective_warp_rgba(img_bgra, jitter_frac: float, rng: random.Random, np, cv2):
    """
    对前景做轻微透视（四角微移），返回 (warped BGRA, M)。
    M 为 3×3，将透视前图像平面上的点映射到 warped 图像像素坐标；未应用透视时返回 (copy, None)。
    jitter_frac：扰动幅度约为 min(w,h) 的比例。
    """
    h, w = img_bgra.shape[:2]
    if jitter_frac <= 0 or min(w, h) < 16:
        return img_bgra.copy(), None

    j = float(max(1.5, min(w, h) * jitter_frac))

    def dj():
        return rng.uniform(-j, j)

    pts_src = np.float32([[0, 0], [w, 0], [w, h], [0, h]])
    pts_dst = np.float32(
        [
            [dj(), dj()],
            [w + dj(), dj()],
            [w + dj(), h + dj()],
            [dj(), h + dj()],
        ]
    )

    xmin = float(pts_dst[:, 0].min())
    ymin = float(pts_dst[:, 1].min())
    pts_shift = pts_dst.copy()
    pts_shift[:, 0] -= xmin
    pts_shift[:, 1] -= ymin
    out_w = max(4, int(np.ceil(float(pts_shift[:, 0].max()))) + 2)
    out_h = max(4, int(np.ceil(float(pts_shift[:, 1].max()))) + 2)

    M = cv2.getPerspectiveTransform(pts_src, pts_shift)
    warped = cv2.warpPerspective(
        img_bgra,
        M,
        (out_w, out_h),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(0, 0, 0, 0),
    )
    return warped, M


def _color_jitter_bgr(comp_bgr, strength: float, rng: random.Random, np, cv2):
    """整图 HSV 抖动：strength∈[0,1] 越大越强。"""
    if strength <= 1e-6:
        return comp_bgr
    strength = min(1.0, max(0.0, strength))
    hsv = cv2.cvtColor(comp_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
    dh = rng.uniform(-18.0 * strength, 18.0 * strength)
    hsv[:, :, 0] = (hsv[:, :, 0] + dh) % 180.0
    sf = rng.uniform(1.0 - 0.22 * strength, 1.0 + 0.22 * strength)
    vf = rng.uniform(1.0 - 0.22 * strength, 1.0 + 0.22 * strength)
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * sf, 0, 255)
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] * vf, 0, 255)
    # 轻微 BGR 通道偏置（模拟白平衡）
    out = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR).astype(np.float32)
    bias = np.array(
        [
            rng.uniform(-12 * strength, 12 * strength),
            rng.uniform(-12 * strength, 12 * strength),
            rng.uniform(-12 * strength, 12 * strength),
        ],
        dtype=np.float32,
    )
    out = np.clip(out + bias, 0, 255).astype(np.uint8)
    return out


def _motion_blur_bgr(
    comp_bgr,
    rng: random.Random,
    k_min: int,
    k_max: int,
    np,
    cv2,
):
    """
    方向随机的线性运动模糊（filter2D）。核为奇数 k×k，沿穿过中心、角度 uniform[0,180°) 的线段归一化求和。
    标注无需改：bbox 仍为物体真实位置，与真实相机「糊图+真框」的训练惯例一致。
    """
    lo = int(max(3, k_min | 1))
    hi = int(max(lo, k_max | 1))
    k = rng.randint(lo, hi)
    if k % 2 == 0:
        k = min(hi, k + 1)
    k = max(3, k)
    ker_u = np.zeros((k, k), dtype=np.uint8)
    ang = rng.uniform(0.0, 180.0)
    rad = float(np.deg2rad(ang))
    c = k // 2
    dx = float(np.cos(rad) * (k // 2))
    dy = float(np.sin(rad) * (k // 2))
    x0 = int(round(c - dx))
    y0 = int(round(c - dy))
    x1 = int(round(c + dx))
    y1 = int(round(c + dy))
    cv2.line(ker_u, (x0, y0), (x1, y1), 255, 1)
    s = float(ker_u.sum())
    if s < 1e-3:
        ker_u[c, c] = 255
        s = 255.0
    ker = ker_u.astype(np.float32) / s
    return cv2.filter2D(comp_bgr, -1, ker)


def _yolo_line(cls: int, xyxy_on_bg, img_w: int, img_h: int) -> str:
    x0, y0, x1, y1 = xyxy_on_bg
    bw, bh = x1 - x0, y1 - y0
    cx = (x0 + x1) / 2.0 / img_w
    cy = (y0 + y1) / 2.0 / img_h
    nw = bw / img_w
    nh = bh / img_h
    cx = max(0.0, min(1.0, cx))
    cy = max(0.0, min(1.0, cy))
    nw = max(1e-6, min(1.0, nw))
    nh = max(1e-6, min(1.0, nh))
    return f"{cls} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}\n"


def _xyxy_exclusive_to_voc_inclusive(
    x0: float, y0: float, x1: float, y1: float, img_w: int, img_h: int
) -> tuple[int, int, int, int]:
    """内部 xyxy 为半开区间 [x0,x1)×[y0,y1)，转为 VOC  inclusive 整数像素框。"""
    iw, ih = max(1, img_w), max(1, img_h)
    xi0 = max(0, min(iw - 1, int(x0)))
    yi0 = max(0, min(ih - 1, int(y0)))
    xi1 = max(xi0, min(iw - 1, int(x1) - 1))
    yi1 = max(yi0, min(ih - 1, int(y1) - 1))
    return xi0, yi0, xi1, yi1


def _write_pascal_voc_xml(
    xml_path: str,
    img_filename: str,
    img_folder: str,
    img_w: int,
    img_h: int,
    depth: int,
    objects: list[tuple],
) -> None:
    """
    objects 每项为 (class_name, xmin, ymin, xmax, ymax) 或
    (class_name, xmin, ymin, xmax, ymax, truncated)，坐标均为 inclusive 整数像素；
    truncated 为 \"0\" 或 \"1\"（省略时默认为 \"0\"）。
    """
    root = ET.Element("annotation")
    ET.SubElement(root, "folder").text = img_folder
    ET.SubElement(root, "filename").text = img_filename
    src = ET.SubElement(root, "source")
    ET.SubElement(src, "database").text = "synthetic_archery"
    ET.SubElement(src, "annotation").text = "Pascal VOC compatible"
    sz = ET.SubElement(root, "size")
    ET.SubElement(sz, "width").text = str(img_w)
    ET.SubElement(sz, "height").text = str(img_h)
    ET.SubElement(sz, "depth").text = str(depth)
    ET.SubElement(root, "segmented").text = "0"
    for item in objects:
        if len(item) == 6:
            name, xmin, ymin, xmax, ymax, truncated = item
        else:
            name, xmin, ymin, xmax, ymax = item
            truncated = "0"
        obj = ET.SubElement(root, "object")
        ET.SubElement(obj, "name").text = name
        ET.SubElement(obj, "pose").text = "Unspecified"
        ET.SubElement(obj, "truncated").text = str(truncated)
        ET.SubElement(obj, "difficult").text = "0"
        bb = ET.SubElement(obj, "bndbox")
        ET.SubElement(bb, "xmin").text = str(xmin)
        ET.SubElement(bb, "ymin").text = str(ymin)
        ET.SubElement(bb, "xmax").text = str(xmax)
        ET.SubElement(bb, "ymax").text = str(ymax)

    tree = ET.ElementTree(root)
    try:
        ET.indent(tree, space="  ")
    except AttributeError:
        pass
    tree.write(xml_path, encoding="utf-8", xml_declaration=True)


def _zip_images_xml(dataset_root: str, zip_path: str) -> None:
    """打包 dataset_root 下的 images/ 与 xml/ 到 zip（根目录含这两个文件夹）。"""
    img_dir = os.path.join(dataset_root, "images")
    xml_dir = os.path.join(dataset_root, "xml")
    if not os.path.isdir(img_dir) or not os.path.isdir(xml_dir):
        raise FileNotFoundError(f"需要存在目录: {img_dir} 与 {xml_dir}")
    zip_path = os.path.abspath(zip_path)
    os.makedirs(os.path.dirname(zip_path) or ".", exist_ok=True)
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for folder, arc_prefix in ((img_dir, "images"), (xml_dir, "xml")):
            for name in sorted(os.listdir(folder)):
                fp = os.path.join(folder, name)
                if os.path.isfile(fp):
                    zf.write(fp, arcname=os.path.join(arc_prefix, name).replace("\\", "/"))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--bg-dir", required=True, help="背景图目录")
    ap.add_argument("--fg", required=True, help="靶子 PNG（推荐 RGBA 抠图）或任意图")
    ap.add_argument("--out", default="./synth_dataset", help="输出根目录")
    ap.add_argument("--num", type=int, default=200, help="请求生成张数（实际不超过 --max-images）")
    ap.add_argument(
        "--max-images",
        type=int,
        default=3000,
        help="最多生成图片张数，超出部分忽略（MaixCam 等平台常见上限 3000）",
    )
    ap.add_argument(
        "--format",
        choices=("voc", "yolo", "both"),
        default="voc",
        help="voc=Pascal VOC（images+xml）；yolo=labels txt；both=两者都写",
    )
    ap.add_argument(
        "--class-name",
        default="黑三角和圆环",
        help="VOC <object><name> 类别名（单类检测默认 target）",
    )
    ap.add_argument("--class-id", type=int, default=0, help="YOLO 类别 id（仅 --format yolo/both）")
    ap.add_argument(
        "--zip",
        default=None,
        metavar="PATH",
        help="完成后将 images/ 与 xml/ 打成 zip（仅 VOC/both 时有 xml；路径如 ./dataset.zip）",
    )
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--scale-min", type=float, default=0.15, help="靶子最短边占背景最短边比例下限")
    ap.add_argument("--scale-max", type=float, default=0.55, help="比例上限")
    ap.add_argument("--blur-max", type=float, default=0.0, help="高斯模糊 sigma 上限，0 关闭")
    ap.add_argument(
        "--motion-prob",
        type=float,
        default=0.45,
        help="运动模糊概率 0~1（默认约一半样本；关模糊用 0）",
    )
    ap.add_argument(
        "--motion-kernel-min",
        type=int,
        default=7,
        help="运动模糊 streak 长度下限（奇数，实际会纠到奇数）",
    )
    ap.add_argument(
        "--motion-kernel-max",
        type=int,
        default=35,
        help="运动模糊 streak 长度上限，越大越像长曝光/手抖",
    )
    ap.add_argument("--jpeg-quality", type=int, default=92)
    ap.add_argument(
        "--perspective",
        type=float,
        default=0.0,
        help="轻微透视：四角扰动约为 min(靶宽,靶高)×该系数，0 关闭（建议 0.02~0.06）",
    )
    ap.add_argument(
        "--perspective-prob",
        type=float,
        default=0.75,
        help="每张图应用透视的概率 0~1",
    )
    ap.add_argument(
        "--color-jitter",
        type=float,
        default=0.0,
        help="合成后整图颜色抖动强度 0~1，0 关闭（建议 0.4~0.8）",
    )
    ap.add_argument(
        "--triangles-json",
        default=None,
        metavar="PATH",
        help="三角形顶点 JSON（test/archery_triangles_default.json）；坐标与 --fg 原图一致，"
        "多三角形时每张图写多个 VOC <object>（透视时顶点同步变换）",
    )
    ap.add_argument(
        "--triangle-bbox-pad-frac",
        type=float,
        default=0.0,
        help="三角形检测框在紧 AABB 四周再加 max(宽,高)×该比例（VOC/YOLO 同步）；"
        "0=贴顶点外接框；Stage2 建议 0.08~0.18，与推理端 margin 接近更易对齐",
    )
    ap.add_argument(
        "--stage2-crop",
        action="store_true",
        help="合成与增强后按靶子外接框+随机边距裁剪，输出与 Stage2（整靶 ROI）构图一致；标注为裁剪后坐标",
    )
    ap.add_argument(
        "--stage2-pad-min",
        type=float,
        default=0.02,
        help="Stage2 裁剪：四边 padding 相对靶 max(宽,高) 的比例下限",
    )
    ap.add_argument(
        "--stage2-pad-max",
        type=float,
        default=0.14,
        help="Stage2 裁剪：padding 比例上限",
    )
    ap.add_argument(
        "--stage2-allow-partial",
        action="store_true",
        help="有 --triangles-json 时允许裁剪后有效三角形数少于 JSON（默认要求数量一致）",
    )
    args = ap.parse_args()

    try:
        import cv2
        import numpy as np
    except ImportError:
        print("[ERR] 需要 opencv-python、numpy")
        sys.exit(1)

    rng = random.Random(args.seed)

    bgs = _collect_images(args.bg_dir)
    if not bgs:
        print(f"[ERR] 背景目录无图片: {args.bg_dir}")
        sys.exit(1)

    fg_path = args.fg
    if not os.path.isfile(fg_path):
        print(f"[ERR] 找不到靶图: {fg_path}")
        sys.exit(1)

    fg = cv2.imread(fg_path, cv2.IMREAD_UNCHANGED)
    if fg is None:
        print(f"[ERR] 无法读取靶图: {fg_path}")
        sys.exit(1)
    if fg.ndim == 2:
        fg = cv2.cvtColor(fg, cv2.COLOR_GRAY2BGRA)
    elif fg.shape[2] == 3:
        b, g, r = cv2.split(fg)
        a = np.full_like(b, 255)
        fg = cv2.merge([b, g, r, a])

    fx0, fy0, fw0, fh0 = _fg_bbox_from_alpha(fg)
    fg_crop = fg[fy0 : fy0 + fh0, fx0 : fx0 + fw0].copy()

    triangles_full = None
    if args.triangles_json:
        tpath = args.triangles_json
        if not os.path.isfile(tpath):
            print(f"[ERR] 找不到 --triangles-json: {tpath}")
            sys.exit(1)
        try:
            triangles_full = _load_triangles_json(tpath)
        except (json.JSONDecodeError, ValueError, OSError) as e:
            print(f"[ERR] 解析三角形 JSON 失败: {e}")
            sys.exit(1)
        print(f"[INFO] 已加载 {len(triangles_full)} 个三角形（每张图多个 VOC 检测框）")

    want_voc = args.format in ("voc", "both")
    want_yolo = args.format in ("yolo", "both")
    n_gen = min(max(0, args.num), max(0, args.max_images))
    if args.num > args.max_images:
        print(f"[INFO] --num={args.num} 大于 --max-images={args.max_images}，仅生成 {n_gen} 张")

    if args.stage2_crop:
        print(
            f"[INFO] Stage2 裁剪: pad∈[{args.stage2_pad_min},{args.stage2_pad_max}]×max(靶宽,靶高)，"
            f"partial={'允许' if args.stage2_allow_partial else '不允许'}"
        )

    out_img = os.path.join(args.out, "images")
    out_xml = os.path.join(args.out, "xml")
    out_lbl = os.path.join(args.out, "labels")
    os.makedirs(out_img, exist_ok=True)
    if want_voc:
        os.makedirs(out_xml, exist_ok=True)
    if want_yolo:
        os.makedirs(out_lbl, exist_ok=True)

    print(f"[INFO] 背景 {len(bgs)} 张，格式={args.format}，生成 {n_gen} 张 → {args.out}")

    i_done = 0
    while i_done < n_gen:
        bg_path = rng.choice(bgs)
        bg = cv2.imread(bg_path, cv2.IMREAD_COLOR)
        if bg is None:
            continue
        bh, bw = bg.shape[:2]
        short_bg = min(bh, bw)
        short_fg = min(fh0, fw0)
        smin = args.scale_min * short_bg / max(short_fg, 1)
        smax = args.scale_max * short_bg / max(short_fg, 1)
        scale = rng.uniform(max(smin, 0.05), max(smax, smin + 0.01))

        new_w = max(4, int(fw0 * scale))
        new_h = max(4, int(fh0 * scale))
        fg_s = cv2.resize(fg_crop, (new_w, new_h), interpolation=cv2.INTER_AREA)

        persp_M = None
        if args.perspective > 0 and rng.random() < args.perspective_prob:
            fg_s, persp_M = _perspective_warp_rgba(fg_s, args.perspective, rng, np, cv2)

        fw2, fh2 = fg_s.shape[1], fg_s.shape[0]
        tx0, ty0, tw, th = _fg_bbox_from_alpha(fg_s)

        max_x = max(0, bw - fw2)
        max_y = max(0, bh - fh2)
        px = rng.randint(0, max_x) if max_x > 0 else 0
        py = rng.randint(0, max_y) if max_y > 0 else 0

        comp = bg.copy()
        _paste_fg_on_bg(comp, px, py, fg_s)

        # 标注：整靶 alpha 框（无 triangles-json 时使用）或多三角形框
        bx0 = px + tx0
        by0 = py + ty0
        bx1 = px + tx0 + tw
        by1 = py + ty0 + th
        bx0 = max(0, min(bx0, bw - 1))
        by0 = max(0, min(by0, bh - 1))
        bx1 = max(bx0 + 1, min(bx1, bw))
        by1 = max(by0 + 1, min(by1, bh))

        tri_pts_full: list[np.ndarray] = []
        if triangles_full is not None:
            for tri in triangles_full:
                pts_c = _warp_triangle_points(
                    tri,
                    float(fx0),
                    float(fy0),
                    float(fw0),
                    float(fh0),
                    new_w,
                    new_h,
                    persp_M,
                    px,
                    py,
                    np,
                    cv2,
                )
                tri_pts_full.append(pts_c)

        if args.color_jitter > 1e-6:
            comp = _color_jitter_bgr(comp, args.color_jitter, rng, np, cv2)

        if args.blur_max > 1e-6:
            sig = rng.uniform(0.3, args.blur_max)
            k = int(sig * 4) | 1
            comp = cv2.GaussianBlur(comp, (k, k), sig)

        if rng.random() < max(0.0, min(1.0, float(args.motion_prob))):
            comp = _motion_blur_bgr(
                comp,
                rng,
                args.motion_kernel_min,
                args.motion_kernel_max,
                np,
                cv2,
            )

        bh, bw = comp.shape[:2]

        if args.stage2_crop:
            win = _stage2_crop_window(
                bx0,
                by0,
                bx1,
                by1,
                bw,
                bh,
                args.stage2_pad_min,
                args.stage2_pad_max,
                rng,
            )
            if win is None:
                continue
            cx0, cy0, cw, ch = win
            comp = comp[cy0 : cy0 + ch, cx0 : cx0 + cw].copy()
            out_w, out_h = cw, ch

            if triangles_full is not None:
                voc_objects = []
                yolo_lines_list = []
                for pts_c in tri_pts_full:
                    p2 = pts_c.copy()
                    p2[:, 0] -= cx0
                    p2[:, 1] -= cy0
                    pair = _triangle_to_voc_tuple(
                        p2,
                        out_w,
                        out_h,
                        args.class_name,
                        args.triangle_bbox_pad_frac,
                    )
                    if pair is None:
                        continue
                    vo, xyxy = pair
                    voc_objects.append(vo)
                    if want_yolo:
                        yolo_lines_list.append(
                            _yolo_line(args.class_id, xyxy, out_w, out_h)
                        )
                if not args.stage2_allow_partial and len(voc_objects) != len(
                    triangles_full
                ):
                    continue
                if want_voc and not voc_objects:
                    continue
                if want_yolo and not yolo_lines_list:
                    continue
            else:
                nbx0, nby0 = bx0 - cx0, by0 - cy0
                nbx1, nby1 = bx1 - cx0, by1 - cy0
                nbx0 = max(0, min(nbx0, out_w - 1))
                nby0 = max(0, min(nby0, out_h - 1))
                nbx1 = max(nbx0 + 1, min(nbx1, out_w))
                nby1 = max(nby0 + 1, min(nby1, out_h))
                if nbx1 <= nbx0 or nby1 <= nby0:
                    continue
                vx0, vy0, vx1, vy1 = _xyxy_exclusive_to_voc_inclusive(
                    nbx0, nby0, nbx1, nby1, out_w, out_h
                )
                voc_objects = [(args.class_name, vx0, vy0, vx1, vy1)]
                yolo_lines_list = (
                    [_yolo_line(args.class_id, (nbx0, nby0, nbx1, nby1), out_w, out_h)]
                    if want_yolo
                    else []
                )
        else:
            out_w, out_h = bw, bh
            if triangles_full is not None:
                voc_objects = []
                yolo_lines_list = []
                for pts_c in tri_pts_full:
                    pair = _triangle_to_voc_tuple(
                        pts_c,
                        out_w,
                        out_h,
                        args.class_name,
                        args.triangle_bbox_pad_frac,
                    )
                    if pair is None:
                        continue
                    vo, xyxy = pair
                    voc_objects.append(vo)
                    if want_yolo:
                        yolo_lines_list.append(
                            _yolo_line(args.class_id, xyxy, out_w, out_h)
                        )
                if want_voc and not voc_objects:
                    continue
                if want_yolo and not yolo_lines_list:
                    continue
            else:
                vx0, vy0, vx1, vy1 = _xyxy_exclusive_to_voc_inclusive(
                    bx0, by0, bx1, by1, out_w, out_h
                )
                voc_objects = [(args.class_name, vx0, vy0, vx1, vy1)]
                yolo_lines_list = (
                    [_yolo_line(args.class_id, (bx0, by0, bx1, by1), out_w, out_h)]
                    if want_yolo
                    else []
                )

        stem = f"synth_{i_done:06d}"
        img_name = stem + ".jpg"
        img_path = os.path.join(out_img, img_name)
        cv2.imwrite(img_path, comp, [int(cv2.IMWRITE_JPEG_QUALITY), args.jpeg_quality])

        if want_voc:
            xml_path = os.path.join(out_xml, stem + ".xml")
            _write_pascal_voc_xml(
                xml_path,
                img_filename=img_name,
                img_folder="images",
                img_w=out_w,
                img_h=out_h,
                depth=3,
                objects=voc_objects,
            )
        if want_yolo:
            lbl_path = os.path.join(out_lbl, stem + ".txt")
            with open(lbl_path, "w", encoding="utf-8") as f:
                f.writelines(yolo_lines_list)

        i_done += 1
        if i_done % 50 == 0:
            print(f"  ... {i_done}/{n_gen}")

    parts = [out_img]
    if want_voc:
        parts.append(out_xml)
    if want_yolo:
        parts.append(out_lbl)
    print(f"[OK] 完成: " + " , ".join(parts))

    if args.zip:
        if not want_voc:
            print("[WARN] --zip 需要 VOC 标注目录 xml/，当前格式未生成 xml，跳过打包")
        else:
            try:
                _zip_images_xml(args.out, args.zip)
                print(f"[OK] 已打包: {os.path.abspath(args.zip)}")
            except OSError as e:
                print(f"[ERR] 打包失败: {e}")
                sys.exit(1)


if __name__ == "__main__":
    main()