water_turbidity_det/roi_creator_mp.py

# -*- coding: utf-8 -*-
"""
代码一：定义数据并保存（支持多个ROI）

功能：
1. 加载图一。
2. 允许用户交互式地绘制多个多边形 ROI，并为每个ROI命名。
3. 允许用户交互式地绘制多个矩形屏蔽区。
4. 将所有ROI坐标、名称、屏蔽区坐标和相关参数打包保存到一个 JSON 文件中。
"""
import cv2
import numpy as np
import json
import os
import base64

# ========== 全局变量 ==========
# --- ROI 绘制相关 ---
current_roi_points = []
all_rois = []  # 存储所有ROI及其名称
drawing_complete = False
drawing_image = None
original_image = None
# --- 屏蔽区绘制相关 ---
ignore_polygons = []
drawing_ignore = None
drawing_rect = False
rect_start_point = None

# ========== 可配置参数 ==========
CONFIG_PARAMS = {
    "MARGIN_RATIO": 0.25,
    "MIN_MATCH_COUNT": 10,
    "USE_CLAHE": True
}


# -------------------------------------------------
# 1. 回调函数：选择 ROI
# -------------------------------------------------
def select_roi(event, x, y, flags, param):
    """用于绘制 ROI 多边形的回调函数"""
    global current_roi_points, drawing_complete, drawing_image

    if drawing_complete:
        return

    if event == cv2.EVENT_LBUTTONDOWN:
        current_roi_points.append((x, y))
        cv2.circle(drawing_image, (x, y), 4, (0, 255, 0), -1)
        if len(current_roi_points) > 1:
            cv2.line(drawing_image, current_roi_points[-2], current_roi_points[-1], (0, 255, 0), 2)
        cv2.imshow("Image 1 - Select ROI", drawing_image)

    elif event == cv2.EVENT_RBUTTONDOWN:
        if len(current_roi_points) < 3:
            print("[提示] ROI 至少需要 3 个点构成多边形!")
            return
        drawing_complete = True
        cv2.line(drawing_image, current_roi_points[-1], current_roi_points[0], (0, 255, 0), 2)
        cv2.imshow("Image 1 - Select ROI", drawing_image)
        print("[提示] 当前ROI选择完毕，按任意键继续...")


# -------------------------------------------------
# 2. 回调函数：选择屏蔽区
# -------------------------------------------------
def select_ignore_rect(event, x, y, flags, param):
    """用于绘制屏蔽矩形的回调函数"""
    global rect_start_point, drawing_rect, ignore_polygons, drawing_ignore

    if event == cv2.EVENT_LBUTTONDOWN:
        rect_start_point = (x, y)
        drawing_rect = True

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing_rect:
            preview_img = drawing_ignore.copy()
            cv2.rectangle(preview_img, rect_start_point, (x, y), (0, 0, 255), 2)
            cv2.imshow("Image 1 - Ignore Area", preview_img)

    elif event == cv2.EVENT_LBUTTONUP:
        if not drawing_rect:
            return
        drawing_rect = False
        x_start, y_start = rect_start_point
        x1, y1 = min(x_start, x), min(y_start, y)
        x2, y2 = max(x_start, x), max(y_start, y)

        if x2 - x1 > 5 and y2 - y1 > 5:
            cv2.rectangle(drawing_ignore, (x1, y1), (x2, y2), (0, 0, 255), 2)
            rect_as_poly = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
            ignore_polygons.append(rect_as_poly)
            print(f"[提示] 已新增屏蔽区 #{len(ignore_polygons)}，可继续拉框，完成后按 Enter 键。")
        cv2.imshow("Image 1 - Ignore Area", drawing_ignore)


# -------------------------------------------------
# 3. 工具函数
# -------------------------------------------------
def read_image_chinese(path):
    """读取可能包含中文路径的图片"""
    try:
        with open(path, 'rb') as f:
            data = np.frombuffer(f.read(), np.uint8)
        return cv2.imdecode(data, cv2.IMREAD_COLOR)
    except Exception as e:
        print(f"读取图片失败: {e}")
        return None


def draw_all_rois(image, rois_list):
    """在图像上绘制所有已保存的ROI"""
    result = image.copy()
    colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255), (255, 255, 0),
              (255, 0, 255), (0, 255, 255), (128, 255, 0), (255, 128, 0)]

    for idx, roi_data in enumerate(rois_list):
        color = colors[idx % len(colors)]
        points = roi_data["points"]
        name = roi_data["name"]

        # 绘制多边形
        cv2.polylines(result, [np.array(points, np.int32)], True, color, 2)

        # 计算中心点并显示名称
        center = np.mean(points, axis=0).astype(int)
        cv2.putText(result, name, tuple(center), cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, color, 2, cv2.LINE_AA)

    return result


# -------------------------------------------------
# 4. 主函数
# -------------------------------------------------
def define_and_save_data(image1_path, output_json_path):
    """
    主流程函数：加载图像，引导用户选择多个ROI并命名，最终保存数据。
    """
    global current_roi_points, drawing_complete, drawing_image, all_rois
    global drawing_ignore, ignore_polygons, original_image

    # --- 步骤 1: 加载图像 ---
    img1 = read_image_chinese(image1_path)
    if img1 is None:
        print(f"错误：无法加载图像 '{image1_path}'，请检查路径！")
        return

    original_image = img1.copy()

    # --- 步骤 2: 循环选择多个ROI ---
    print("\n========== 开始选择ROI ==========")
    roi_count = 0

    while True:
        roi_count += 1
        current_roi_points = []
        drawing_complete = False

        # 显示已有的ROI
        if all_rois:
            drawing_image = draw_all_rois(original_image, all_rois)
        else:
            drawing_image = original_image.copy()

        cv2.namedWindow("Image 1 - Select ROI")
        cv2.setMouseCallback("Image 1 - Select ROI", select_roi)

        print(f"\n【ROI #{roi_count}】")
        print("用鼠标左键多点圈出ROI，右键闭合完成当前ROI。")
        print("按 ESC 键结束所有ROI选择，按其他键继续添加下一个ROI。")
        cv2.imshow("Image 1 - Select ROI", drawing_image)

        key = cv2.waitKey(0)

        if key == 27:  # ESC键，结束ROI选择
            cv2.destroyWindow("Image 1 - Select ROI")
            if len(current_roi_points) >= 3:
                # 保存最后一个ROI
                roi_name = input(f"请输入ROI #{roi_count} 的名称: ").strip()
                if not roi_name:
                    roi_name = f"ROI_{roi_count}"
                all_rois.append({
                    "name": roi_name,
                    "points": current_roi_points
                })
                print(f"已保存ROI: {roi_name}")
            break
        else:
            if len(current_roi_points) >= 3:
                # 为当前ROI输入名称
                roi_name = input(f"请输入ROI #{roi_count} 的名称: ").strip()
                if not roi_name:
                    roi_name = f"ROI_{roi_count}"

                all_rois.append({
                    "name": roi_name,
                    "points": current_roi_points
                })
                print(f"已保存ROI: {roi_name}")
            else:
                print("当前ROI点数不足，跳过保存。")
                roi_count -= 1

    if not all_rois:
        print("错误：未选择任何ROI，程序退出。")
        return

    print(f"\n总共选择了 {len(all_rois)} 个ROI")

    # --- 步骤 3: 交互式选择屏蔽区 ---
    drawing_ignore = draw_all_rois(original_image, all_rois)
    cv2.namedWindow("Image 1 - Ignore Area")
    cv2.setMouseCallback("Image 1 - Ignore Area", select_ignore_rect)
    print("\n【步骤 2】在弹窗中，用鼠标左键拖拽拉框来选择屏蔽区，可画多个。")
    print("         所有屏蔽区画完后，按 Enter 键保存并结束，或按 Esc 键退出。")
    cv2.imshow("Image 1 - Ignore Area", drawing_ignore)

    while True:
        key = cv2.waitKey(0)
        if key in [13, 10]:  # Enter键
            break
        elif key == 27:  # Esc
            cv2.destroyAllWindows()
            print("用户取消操作。")
            return
    cv2.destroyAllWindows()

    # --- 步骤 4: 封装数据并保存到 JSON ---
    try:
        with open(image1_path, 'rb') as f:
            image_bytes = f.read()
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        image_filename = os.path.basename(image1_path)
    except Exception as e:
        print(f"\n[失败] 读取并编码图像文件时出错: {e}")
        return

    # 封装数据
    data_to_save = {
        "image1_filename": image_filename,
        "image1_data_base64": image_base64,
        "rois": all_rois,  # 保存所有ROI及其名称
        "ignore_polygons": ignore_polygons,
        "parameters": CONFIG_PARAMS
    }

    try:
        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(data_to_save, f, indent=4, ensure_ascii=False)
        print(f"\n[成功] {len(all_rois)} 个ROI、屏蔽区和图像数据已成功保存到: {output_json_path}")
    except Exception as e:
        print(f"\n[失败] 保存 JSON 文件时出错: {e}")


# -------------------------------------------------
if __name__ == "__main__":
    IMG1_PATH = r"20250615000002.jpg"
    OUTPUT_JSON = "roi_json_dir/matching_data_multi_roi.json"

    # 创建输出目录
    os.makedirs(os.path.dirname(OUTPUT_JSON), exist_ok=True)

    define_and_save_data(IMG1_PATH, OUTPUT_JSON)