# -*- coding: utf-8 -*-
"""
flat_grading.py - 平铺目录版评级检测(YOLO26 best.pt)
适配: 输入是平铺图片目录(无子文件夹), 图片可能是 jpg/png, 无 true_label(未知评级公司)

输入:  IN_DIR 下所有 jpg/png
输出:  OUT/crops/<name>.jpg   评级标签条裁剪(供OCR)
       OUT/predict/<name>     best.pt 标注图
       OUT/crop_records.json  裁剪记录
       OUT/grading.csv        评级结果(filename, 评级公司, is_graded, conf)
环境:  pytorch (需 CUDA_VISIBLE_DEVICES 指向 V100)
"""
import os, sys, json, csv, time, datetime
from pathlib import Path
import numpy as np
import cv2
sys.stdout.reconfigure(encoding="utf-8")

IN_DIR = os.environ.get("IN_DIR", "/home/martin/顾工交接/wzj/_infer_in")
OUT = os.environ.get("OUT_DIR", "/home/martin/顾工交接/wzj/_infer_out")
BEST = "/home/martin/顾工交接/wzj/ultralytics/runs/detect/card_v1/weights/best.pt"
CONF = 0.25

os.makedirs(os.path.join(OUT, "crops"), exist_ok=True)
os.makedirs(os.path.join(OUT, "predict"), exist_ok=True)

print("[STEP] 加载 best.pt ...")
sys.path.insert(0, "/home/martin/顾工交接/wzj")
from modules.ultralytics_compat import import_yolo
YOLO = import_yolo()
_t = time.perf_counter()
model = YOLO(BEST, task="detect")
t_load = time.perf_counter() - _t
print(f"[STEP] 加载完成 ({t_load:.2f}s)\n")

# 收集平铺目录所有图(jpg/jpeg/png)
imgs = []
for ext in ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"):
    imgs += list(Path(IN_DIR).glob(ext))
imgs = sorted(set(imgs))
print(f"[STEP] 待处理 {len(imgs)} 张 (平铺目录)\n", flush=True)

crop_records, csv_rows = [], []
T = {"infer": 0.0, "imread": 0.0, "write": 0.0}
_t_all = time.perf_counter()

for i, imgp in enumerate(imgs):
    name = imgp.name
    _t = time.perf_counter()
    results = model.predict(str(imgp), conf=CONF, imgsz=640, verbose=False)
    T["infer"] += time.perf_counter() - _t
    box, conf_val, pred_label = None, 0.0, "(无检测)"
    if results and results[0].boxes is not None and len(results[0].boxes) > 0:
        confs = results[0].boxes.conf.cpu().numpy()
        clses = results[0].boxes.cls.cpu().numpy().astype(int)
        xyxy = results[0].boxes.xyxy.cpu().numpy()
        bi = int(np.argmax(confs))
        conf_val = float(confs[bi]); box = tuple(float(v) for v in xyxy[bi])
        pred_label = model.names[clses[bi]]
    company = pred_label.split("-", 1)[0] if pred_label != "(无检测)" else ""
    is_graded = 1 if company else 0

    _t = time.perf_counter()
    img = cv2.imread(str(imgp))
    T["imread"] += time.perf_counter() - _t
    crop = None
    if box and img is not None:
        x1, y1, x2, y2 = box; h, w = img.shape[:2]
        x1i, y1i, x2i, y2i = max(0, int(x1)), max(0, int(y1)), min(w, int(x2)), min(h, int(y2))
        if x2i > x1i and y2i > y1i:
            crop = img[y1i:y2i, x1i:x2i]
    if crop is None:
        crop = img
    stem = imgp.stem
    crop_path = os.path.join(OUT, "crops", f"{stem}.jpg")
    _t = time.perf_counter()
    if crop is not None:
        cv2.imwrite(crop_path, crop)
        annotated = results[0].plot() if (results and results[0].boxes is not None and len(results[0].boxes) > 0) else img
        cv2.imwrite(os.path.join(OUT, "predict", f"{stem}.jpg"), annotated)
    T["write"] += time.perf_counter() - _t

    crop_records.append({"file": name, "stem": stem, "box": box,
                         "conf": round(conf_val, 4) if conf_val else None, "crop_path": crop_path})
    csv_rows.append({"filename": name, "评级公司": company or "非评级卡",
                     "is_graded": is_graded, "conf": round(conf_val, 4) if conf_val else ""})
    if (i + 1) % 200 == 0:
        print(f"  [{i+1}/{len(imgs)}] 完成", flush=True)

with open(os.path.join(OUT, "crop_records.json"), "w", encoding="utf-8") as f:
    json.dump(crop_records, f, ensure_ascii=False, indent=2)
with open(os.path.join(OUT, "grading.csv"), "w", encoding="utf-8-sig", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["filename", "评级公司", "is_graded", "conf"])
    w.writeheader(); w.writerows(csv_rows)

n = max(len(imgs), 1)
n_graded = sum(1 for r in csv_rows if r["is_graded"] == 1)
t_all = time.perf_counter() - _t_all
print(f"\n[RESULT] 评级卡 {n_graded} / 非评级卡 {n - n_graded}")
print("=" * 50)
print(f"⏱ flat_grading (YOLO26 best.pt) | {n} 张")
print(f"  模型加载 {t_load:.2f}s | best.pt推理 {T['infer']:.2f}s (均{T['infer']/n*1000:.1f}ms)")
print(f"  读图 {T['imread']:.2f}s | 写图 {T['write']:.2f}s | 总 {t_all:.2f}s ({n/t_all:.1f}张/s)")
print("=" * 50)
with open(os.path.join(OUT, "timing.log"), "a", encoding="utf-8") as lf:
    lf.write(f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] === 步骤① flat_grading.py (评级检测 YOLO26) ===\n")
    lf.write(f"  样本 {n} 张 | 评级卡 {n_graded} / 非评级卡 {n-n_graded}\n")
    lf.write(f"  模型加载 best.pt     : {t_load:.2f}s\n")
    lf.write(f"  best.pt 推理累计     : {T['infer']:.2f}s (均 {T['infer']/n*1000:.1f} ms/张)\n")
    lf.write(f"  读图/写图            : {T['imread']:.2f}s / {T['write']:.2f}s\n")
    lf.write(f"  总耗时               : {t_all:.2f}s (均 {t_all/n*1000:.1f} ms/张, 吞吐 {n/t_all:.1f} 张/s)\n")
print(f"✓ 输出: {OUT} | 计时→ {OUT}/timing.log")