import cv2
import numpy as np
import os

# 全局变量
drawing = False  # 是否正在绘制
ix, iy = -1, -1  # 起始点
roi_list = []    # 存储多个 ROI 坐标 [(x, y, w, h), ...]
image_path = "1.png"  # <<< 修改为你自己的图像路径
save_dir = "./roi_1/1.txt"  # 保存坐标的目录

# 创建保存目录
os.makedirs(save_dir, exist_ok=True)

def draw_rectangle(event, x, y, flags, param):
    global ix, iy, drawing, img_copy, roi_list

    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            # 每次移动都恢复原始图像，重新画矩形
            img_copy = img.copy()
            cv2.rectangle(img_copy, (ix, iy), (x, y), (0, 255, 0), 2)
            cv2.imshow("Select ROI", img_copy)

    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        w = x - ix
        h = y - iy
        if w != 0 and h != 0:
            # 确保宽高为正
            x_start = min(ix, x)
            y_start = min(iy, y)
            w = abs(w)
            h = abs(h)
            cv2.rectangle(img_copy, (x_start, y_start), (x_start + w, y_start + h), (0, 255, 0), 2)
            cv2.imshow("Select ROI", img_copy)
            # 添加到列表
            roi_list.append((x_start, y_start, w, h))
            print(f"已选择 ROI: (x={x_start}, y={y_start}, w={w}, h={h})")

# 保存坐标到 .txt 文件的函数
def save_rois_to_txt(rois, filepath):
    with open(filepath, 'w') as file:
        for roi in rois:
            # 将每个 ROI 转换为字符串并写入文件，每行一个 ROI
            line = ','.join(map(str, roi)) + '\n'
            file.write(line)
    print(f"💾 ROI 坐标已保存至: {filepath}")

def select_roi(image_path):
    global img, img_copy

    img = cv2.imread(image_path)
    if img is None:
        print(f"❌ 无法读取图像: {image_path}")
        return

    img_copy = img.copy()
    cv2.namedWindow("Select ROI")
    cv2.setMouseCallback("Select ROI", draw_rectangle)

    print("📌 使用鼠标左键拖拽选择 ROI")
    print("✅ 选择完成后按 's' 键保存坐标")
    print("⏭️  按 'n' 键跳过/下一步（可自定义）")
    print("🚪 按 'q' 键退出")

    while True:
        cv2.imshow("Select ROI", img_copy)
        key = cv2.waitKey(1) & 0xFF

        if key == ord('s'):
            # 保存坐标
            base_name = os.path.splitext(os.path.basename(image_path))[0]
            save_path = os.path.join(save_dir, f"{base_name}_rois1.txt")  # 修改了扩展名为 .txt
            save_rois_to_txt(roi_list, save_path)  # 使用新的保存函数

        elif key == ord('n'):
            print("⏭️  跳到下一张图片（此处可扩展）")
            break

        elif key == ord('q'):
            print("👋 退出程序")
            cv2.destroyAllWindows()
            return

    cv2.destroyAllWindows()

if __name__ == "__main__":
    select_roi(image_path)