"""
flat_card_no_detect.py - 步骤①：卡牌分割+编号检测+裁剪（pytorch环境）
输入：IN_DIR（平铺原图）
输出：OUT_DIR/card_no_crops/ + OUT_DIR/card_no_crop_records.json
用法：
  export IN_DIR=/path/to/images
  export OUT_DIR=/path/to/output
  export CUDA_VISIBLE_DEVICES=db470791-15e0-5fc2-3408-64ca8e9881dc
  ~/miniconda3/envs/pytorch/bin/python flat_card_no_detect.py
"""
import os, sys, json, time, datetime
from pathlib import Path
sys.stdout.reconfigure(encoding="utf-8")

IN_DIR = os.environ.get("IN_DIR", "/home/martin/顾工交接/wzj/_infer_in")
OUT_DIR = os.environ.get("OUT_DIR", "/home/martin/顾工交接/wzj/_infer_out")
MARGIN = float(os.environ.get("MARGIN", "0.05"))  # v6_medium_rec 在紧裁剪下已验证优于 v5+放宽margin
CROP_DIR = os.path.join(OUT_DIR, "card_no_crops")
os.makedirs(CROP_DIR, exist_ok=True)

sys.path.insert(0, "/home/martin/顾工交接/wzj")
import cv2
import config
from modules.yolo_detector import CardDetector
from modules.rectifier import CardRectifier
from modules.ultralytics_compat import import_yolo

# 模型路径
SEG_MODEL = "/home/martin/顾工交接/wzj/ultralytics/runs/segment/card_seg_pn_v2/weights/best.pt"
CARD_NO_MODEL = "/home/martin/顾工交接/wzj/ultralytics/runs/detect/yolo26s_card_no/weights/best.pt"

print("[INIT] 加载 card_seg_v2 + yolo26s_card_no ...")
_t = time.perf_counter()
seg_detector = CardDetector(SEG_MODEL, conf=0.25)
rectifier = CardRectifier(use_perspective=True)
YOLO = import_yolo()
card_no_model = YOLO(CARD_NO_MODEL, task="detect")
t_load = time.perf_counter() - _t
print(f"[INIT] 模型加载完成 ({t_load:.2f}s)\n")

# 读原图
imgs = []
for ext in ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"):
    imgs += list(Path(IN_DIR).glob(ext))
imgs = sorted(set(imgs))
print(f"[RUN] {len(imgs)} 张原图\n")

crop_records = []
t_detect = 0.0
_t_all = time.perf_counter()
for i, imgp in enumerate(imgs, 1):
    name = imgp.name
    try:
        img_bgr = cv2.imread(str(imgp))
        if img_bgr is None:
            continue
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        # card_seg_v2 分割+warp
        box, mask, orig_rgb = seg_detector.detect_box_mask(img_rgb)
        if box is None:
            crop_rgb = seg_detector.detect_and_crop(img_rgb)
        else:
            crop_rgb = rectifier.rectify(orig_rgb, mask=mask, box=box, angle=0)
        if crop_rgb is None or crop_rgb.size == 0:
            continue

        crop_bgr = cv2.cvtColor(crop_rgb, cv2.COLOR_RGB2BGR)
        h, w = crop_bgr.shape[:2]

        # yolo26s_card_no 检测编号区域
        # conf=0.3；多框时下方 argmax 取最高置信度
        res = card_no_model.predict(crop_bgr, conf=0.3, imgsz=640, verbose=False)[0]
        boxes = res.boxes
        if boxes is None or len(boxes) == 0:
            continue

        # 取置信度最高的框
        xyxy = boxes.xyxy.cpu().numpy()
        confs = boxes.conf.cpu().numpy()
        idx = int(confs.argmax())
        x1, y1, x2, y2 = xyxy[idx]
        conf = float(confs[idx])

        # 裁剪编号区域（margin 为检测框外扩比例，太紧会吃掉斜杠/截断右半）
        margin = MARGIN
        bw = x2 - x1
        bh = y2 - y1
        x1 = max(0, int(x1 - bw * margin))
        y1 = max(0, int(y1 - bh * margin))
        x2 = min(w, int(x2 + bw * margin))
        y2 = min(h, int(y2 + bh * margin))
        crop_num = crop_bgr[y1:y2, x1:x2]

        if crop_num.size == 0 or crop_num.shape[0] == 0 or crop_num.shape[1] == 0:
            continue

        # 保存裁剪图（文件名带 conf 前缀，便于按名称排序挑低置信度坏例）
        crop_path = os.path.join(CROP_DIR, f"{conf:.2f}__{name}")
        cv2.imwrite(crop_path, crop_num)
        crop_records.append({"file": name, "crop_path": crop_path, "conf": conf,
                             "box": [int(x1), int(y1), int(x2), int(y2)],
                             "crop_wh": [crop_num.shape[1], crop_num.shape[0]]})

        if i % 200 == 0:
            print(f"  [{i}/{len(imgs)}] 完成", flush=True)

    except Exception as e:
        print(f"  [ERR] {name}: {e}")

t_detect = time.perf_counter() - _t_all
print(f"\n[RESULT] {len(crop_records)} 张检测成功")
print("=" * 50)
print(f"⏱ flat_card_no_detect (分割+检测+裁剪) | {len(crop_records)} 张")
print(f"  模型加载 {t_load:.2f}s")
print(f"  检测+裁剪 {t_detect:.2f}s (均{t_detect/max(len(crop_records),1)*1000:.1f}ms)")
print("=" * 50)
with open(os.path.join(OUT_DIR, "timing.log"), "a", encoding="utf-8") as lf:
    lf.write(f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] === flat_card_no_detect.py (分割+检测+裁剪) ===\n")
    lf.write(f"  样本 {len(crop_records)} 张\n")
    lf.write(f"  模型加载 {t_load:.2f}s\n")
    lf.write(f"  检测+裁剪 {t_detect:.2f}s (均 {t_detect/max(len(crop_records),1)*1000:.1f} ms/张)\n")
print(f"✓ 裁剪图: {CROP_DIR} | 记录: {OUT_DIR}/card_no_crop_records.json | 计时→ timing.log")

# 保存裁剪记录
with open(os.path.join(OUT_DIR, "card_no_crop_records.json"), "w", encoding="utf-8") as f:
    json.dump(crop_records, f, ensure_ascii=False, indent=2)