import cv2
import numpy as np
from ultralytics import YOLO

# ====================== 用户配置 ======================
MODEL_PATH = 'best.pt'
IMAGE_PATH = '/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/train/1.jpg'  # 👈 修改为你的具体图像路径
OUTPUT_DIR = './output_images'

# 固定点（例如标定得到的理论位置）
FIXED_REF_POINT = (535.0, 605)# (x, y)，单位：像素
def calculate_scale(width_mm, width_px):
    """
    计算缩放因子（单位：mm/px）
    :param width_mm: 实际宽度（单位：毫米）
    :param width_px: 宽度的像素数量
    :return: 缩放因子（单位：mm/px）
    """
    if width_px == 0:
        print("像素宽度不能为0")
        return None
    return width_mm / float(width_px)

# 示例使用：
# 假设我们知道一个参考物体的实际宽度是50毫米，在图像中占据100个像素
width_mm = 70.0  # 实际宽度（单位：毫米）
width_px = 42 # 在图像中的宽度（单位：像素）

SCALE_X= calculate_scale(width_mm, width_px)
print(f"水平方向的缩放因子为: {SCALE_X:.3f} mm/px")

def calculate_scale_y(height_mm, height_px):
    """
    计算垂直方向的缩放因子（单位：mm/px）
    :param height_mm: 实际高度（单位：毫米）
    :param height_px: 高度的像素数量
    :return: 缩放因子（单位：mm/px）
    """
    if height_px == 0:
        print("像素高度不能为0")
        return None
    return height_mm / float(height_px)

# 同样地，对于高度来说
height_mm = 890.0 # 实际高度（单位：毫米）
height_px = 507  # 在图像中的高度（单位：像素）

SCALE_Y = calculate_scale_y(height_mm, height_px)
print(f"垂直方向的缩放因子为: {SCALE_Y:.3f} mm/px")
# 创建输出目录
import os
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ====================== 可视化函数（增强版）======================
def draw_keypoints_and_offset(image, kpts_xy, kpts_conf, orig_shape, fixed_point, scale_x, scale_y):
    """
    在图像上绘制关键点、中心点、参考点、偏移箭头和文本
    :param image: OpenCV 图像
    :param kpts_xy: (N, K, 2) 坐标
    :param kpts_conf: (N, K) 置信度
    :param orig_shape: 原图尺寸 (H, W)
    :param fixed_point: 固定参考点 (fx, fy)
    :param scale_x: x方向缩放 mm/px
    :param scale_y: y方向缩放 mm/px
    :return: 处理后的图像，偏移信息列表
    """
    colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0)]  # 1红, 2蓝, 3绿, 4青
    results_info = []

    for i in range(len(kpts_xy)):
        xy = kpts_xy[i]  # (K, 2)
        conf = kpts_conf[i] if kpts_conf.ndim == 2 else kpts_conf[i:i+1]

        # 检查是否有至少两个关键点
        if len(xy) < 2:
            print(f"⚠️ 实例 {i} 的关键点数量不足2个")
            continue

        p1 = xy[0]  # 第一个关键点
        p2 = xy[1]  # 第二个关键点

        c1 = conf[0] if hasattr(conf, '__len__') else conf
        c2 = conf[1] if hasattr(conf, '__len__') else conf
        if c1 < 0.5 or c2 < 0.5:
            print(f"⚠️ 实例 {i} 的前两个关键点置信度过低: c1={c1:.3f}, c2={c2:.3f}")
            continue

        # 转为整数坐标（仅用于绘制）
        p1_int = tuple(map(int, p1))
        p2_int = tuple(map(int, p2))

        h, w = orig_shape
        valid = all(0 <= x < w and 0 <= y < h for x, y in [p1, p2])
        if not valid:
            print(f"⚠️ 实例 {i} 的关键点超出图像边界")
            continue

        # 绘制前两个关键点
        cv2.circle(image, p1_int, radius=15, color=colors[0], thickness=-1)  # 红色
        cv2.circle(image, p2_int, radius=15, color=colors[1], thickness=-1)  # 蓝色

        # 标注编号
        cv2.putText(image, "1", (p1_int[0] + 20, p1_int[1] - 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.5, colors[0], 5)
        cv2.putText(image, "2", (p2_int[0] + 20, p2_int[1] - 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.5, colors[1], 5)

        center_x = (p1[0] + p2[0]) / 2.0
        center_y = (p1[1] + p2[1]) / 2.0
        dynamic_center = (int(center_x), int(center_y))

        cv2.circle(image, dynamic_center, radius=18, color=(0, 255, 0), thickness=3)
        cv2.putText(image, "Center", (dynamic_center[0] + 30, dynamic_center[1]),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 3)

        fx, fy = map(int, fixed_point)
        cv2.circle(image, (fx, fy), radius=20, color=(255, 255, 0), thickness=3)
        cv2.putText(image, "Ref", (fx + 30, fy), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 0), 3)

        dx_px = center_x - fixed_point[0]
        dy_px = center_y - fixed_point[1]
        dx_mm = dx_px * scale_x
        dy_mm = dy_px * scale_y

        cv2.arrowedLine(image, (fx, fy), dynamic_center, (0, 255, 255), 3, tipLength=0.05)

        cv2.putText(image, f"ΔX={dx_mm:+.1f}mm", (fx + 40, fy - 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 3)
        cv2.putText(image, f"ΔY={dy_mm:+.1f}mm", (fx + 40, fy + 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 3)

        results_info.append({
            'instance': i,
            'center': (center_x, center_y),
            'dx_px': dx_px, 'dy_px': dy_px,
            'dx_mm': dx_mm, 'dy_mm': dy_mm
        })

    return image, results_info


if __name__ == "__main__":
    img = cv2.imread(IMAGE_PATH)
    if img is None:
        print(f"❌ 无法读取图像，检查路径: {IMAGE_PATH}")
        exit(1)

    model = YOLO(MODEL_PATH)
    results = model(img)

    for i, result in enumerate(results):
        if result.keypoints is not None:
            kpts = result.keypoints
            orig_shape = kpts.orig_shape

            kpts_xy = kpts.xy.cpu().numpy()
            kpts_conf = kpts.conf.cpu().numpy() if kpts.conf is not None else np.ones(kpts_xy.shape[:2])

            img_with_kpts = img.copy()

            img_with_kpts, offset_results = draw_keypoints_and_offset(
                img_with_kpts, kpts_xy, kpts_conf, orig_shape,
                fixed_point=FIXED_REF_POINT,
                scale_x=SCALE_X, scale_y=SCALE_Y
            )

            for info in offset_results:
                print(f"   📌 实例 {info['instance']}: "
                      f"ΔX={info['dx_mm']:+.2f}mm, ΔY={info['dy_mm']:+.2f}mm")

            save_filename = f"offset_{os.path.basename(IMAGE_PATH)}"
            save_path = os.path.join(OUTPUT_DIR, save_filename)
            cv2.imwrite(save_path, img_with_kpts)
            print(f"   💾 结果已保存: {save_path}")