import cv2
import numpy as np
import os
from ultralytics import YOLO

# ================= 配置 =================
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/60seg/exp3/weights/best.pt"
IMAGE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/分割60/class4/1"
OUT_DIR = "./outputs"

IMG_SIZE = 640
CONF_THRES = 0.25

# 多边形简化比例（点数控制核心参数）
EPSILON_RATIO = 0.001

IMG_EXTS = (".jpg", ".jpeg", ".png", ".bmp")

# -------- 保存开关 --------
SAVE_LABELS = True     # 是否保存 YOLO seg 标签
SAVE_VIS = True       # 是否保存可视化结果
# ======================================


def simplify_polygon(poly, epsilon_ratio):
    """使用 approxPolyDP 简化多边形"""
    poly = poly.astype(np.int32)
    perimeter = cv2.arcLength(poly, True)
    epsilon = epsilon_ratio * perimeter
    approx = cv2.approxPolyDP(poly, epsilon, True)
    return approx.reshape(-1, 2)


def extract_simplified_masks(result):
    """
    提取并简化 YOLO mask
    返回: [(cls_id, poly), ...]
    """
    simplified = []

    if result.masks is None:
        return simplified

    boxes = result.boxes

    for i, poly in enumerate(result.masks.xy):
        cls_id = int(boxes.cls[i])
        conf = float(boxes.conf[i])

        if conf < CONF_THRES:
            continue

        poly = simplify_polygon(poly, EPSILON_RATIO)

        if len(poly) < 3:
            continue

        simplified.append((cls_id, poly))

    return simplified


def save_yolo_seg_labels(masks, img_shape, save_path):
    """保存 YOLO segmentation 标签（无目标也生成空 txt）"""
    h, w = img_shape[:2]
    lines = []

    for cls_id, poly in masks:
        poly_norm = []
        for x, y in poly:
            poly_norm.append(f"{x / w:.6f}")
            poly_norm.append(f"{y / h:.6f}")

        lines.append(str(cls_id) + " " + " ".join(poly_norm))

    with open(save_path, "w") as f:
        if lines:
            f.write("\n".join(lines))


def draw_polygons(img, masks):
    """在图像上绘制 segmentation 多边形"""
    vis = img.copy()

    for cls_id, poly in masks:
        poly = poly.astype(np.int32)

        cv2.polylines(
            vis,
            [poly],
            isClosed=True,
            color=(0, 255, 0),
            thickness=2
        )

        x, y = poly[0]
        cv2.putText(
            vis,
            str(cls_id),
            (int(x), int(y)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (0, 255, 0),
            2
        )

    return vis


def run_folder_inference():
    # 输出目录
    out_lbl_dir = os.path.join(OUT_DIR, "labels")
    out_img_dir = os.path.join(OUT_DIR, "images")

    if SAVE_LABELS:
        os.makedirs(out_lbl_dir, exist_ok=True)
    if SAVE_VIS:
        os.makedirs(out_img_dir, exist_ok=True)

    # 加载模型（只一次）
    model = YOLO(MODEL_PATH)

    img_files = sorted([
        f for f in os.listdir(IMAGE_DIR)
        if f.lower().endswith(IMG_EXTS)
    ])

    print(f"📂 共检测 {len(img_files)} 张图片")

    for idx, img_name in enumerate(img_files, 1):
        img_path = os.path.join(IMAGE_DIR, img_name)
        img = cv2.imread(img_path)

        if img is None:
            print(f"⚠️ 跳过无法读取: {img_name}")
            continue

        results = model(
            img,
            imgsz=IMG_SIZE,
            conf=CONF_THRES,
            verbose=False
        )
        result = results[0]

        masks = extract_simplified_masks(result)
        base_name = os.path.splitext(img_name)[0]

        # ---------- 保存标签 ----------
        if SAVE_LABELS:
            label_path = os.path.join(out_lbl_dir, base_name + ".txt")
            save_yolo_seg_labels(masks, img.shape, label_path)

        # ---------- 保存可视化 ----------
        if SAVE_VIS:
            vis_img = draw_polygons(img, masks)
            vis_path = os.path.join(out_img_dir, img_name)
            cv2.imwrite(vis_path, vis_img)

        print(f"[{idx}/{len(img_files)}] ✅ {img_name}")

    print("🎉 推理完成")


if __name__ == "__main__":
    run_folder_inference()