#!/usr/bin/env python3
"""RF-DETR 整图批量检测野鸡,生成候选标注(初标)

- 整图等比缩放至最长边 INPUT_SIZE(700)送检测,**不做任何裁剪**,坐标按比例映射回原图
- 只保留 COCO bird 类;同目标重复框用 NMS 去重
- 藏身点:图像分析(植被密度/地形边缘/光线)生成 class 1 疑似藏匿位置
- 输出: YOLO txt (class 0 = 确认野鸡, class 1 = 疑似(低置信检测+藏身点))
       -> datasets/labels/pheasant_rfdetr/
- 分析JSON -> datasets/analysis/pheasant_rfdetr/
- 预览图 -> datasets/previews/pheasant_rfdetr/
"""
import argparse
import base64
import glob
import json
import os
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed

import cv2
import numpy as np

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
SRC_IMAGES = os.path.join(BASE_DIR, "datasets/images/pheasant")
OUT_LABELS = os.path.join(BASE_DIR, "datasets/labels/pheasant_rfdetr")
OUT_JSON = os.path.join(BASE_DIR, "datasets/analysis/pheasant_rfdetr")
OUT_PREVIEW = os.path.join(BASE_DIR, "datasets/previews/pheasant_rfdetr")

ENDPOINT = "http://192.168.3.210:18080/v1/detection"
MODEL = "rfdetr-xlarge"
TIMEOUT = 300
THRESHOLD = 0.08
CONF_CONFIRMED = 0.2  # 置信度高于此为明确(class 0 红框),否则为疑似(class 1 蓝框)
# RF-DETR 置信度分布:强目标 0.3-0.9,0.05-0.14 多为噪声尾;
# THRESHOLD=0.08 砍掉纯噪声,CONF_CONFIRMED=0.2 只把强目标标红,宁多勿漏
INPUT_SIZE = 700       # 整图等比缩放后的最长边,不裁剪
NMS_IOU = 0.4
WORKERS = 4
KEEP_CLASSES = set()  # 空 = 保留全部类别(RF-DETR 对小目标类别不稳定,定位一致即可)


IMG_MIME = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png"}


def detect(buf, mime):
    body = json.dumps({"model": MODEL, "image": f"data:{mime};base64,{base64.b64encode(buf.tobytes()).decode()}",
                       "threshold": THRESHOLD}).encode()
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    req = urllib.request.Request(ENDPOINT, data=body,
                                 headers={"Content-Type": "application/json"}, method="POST")
    with opener.open(req, timeout=TIMEOUT) as r:
        resp = json.loads(r.read().decode())
    return resp.get("detections", [])


def nms(boxes, iou_thr):
    """boxes: [(x, y, w, h, conf)], 按置信度降序去重"""
    boxes = sorted(boxes, key=lambda b: -b[4])
    keep = []
    for b in boxes:
        overlap = False
        for k in keep:
            ix = min(b[0] + b[2], k[0] + k[2]) - max(b[0], k[0])
            iy = min(b[1] + b[3], k[1] + k[3]) - max(b[1], k[1])
            if ix <= 0 or iy <= 0:
                continue
            inter = ix * iy
            union = b[2] * b[3] + k[2] * k[3] - inter
            if inter / union > iou_thr:
                overlap = True
                break
        if not overlap:
            keep.append(b)
    return keep


def find_hiding_spots(img, H, W, n_spots=3):
    """图像分析疑似藏身点:植被密度/地形边缘(田埂、林田交界、沟渠)/光线调制。

    野鸡习性:地面活动,藏身于茂密低矮遮蔽物,避开开阔裸地;
    强光场景偏好阴影遮蔽,弱光场景偏好相对亮且有遮蔽的区域。
    返回 [(x, y, w, h)] 像素坐标,框面积约占画面 2%~8%。
    """
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    s, v = hsv[..., 1], hsv[..., 2]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
    b = img[..., 0].astype(np.float32) / 255.0
    g = img[..., 1].astype(np.float32) / 255.0
    r = img[..., 2].astype(np.float32) / 255.0
    veg = np.clip(g - np.maximum(r, b), 0, 1)
    veg[s / 255.0 < 0.12] = 0
    tex = np.abs(gray - cv2.GaussianBlur(gray, (0, 0), 3))
    base = veg + tex
    grad = cv2.magnitude(cv2.Sobel(base, cv2.CV_32F, 1, 0, ksize=3),
                         cv2.Sobel(base, cv2.CV_32F, 0, 1, ksize=3))
    light = v / 255.0
    mean_light = float(gray.mean())
    if mean_light > 0.55:
        light_w = 1.0 - np.abs(light - 0.25)
    else:
        light_w = 1.0 - np.abs(light - 0.6)
    cover = (veg * 0.55 + tex * 0.30 + grad * 0.15) * (0.4 + 0.6 * light_w)
    cover[light < 0.12] = 0
    cover[light > 0.93] *= 0.15
    cover = cv2.GaussianBlur(cover, (0, 0), 7)
    windows = []
    for area_frac in (0.02, 0.04, 0.07):
        for aspect in (1.0, 1.6, 2.5):
            bw = int(round((area_frac * aspect * H * W) ** 0.5))
            bh = int(round((area_frac * H * W / aspect) ** 0.5))
            bw, bh = min(bw, W), min(bh, H)
            windows.append((cv2.boxFilter(cover, -1, (bw, bh)), bw, bh))
    stack = np.stack([w[0] for w in windows])
    best = stack.max(axis=0)
    best_idx = stack.argmax(axis=0)
    floor = max(float(best.max()) * 0.35, 0.02)
    spots = []
    for _ in range(n_spots):
        idx = np.unravel_index(int(np.argmax(best)), best.shape)
        y, x = idx
        if best[y, x] < floor:
            break
        _, bw, bh = windows[int(best_idx[y, x])]
        x0 = max(0, min(x - bw // 2, W - bw))
        y0 = max(0, min(y - bh // 2, H - bh))
        spots.append((x0, y0, bw, bh))
        mx, my = max(bw, 64), max(bh, 64)
        best[max(0, y - my):min(H, y + my), max(0, x - mx):min(W, x + mx)] = -1
    return [(int(x), int(y), int(w), int(h)) for x, y, w, h in spots]


def suppress_overlap(lines):
    """class 1(疑似)框与任一 class 0(确认)框重叠时删除,只保留 class 0"""
    cls0 = [l for l in lines if l.startswith("0 ")]
    if not cls0:
        return lines
    boxes0 = [tuple(map(float, l.split())) for l in cls0]
    keep1 = []
    for l in lines:
        if l.startswith("0 "):
            continue
        c, cx, cy, bw, bh = map(float, l.split())
        x1, y1, x2, y2 = cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2
        overlap = False
        for kc, kx, ky, kw, kh in boxes0:
            ix = min(x2, kx + kw / 2) - max(x1, kx - kw / 2)
            iy = min(y2, ky + kh / 2) - max(y1, ky - kh / 2)
            if ix > 0 and iy > 0:
                overlap = True
                break
        if not overlap:
            keep1.append(l)
    return cls0 + keep1


def process(name):
    path = os.path.join(SRC_IMAGES, name)
    img = cv2.imread(path)
    H, W = img.shape[:2]
    ext = os.path.splitext(path)[1].lower()
    mime = IMG_MIME.get(ext, "image/jpeg")
    jpeg_params = [cv2.IMWRITE_JPEG_QUALITY, 100] if ext in (".jpg", ".jpeg") else []
    # 整图等比缩放至最长边 INPUT_SIZE,不裁剪;检测坐标按比例映射回原图
    scale = min(1.0, INPUT_SIZE / max(H, W))
    if scale < 1.0:
        small = cv2.resize(img, (round(W * scale), round(H * scale)),
                           interpolation=cv2.INTER_AREA)
    else:
        small = img
    ok, buf = cv2.imencode(ext, small, jpeg_params)
    boxes = []
    for d in detect(buf, mime):
        if not KEEP_CLASSES or d.get("class_name") in KEEP_CLASSES:
            boxes.append((d["x"] / scale, d["y"] / scale,
                          d["width"] / scale, d["height"] / scale, d["confidence"]))
    # 野鸡目标 20-100px(近大远小);排除大物体误检(田地纹理)与细条误检
    boxes = [b for b in nms(boxes, NMS_IOU)
             if 4 <= b[3] <= H * 0.06 and 4 <= b[2] <= W * 0.1
             and 0.3 <= b[2] / b[3] <= 3.0]
    hiding = find_hiding_spots(img, H, W)
    return name, W, H, boxes, hiding


def reclassify():
    """不重新检测:按当前 CONF_CONFIRMED 从分析 JSON 重出 txt 标签与预览"""
    total_boxes = total_empty = 0
    for p in sorted(glob.glob(os.path.join(OUT_JSON, "*.json"))):
        j = json.load(open(p, encoding="utf-8"))
        W, H = j["size"]
        stem = os.path.splitext(os.path.basename(p))[0]
        lines = []
        for b in j["candidates"]:
            x, y, w, h, conf = b["x"], b["y"], b["width"], b["height"], b["confidence"]
            cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
            cls = 0 if conf >= CONF_CONFIRMED else 1
            lines.append(f"{cls} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
        for b in j.get("hiding_spots", []):
            x, y, w, h = b["x"], b["y"], b["width"], b["height"]
            cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
            lines.append(f"1 {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
        lines = suppress_overlap(lines)
        with open(os.path.join(OUT_LABELS, f"{stem}.txt"), "w", encoding="utf-8") as f:
            f.write("\n".join(lines) + ("\n" if lines else ""))
        draw_preview(os.path.join(SRC_IMAGES, j["image"]), lines,
                     os.path.join(OUT_PREVIEW, f"{stem}.jpg"))
        total_boxes += len(lines)
        if not lines:
            total_empty += 1
    print(f"重分类完成: {total_boxes} 框, 空图 {total_empty} 张", flush=True)


def main():
    ap = argparse.ArgumentParser(description="RF-DETR 切块放大批量检测野鸡候选框")
    ap.add_argument("--names", type=str, default="", help="只处理指定编号,逗号分隔")
    ap.add_argument("--reclassify", action="store_true",
                    help="不重新检测,按当前 CONF_CONFIRMED 从分析JSON重出标签与预览")
    args = ap.parse_args()

    if args.reclassify:
        reclassify()
        return

    for d in (OUT_LABELS, OUT_JSON, OUT_PREVIEW):
        os.makedirs(d, exist_ok=True)
    names = sorted(p.name for p in os.scandir(SRC_IMAGES)
                   if p.name.startswith("pheasant_") and p.name.endswith(".jpg"))
    if args.names:
        want = {f"pheasant_{n}.jpg" for n in args.names.split(",")}
        names = [n for n in names if n in want]

    total_boxes = total_empty = 0
    with ThreadPoolExecutor(max_workers=WORKERS) as pool:
        futures = [pool.submit(process, n) for n in names]
        for fut in as_completed(futures):
            name, W, H, boxes, hiding = fut.result()
            stem = os.path.splitext(name)[0]
            lines = []
            for x, y, w, h, conf in boxes:
                cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
                cls = 0 if conf >= CONF_CONFIRMED else 1
                lines.append(f"{cls} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
            for x, y, w, h in hiding:
                cx, cy, bw, bh = (x + w / 2) / W, (y + h / 2) / H, w / W, h / H
                lines.append(f"1 {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
            lines = suppress_overlap(lines)
            with open(os.path.join(OUT_JSON, f"{stem}.json"), "w", encoding="utf-8") as f:
                json.dump({"image": name, "size": [W, H],
                           "candidates": [{"x": b[0], "y": b[1], "width": b[2],
                                           "height": b[3], "confidence": b[4]} for b in boxes],
                           "hiding_spots": [{"x": b[0], "y": b[1], "width": b[2],
                                             "height": b[3]} for b in hiding]},
                          f, ensure_ascii=False, indent=1)
            with open(os.path.join(OUT_LABELS, f"{stem}.txt"), "w", encoding="utf-8") as f:
                f.write("\n".join(lines) + ("\n" if lines else ""))
            draw_preview(os.path.join(SRC_IMAGES, name), lines,
                         os.path.join(OUT_PREVIEW, f"{stem}.jpg"))
            total_boxes += len(lines)
            if not lines:
                total_empty += 1
            print(f"{name}: 候选{len(lines)}", flush=True)

    print(f"完成: {len(names)} 张, 候选框 {total_boxes}, 无候选 {total_empty} 张", flush=True)


def draw_preview(img_path, lines, out_path):
    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    for line in lines:
        cls, cx, cy, bw, bh = map(float, line.split())
        x1, y1 = int((cx - bw / 2) * w), int((cy - bh / 2) * h)
        x2, y2 = int((cx + bw / 2) * w), int((cy + bh / 2) * h)
        color = (0, 0, 255) if cls == 0 else (255, 0, 0)
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
    cv2.imwrite(out_path, img)


if __name__ == "__main__":
    main()