# -*- coding: utf-8 -*-
"""
step1: yolo26s-seg 裁切 + 透视拉直 + 纯几何横屏(dx>dy 保留, dx<dy 逆时针 90°)
不做任何 OCR 方向判断、不删任何图、不做语种判断。
输出: cropped_raw/<cls>/<name>__c<conf>.jpg
"""
import os, time, cv2, numpy as np
from pathlib import Path
from ultralytics import YOLO

WEIGHTS  = "/home/martin/顾工交接/wzj/ultralytics/runs/segment/yolo26s_seg_grading/weights/best.pt"
SRC_DIR  = "/home/martin/顾工交接/wzj/data_all/test"
OUT_DIR  = "/home/martin/顾工交接/wzj/data_all/cropped_raw"
CONF     = 0.5
IMGSZ    = 640
DEVICE   = 0

CLASS_NAMES = ['PSA', 'BGS', 'BGS-AU', 'BGS-AUTO', 'BGS-ONLY-AUTO',
               'CGC', 'CGC-OLD', 'SGC', 'SGC-AUTO', 'SGC-OLD', 'TAG', 'Other']


def mask_to_corners(mask, min_area=30.0):
    m = (mask > 0).astype(np.uint8)
    cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts: return None
    c = max(cnts, key=cv2.contourArea)
    if cv2.contourArea(c) < min_area: return None
    hull = cv2.convexHull(c)
    peri = cv2.arcLength(hull, True)
    for eps in (0.01, 0.02, 0.03, 0.04, 0.06, 0.08, 0.10, 0.15):
        approx = cv2.approxPolyDP(hull, eps * peri, True)
        if len(approx) == 4:
            return approx.reshape(4, 2).astype(np.float32)
    return cv2.boxPoints(cv2.minAreaRect(hull)).astype(np.float32)


def order_corners(pts):
    pts = np.asarray(pts, dtype=np.float32).reshape(4, 2)
    r = np.zeros((4, 2), dtype=np.float32)
    s = pts.sum(1); d = np.diff(pts, 1).ravel()
    r[0]=pts[s.argmin()]; r[2]=pts[s.argmax()]
    r[1]=pts[d.argmin()]; r[3]=pts[d.argmax()]
    return r


def warp_perspective(img, corners):
    o = order_corners(corners)
    wt=np.linalg.norm(o[1]-o[0]); wb=np.linalg.norm(o[2]-o[3])
    hl=np.linalg.norm(o[3]-o[0]); hr=np.linalg.norm(o[2]-o[1])
    W=max(int(round((wt+wb)/2)),1); H=max(int(round((hl+hr)/2)),1)
    dst=np.array([[0,0],[W-1,0],[W-1,H-1],[0,H-1]],np.float32)
    return cv2.warpPerspective(img, cv2.getPerspectiveTransform(o, dst), (W, H))


def ensure_landscape(img):
    """纯几何横屏: dx>dy 保留; dx<dy(竖向)逆时针 90° → 横向"""
    h, w = img.shape[:2]
    if w < h:
        return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    return img


def main():
    Path(OUT_DIR).mkdir(parents=True, exist_ok=True)
    for n in CLASS_NAMES:
        Path(OUT_DIR, n).mkdir(exist_ok=True)

    model = YOLO(WEIGHTS)
    img_paths = []
    for ext in ("*.jpg", "*.jpeg", "*.png"):
        img_paths.extend(Path(SRC_DIR).rglob(ext))
    print(f"Found {len(img_paths)} images", flush=True)

    t0 = time.time()
    total = 0
    for img_path in img_paths:
        try:
            results = model.predict(source=str(img_path), conf=CONF, imgsz=IMGSZ,
                                      device=DEVICE, verbose=False, retina_masks=True)
            if not results: continue
            res = results[0]
            H, W = res.orig_img.shape[:2]
            if res.boxes is None or len(res.boxes) == 0: continue
            xyxy   = res.boxes.xyxy.cpu().numpy()
            clss   = res.boxes.cls.cpu().numpy().astype(int)
            confs  = res.boxes.conf.cpu().numpy()
            masks  = res.masks.data.cpu().numpy() if res.masks is not None else None
            best_i = int(np.argmax(confs))

            cropped = None
            if masks is not None and best_i < len(masks):
                m = (masks[best_i] > 0.5).astype(np.uint8) * 255
                if m.shape != (H, W):
                    m = cv2.resize(m, (W, H), interpolation=cv2.INTER_NEAREST)
                corners = mask_to_corners(m)
                if corners is not None:
                    warped = warp_perspective(res.orig_img, corners)
                    if warped is not None and warped.size > 0:
                        cropped = warped
            if cropped is None:
                box = xyxy[best_i]
                h, w = res.orig_img.shape[:2]
                x1, y1 = max(0, int(box[0])), max(0, int(box[1]))
                x2, y2 = min(w, int(box[2])), min(h, int(box[3]))
                if x2 > x1 and y2 > y1:
                    cropped = res.orig_img[y1:y2, x1:x2]
            if cropped is None: continue

            cropped = ensure_landscape(cropped)

            cls_id   = int(clss[best_i])
            cls_name = CLASS_NAMES[cls_id] if 0 <= cls_id < len(CLASS_NAMES) else f"cls{cls_id}"
            out_path = Path(OUT_DIR) / cls_name / f"{img_path.stem}__c{confs[best_i]:.2f}.jpg"
            cv2.imwrite(str(out_path), cropped, [cv2.IMWRITE_JPEG_QUALITY, 92])
            total += 1
            if total % 200 == 0:
                print(f"  {total} done, elapsed={time.time()-t0:.0f}s", flush=True)
        except Exception as e:
            print(f"  err {img_path.name}: {e}", flush=True)

    print(f"DONE  total={total}  time={time.time()-t0:.1f}s", flush=True)


if __name__ == "__main__":
    main()