pokouqiege/utils/data_process/zyh/rgbd/mask_xishu/mask.py

import os
import json
import shutil
import numpy as np
import cv2
from scipy.optimize import leastsq
from tqdm import tqdm


def calc_R(x, y, xc, yc):
    return np.sqrt((x - xc) ** 2 + (y - yc) ** 2)


def f(c, x, y):
    Ri = calc_R(x, y, *c)
    return Ri - Ri.mean()


def fit_arc_and_create_mask(points, shape, point_step_deg=0.5, radius=2):
    x_data = points[:, 0]
    y_data = points[:, 1]
    center_estimate = np.mean(x_data), np.mean(y_data)
    center, _ = leastsq(f, center_estimate, args=(x_data, y_data))
    xc, yc = center
    Ri = calc_R(x_data, y_data, xc, yc)
    R = Ri.mean()

    # 极角
    angles = np.arctan2(points[:, 1] - yc, points[:, 0] - xc)
    start_angle = np.min(angles)
    end_angle = np.max(angles)

    # 确保角度顺序正确
    if end_angle - start_angle > np.pi:
        start_angle, end_angle = end_angle, start_angle + 2 * np.pi

    arc_angles = np.arange(start_angle, end_angle, np.deg2rad(point_step_deg))
    arc_points = np.stack([
        xc + R * np.cos(arc_angles),
        yc + R * np.sin(arc_angles)
    ], axis=-1).astype(np.int32)

    # 不返回mask了，直接返回点列表
    return None, arc_points.tolist()


def build_shapes_from_points(points, label="arc"):
    points_np = np.array(points)
    if points_np.shape[0] < 3:
        return []

    xs = points_np[:, 0]
    ys = points_np[:, 1]

    shape = {
        "label": label,
        "points": points,
        "shape_type": "polygon",
        "flags": {},
        "xmin": int(xs.min()),
        "ymin": int(ys.min()),
        "xmax": int(xs.max()),
        "ymax": int(ys.max())
    }
    return [shape]


def process_folder_labelme(input_dir, output_dir, point_step_deg=0.5, radius=2):
    os.makedirs(output_dir, exist_ok=True)

    for file_name in tqdm(os.listdir(input_dir)):
        if not file_name.endswith(".json"):
            continue

        json_path = os.path.join(input_dir, file_name)
        image_path = json_path.replace(".json", ".jpg")
        if not os.path.exists(image_path):
            print(f"图像不存在，跳过：{image_path}")
            continue

        with open(json_path, "r") as f:
            label_data = json.load(f)

        points_list = [item['points'][0] for item in label_data['shapes'] if item['label'] == 'arc']
        if len(points_list) != 3:
            print(f"{file_name} 中 arc 点数不足 3，跳过")
            continue

        image = cv2.imread(image_path)
        h, w = image.shape[:2]
        points = np.array(points_list)
        _, arc_points = fit_arc_and_create_mask(points, (h, w), point_step_deg, radius)

        shapes = build_shapes_from_points(arc_points, label="arc")
        if len(shapes) == 0:
            print(f"{file_name} 拟合失败，跳过")
            continue

        output_json = {
            "version": "5.0.1",
            "flags": {},
            "shapes": shapes,
            "imagePath": os.path.basename(image_path),
            "imageHeight": h,
            "imageWidth": w
        }

        base_name = os.path.splitext(os.path.basename(image_path))[0]
        out_json_path = os.path.join(output_dir, base_name + ".json")
        out_img_path = os.path.join(output_dir, base_name + ".jpg")

        with open(out_json_path, "w") as f:
            json.dump(output_json, f, indent=4)

        shutil.copy(image_path, out_img_path)

    print(f"\n✅ 所有 arc 区域已保存为 labelme 格式（图像 + json）到目录：{output_dir}")


# ===== 修改为你自己的输入输出路径 =====
input_dir = r"G:\python_ws_g\data\arc"
output_dir = r"G:\python_ws_g\data\arcout"
point_step_deg = 0.5
draw_radius = 2

if __name__ == "__main__":
    process_folder_labelme(input_dir, output_dir, point_step_deg, draw_radius)