#!/usr/bin/env python3
"""qwen3.8-9b (LocalAI 192.168.3.210:18080) 多模态自动标注 → YOLO 格式

类别: pheasant(0) hare(1) dove(2) fish(3) cover(4, 生境区域)

标注两路:
- 动物框(红): 描述门控(2 采样, 均明确"无鸟"才判负) + 定位共识(3 采样, 中位数),
  样本分歧大(两两 IoU < 0.5) 的图像进复核清单
- 生境框(黄, cover): 植被覆盖地带识别 + NMS 去重

经验(2026-08-20 实测):
- 模型对"有无野鸡"的布尔判断有"有"偏置, 全图 3/3 答"有", 不可用
- 描述模式能区分大部分正负样本(负样本会明确说"没有鸟类或动物"), 但个别图多次描述矛盾
- 模型对负样本图像会稳定幻觉出小框, 所以不能只用定位采样判断有无
- 定位提示词强调"完整包住/从头到尾"比"贴合身体"更准(原提示词框偏小偏右)
- 隐藏目标(只露头/尾巴)会被描述门控误判负, 判负后需定位保险(42 图实测)
- 模型可能多采样一致地错(系统性偏差, 41/44 图实测), 需双提示词交叉验证
- 框面积 <1% 多为局部误检, 但真目标(只露头)也可能很小, 一律进复核

用法:
  venv/bin/python auto_label.py --input datasets/images --output datasets/labels
  venv/bin/python auto_label.py --input datasets/images --output datasets/labels --habitat
  venv/bin/python auto_label.py --input datasets/images --output datasets/labels --animal
  venv/bin/python auto_label.py --input datasets/images --output datasets/labels --dry-run
"""
import argparse
import base64
import io
import json
import os
import re
import sys
import time
import urllib.request
from pathlib import Path

from PIL import Image

API_URL = os.environ.get("OBSERVER_AI_URL", "http://192.168.3.210:18080")
MODEL = os.environ.get("OBSERVER_AI_MODEL", "qwen3.8-9b")
MAX_SIDE = 800          # 预处理尺寸, 对齐 LocalAI 聊天界面(100K 上下文下约 600~700 tokens/图)
JPEG_QUALITY = 85
TIMEOUT = 300
MAX_RETRIES = 4
RETRY_BACKOFF = 3       # 秒, 指数退避
DESC_SAMPLES = 2        # 描述门控采样数
BOX_SAMPLES = 3         # 定位采样数
CROSS_SAMPLES = 2       # 交叉提示词采样数(双提示词交叉验证)
CONSENSUS_IOU = 0.5     # 定位共识: 低于此值=采样分歧大, 进复核
CONSENSUS_OK_IOU = 0.7  # 定位共识: 高于此值才算可信, 中间段=勉强一致, 进复核
CROSS_IOU = 0.5         # 双提示词共识框 IoU 低于此值=交叉不一致, 进复核
CENTER_AGREE = 0.08     # 小框对 IoU 敏感, 中心距离不超过此值也视为一致
MIN_ANIMAL_AREA = 0.002 # 动物框面积低于 0.2%=极小疑似误检, 进复核

# 项目范围: 仅野鸡 + 生境区域 cover(2026-08-20 确认, 不含野兔/斑鸠/鱼)
CLASSES = ["pheasant", "cover"]
CLASS_CN = {"pheasant": "野鸡（环颈雉）"}

PHEASANT_FEATURES = ("野鸡（环颈雉）识别特征（雄性个体）：黑脑袋、红色脸颊、白色颈环、细长尾羽；"
                     "体型似鸡，站立或行走姿态")

DESC_PROMPT = ("请客观描述这张图片的内容（50字以内）：画面里有什么？"
               "是否有任何鸟类或动物？如果看到鸟类或动物，请明确说出来。"
               f"注意：{PHEASANT_FEATURES}。"
               "野鸡可能藏在草丛灌木中，只露出头部或尾巴，这样也算看到野鸡，要明确说出来。")

ANIMAL_PROMPT_TMPL = ("图片中有{cn}。注意识别特征：{feat}。"
                      "给出完整包住{cn}的边界框，尽量贴合，不要切掉身体任何部分（含尾巴）。"
                      "如果{cn}被遮挡、只露出部分（头部/尾巴/局部身体），也要框住可见部分，"
                      "不要因为没有全身就漏标。"
                      "每只{cn}一个框。如果图片中没有{cn}，boxes 输出空数组。"
                      "只输出JSON: {{\"boxes\": [[ymin, xmin, ymax, xmax], ...]}}。"
                      "坐标必须是0到1之间的归一化小数，禁止输出像素坐标")

# 交叉验证提示词: 结构与主提示词不同, 用于拆穿"多采样一致的错"(系统性偏差)
ANIMAL_PROMPT2_TMPL = ("画面中可能有{cn}。先在脑海中定位：{cn}的头、身体、尾巴各在什么位置？"
                       "然后给出完整包住{cn}的边界框；如果只露出部分，框住可见部分。"
                       "只输出JSON: {{\"boxes\": [[ymin, xmin, ymax, xmax], ...]}}，没有则空数组。"
                       "坐标0~1，禁止像素坐标")

HABITAT_PROMPT = """你是野生动物观察辅助工具。请找出野鸡最可能藏身或出现的具体位置。

规则：
- 只标注【具体的可疑位置】（如浓密草丛、植被边缘、沟渠边、倒木旁），不要框大片地带
- 每个位置一个小框，框住该可疑处即可，框的面积适中（约占画面 2%~15%）
- 按可疑度从高到低排列，最多 3 个
- 综合植被密度、地形、光线判断：植被浓密、能藏身、光线被遮挡处优先
- 不要把裸露地面、道路、天空框进去
- 只输出JSON: {"regions": [[ymin, xmin, ymax, xmax], ...]}，没有则 {"regions": []}"""

NO_BIRD_RE = re.compile(r"(没有|未发现|未看到|没有任何|看不到|不见).{0,10}(鸟类|鸟|动物|野鸡|雉)")

opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def preprocess_image(path: Path) -> tuple[bytes, int, int]:
    """返回 (jpeg字节, 缩放后宽, 缩放后高)"""
    with Image.open(path) as im:
        im = im.convert("RGB")
        w, h = im.size
        if max(w, h) > MAX_SIDE:
            scale = MAX_SIDE / max(w, h)
            w, h = round(w * scale), round(h * scale)
            im = im.resize((w, h), Image.LANCZOS)
        buf = io.BytesIO()
        im.save(buf, format="JPEG", quality=JPEG_QUALITY)
        return buf.getvalue(), w, h


def call_vision(jpeg: bytes, prompt: str, max_tokens: int = 400,
                temperature: float = 0.1) -> str:
    b64 = base64.b64encode(jpeg).decode()
    payload = {
        "model": MODEL,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url",
                 "image_url": {"url": "data:image/jpeg;base64," + b64}},
            ],
        }],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    req = urllib.request.Request(
        f"{API_URL}/v1/chat/completions",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    last_err = None
    for attempt in range(MAX_RETRIES):
        try:
            with opener.open(req, timeout=TIMEOUT) as resp:
                data = json.load(resp)
            return data["choices"][0]["message"]["content"]
        except Exception as e:
            last_err = e
            time.sleep(RETRY_BACKOFF * (2 ** attempt))
    raise RuntimeError(f"请求失败: {last_err}")


def parse_json_box_list(content: str, key: str, disp_w: int, disp_h: int):
    """解析 {"key": [[ymin,xmin,ymax,xmax],...]} 返回 [(ymin,xmin,ymax,xmax)]

    模型偶发输出像素坐标(相对缩放后输入图), 自动按实际尺寸归一化
    """
    text = re.sub(r"^```(?:json)?\s*|\s*```$", "", content.strip(), flags=re.MULTILINE)
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m:
        raise ValueError(f"响应中无 JSON: {content[:200]}")
    data = json.loads(m.group(0))
    boxes = []
    for b in data.get(key, []):
        ymin, xmin, ymax, xmax = (float(v) for v in b)
        if max(ymin, xmin, ymax, xmax) > 1:
            # 模型偶发输出像素坐标(相对缩放后输入图), 超出部分裁剪到边界
            xmin, xmax = xmin / disp_w, xmax / disp_w
            ymin, ymax = ymin / disp_h, ymax / disp_h
        xmin, xmax = min(max(xmin, 0), 1), min(max(xmax, 0), 1)
        ymin, ymax = min(max(ymin, 0), 1), min(max(ymax, 0), 1)
        # 模型偶发输出坐标序颠倒的框
        if xmin > xmax:
            xmin, xmax = xmax, xmin
        if ymin > ymax:
            ymin, ymax = ymax, ymin
        if xmin >= xmax or ymin >= ymax:
            continue  # 裁剪后退化(整体越界)的框直接丢弃
        boxes.append((ymin, xmin, ymax, xmax))
    return boxes


def iou(a, b):
    ymin1, xmin1, ymax1, xmax1 = a
    ymin2, xmin2, ymax2, xmax2 = b
    iw = min(xmax1, xmax2) - max(xmin1, xmin2)
    ih = min(ymax1, ymax2) - max(ymin1, ymin2)
    if iw <= 0 or ih <= 0:
        return 0.0
    inter = iw * ih
    union = (xmax1 - xmin1) * (ymax1 - ymin1) + (xmax2 - xmin2) * (ymax2 - ymin2) - inter
    return inter / union if union > 0 else 0.0


def nms(boxes, thr=0.5):
    """按面积降序贪心去重(用于生境多框)"""
    kept = []
    for b in sorted(boxes, key=lambda x: (x[2] - x[0]) * (x[3] - x[1]), reverse=True):
        if all(iou(b, k) < thr for k in kept):
            kept.append(b)
    return kept


def boxes_agree(a, b, iou_thr=CROSS_IOU, center_thr=CENTER_AGREE):
    """小框 IoU 敏感(同位置小框 IoU 可能很低), 中心距离足够近也视为一致"""
    if iou(a, b) >= iou_thr:
        return True
    ca = ((a[1] + a[3]) / 2, (a[0] + a[2]) / 2)
    cb = ((b[1] + b[3]) / 2, (b[0] + b[2]) / 2)
    return max(abs(ca[0] - cb[0]), abs(ca[1] - cb[1])) <= center_thr


def median_box(boxes):
    """按坐标分量的中位数合成框(定位共识)"""
    return tuple(
        sorted(v)[len(v) // 2] for v in zip(*boxes)
    )


def consensus_box(nonempty: list[list]) -> tuple[float, tuple]:
    """投票法共识: 选出其他采样中同意率最高的框, 返回 (同意率, 框)

    同意 = boxes_agree(IoU 达标或中心距离足够近), 对小框友好(IoU 均值法会误伤小框)
    """
    if len(nonempty) == 1:
        return 0.0, nonempty[0][0]
    best = None
    for b in nonempty[0]:
        agree = [any(boxes_agree(b, o) for o in s) for s in nonempty[1:]]
        score = sum(agree) / len(agree)
        if best is None or score > best[0]:
            best = (score, b)
    return best


def annotate_animal(jpeg: bytes, cls: str, disp_w: int, disp_h: int) -> tuple[bool, list, str, str]:
    """返回 (检出, 框列表[(ymin,xmin,ymax,xmax)], 状态 ok|review, 复核原因)"""
    cn = CLASS_CN[cls]
    prompt = ANIMAL_PROMPT_TMPL.format(cn=cn, feat=PHEASANT_FEATURES)
    prompt2 = ANIMAL_PROMPT2_TMPL.format(cn=cn)

    # 1. 描述门控: 2 采样, 均明确无鸟才判负
    descs = [call_vision(jpeg, DESC_PROMPT, max_tokens=150) for _ in range(DESC_SAMPLES)]
    no_birds = [bool(NO_BIRD_RE.search(d)) for d in descs]
    gate_blocked = False
    if all(no_birds):
        # 门控保险: 判负后仍各跑 1 次主/交叉定位, 任一有框则门控不可信
        ins1 = parse_json_box_list(
            call_vision(jpeg, prompt, temperature=0.3), "boxes", disp_w, disp_h)
        ins2 = parse_json_box_list(
            call_vision(jpeg, prompt2, temperature=0.3), "boxes", disp_w, disp_h)
        if not ins1 and not ins2:
            return False, [], "ok", ""
        gate_blocked = True

    # 2. 定位: 主提示词 3 采样 + 交叉提示词 2 采样
    samples = [parse_json_box_list(call_vision(jpeg, prompt, temperature=0.3),
                                   "boxes", disp_w, disp_h) for _ in range(BOX_SAMPLES)]
    cross = [parse_json_box_list(call_vision(jpeg, prompt2, temperature=0.3),
                                 "boxes", disp_w, disp_h) for _ in range(CROSS_SAMPLES)]

    nonempty = [s for s in samples if s]
    if not nonempty:
        cn2 = [s for s in cross if s]
        if cn2 and (gate_blocked or len(cn2) >= CROSS_SAMPLES):
            _, b2 = consensus_box(cn2)
            return True, [b2], "review", "主提示词未检出但交叉提示词有框"
        return False, [], "ok", ""

    score, box = consensus_box(nonempty)
    if len(nonempty) < BOX_SAMPLES:
        return True, [box], "review", "定位采样检出不一致"

    # 3. 双提示词交叉验证: 拆穿多采样一致的系统性偏差
    cn2 = [s for s in cross if s]
    if cn2:
        _, b2 = consensus_box(cn2)
        if not boxes_agree(box, b2):
            return True, [box], "review", "双提示词交叉不一致"

    # 4. 共识分级: <0.5 分歧大, 0.5~0.7 勉强一致, 均进复核
    if score < CONSENSUS_IOU:
        return True, [box], "review", "定位采样分歧大"
    if score < CONSENSUS_OK_IOU:
        return True, [box], "review", "定位采样勉强一致"

    # 5. 小框复核: 面积 <1% 疑似局部误检
    if (box[2] - box[0]) * (box[3] - box[1]) < MIN_ANIMAL_AREA:
        return True, [box], "review", "框过小,疑似局部或误检"
    if gate_blocked:
        return True, [box], "review", "门控判负但定位有框"
    return True, [box], "ok", ""


MAX_HABITAT_BOXES = 3    # 生境只保留可疑度最高的 3 个位置
MIN_HABITAT_AREA = 0.01  # 太小(碎点, <1% 面积)对训练无意义
MAX_HABITAT_AREA = 0.25  # 太大(大面积地带)不是"可疑点"
MAX_HABITAT_SIDE = 0.7   # 全宽/全高条带排除

def annotate_habitat(jpeg: bytes, disp_w: int, disp_h: int) -> list:
    content = call_vision(jpeg, HABITAT_PROMPT, max_tokens=400)
    boxes = parse_json_box_list(content, "regions", disp_w, disp_h)
    boxes = nms(boxes, 0.5)          # 去重(模型偶发输出重复框)
    filtered = []
    for ymin, xmin, ymax, xmax in boxes:
        w, h = xmax - xmin, ymax - ymin
        area = w * h
        if area < MIN_HABITAT_AREA or area > MAX_HABITAT_AREA:
            continue
        if w > MAX_HABITAT_SIDE or h > MAX_HABITAT_SIDE:
            continue
        filtered.append((ymin, xmin, ymax, xmax))
    return filtered[:MAX_HABITAT_BOXES]


def to_yolo(class_id: int, boxes) -> str:
    lines = []
    for ymin, xmin, ymax, xmax in boxes:
        cx = min(max((xmin + xmax) / 2, 0), 1)
        cy = min(max((ymin + ymax) / 2, 0), 1)
        w = min(max(xmax - xmin, 0), 1)
        h = min(max(ymax - ymin, 0), 1)
        lines.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")
    return "\n".join(lines)


def main():
    ap = argparse.ArgumentParser(description="qwen 自动标注 → YOLO(动物+生境)")
    ap.add_argument("--input", default="datasets/images")
    ap.add_argument("--output", default="datasets/labels")
    ap.add_argument("--review-file", default="datasets/review.txt")
    ap.add_argument("--mode", choices=["both", "animal", "habitat"], default="both")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    img_dir = Path(args.input)
    out_dir = Path(args.output)
    review_file = Path(args.review_file)
    if not img_dir.is_dir():
        sys.exit(f"输入目录不存在: {img_dir}")

    class_ids = {name: i for i, name in enumerate(CLASSES)}
    images = sorted(p for p in img_dir.rglob("*")
                    if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp", ".bmp"))
    todo = []
    for p in images:
        rel = p.relative_to(img_dir)
        cls = rel.parts[0] if len(rel.parts) > 1 else ""
        if cls not in class_ids or cls == "cover":
            print(f"跳过: 未知类别目录 {rel}", file=sys.stderr)
            continue
        label_path = out_dir / rel.with_suffix(".txt")
        if not args.force and label_path.exists():
            existing = label_path.read_text().splitlines()
            if args.mode == "both":
                continue  # 完整模式,已有结果即跳过
            if args.mode == "animal" and any(
                    not l.startswith(f"{class_ids['cover']} ") for l in existing if l):
                continue  # 已有动物标注
            if args.mode == "habitat" and any(
                    l.startswith(f"{class_ids['cover']} ") for l in existing if l):
                continue  # 已有生境标注
        todo.append((p, cls, label_path))

    if args.limit > 0:
        todo = todo[:args.limit]
    if not todo:
        print("没有待标注的图片。")
        return
    print(f"待标注 {len(todo)} 张 [模式: {args.mode}]")
    if args.dry_run:
        return

    ok = no_target = fail = review = 0
    t0 = time.time()
    for i, (img_path, cls, label_path) in enumerate(todo, 1):
        rel = img_path.relative_to(img_dir)
        label_path.parent.mkdir(parents=True, exist_ok=True)
        try:
            jpeg, disp_w, disp_h = preprocess_image(img_path)
            cover_id = class_ids["cover"]
            status = "ok"
            # 单模式重跑时保留另一类已有标注
            existing = (label_path.read_text().splitlines()
                        if label_path.exists() else [])
            animal_lines, cover_lines = [], []
            if args.mode in ("both", "animal"):
                found, boxes, status, reason = annotate_animal(jpeg, cls, disp_w, disp_h)
                animal_lines = to_yolo(class_ids[cls], boxes).splitlines()
            if args.mode in ("both", "habitat"):
                hboxes = annotate_habitat(jpeg, disp_w, disp_h)
                cover_lines = to_yolo(cover_id, hboxes).splitlines()
            if args.mode == "animal":
                cover_lines = [l for l in existing
                               if l.startswith(f"{cover_id} ")]
            if args.mode == "habitat":
                animal_lines = [l for l in existing
                                if not l.startswith(f"{cover_id} ")]
            lines = animal_lines + cover_lines
            label_path.write_text("\n".join(lines))
            if status == "review":
                review += 1
                with review_file.open("a") as f:
                    f.write(f"{rel}\t{cls}\t{reason}\n")
                tag = f"复核[{reason[:10]}]"
            elif lines:
                ok += 1
                tag = "OK"
            else:
                no_target += 1
                tag = "无目标"
            print(f"[{i}/{len(todo)}] {rel}: {tag} ({len(lines)} 行)")
        except Exception as e:
            fail += 1
            with review_file.open("a") as f:
                f.write(f"{rel}\t{cls}\t失败: {e}\n")
            print(f"[{i}/{len(todo)}] {rel}: 失败 - {e}", file=sys.stderr)

    print(f"\n完成: 有目标 {ok} 张, 无目标 {no_target} 张, 需复核 {review} 张, "
          f"失败 {fail} 张, 耗时 {time.time() - t0:.0f}s")
    if review:
        print(f"复核清单: {review_file}")


if __name__ == "__main__":
    main()