sam2segment_structure.py

import numpy as np
import sys,os
# sys.path.append("/home/yifan/sam2")
# sys.path.append("/data_sdf/yifan/miniconda3/envs/sam2/lib/python3.10/site-packages")
from huggingface_hub import hf_hub_download
sys.path.append(os.path.join(os.path.dirname(__file__), "sam2"))
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
import torch
import matplotlib.pyplot as plt
from PIL import Image
import cv2
import random
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
device = torch.device("cuda")
sam2_checkpoint = hf_hub_download(
    repo_id="Evan73/sam2-models",
    filename="sam2.1_hiera_large.pt"
)
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
# global sam2_model
predictor = SAM2ImagePredictor(sam2_model)
from ultralytics import YOLO
from diffusers.utils import load_image
import pickle
import os
import math
heatmap_zip = hf_hub_download(
    repo_id="Evan73/attention-heatmaps",
    filename="attention_heatmaps.zip"
)
import zipfile
import os

with zipfile.ZipFile(heatmap_zip, 'r') as zip_ref:
    zip_ref.extractall("heatmaps_lda")

with open("heatmaps_lda/attention_heatmaps.pkl", "rb") as f:
    heatmap_dict = pickle.load(f)

def load_yolov5_model():
    # 使用YOLOv5官方模型加载器（需要安装yolov5）
    # model = torch.hub.load('ultralytics/yolov11', 'yolov11s')  # 可以根据需要选择不同大小的模型
    model = YOLO("yolo11n.pt")
    class_names = model.names  # class index to name mapping
    print("YOLOv11 Class Names:")
    for idx, name in class_names.items():
        print(f"{idx}: {name}")
    return model

# 检查点是否在汽车区域内
def is_point_in_car_area(point, model, image):
    """
    检查给定的点是否在车辆区域内
    - point: 点的坐标 (x, y)
    - model: YOLO模型
    - image: 输入的图像
    """
    # 使用YOLO模型进行物体检测
    results = model(image)  # 获取检测结果
    
    # 获取汽车类别（根据模型调整类别ID）
    # print("Detected classes:", results[0].boxes.cls.cpu().numpy())
    car_class_id = [2, 5, 7]  # COCO数据集中汽车类别通常为2，但需确认
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # 遍历每个检测结果（支持批量处理，这里假设单张图像）
    for result in results:
        # 提取检测框的xyxy坐标、置信度、类别
        boxes = result.boxes.xyxy.cpu().numpy()  # 转换为左上和右下坐标
        confidences = result.boxes.conf.cpu().numpy()
        class_ids = result.boxes.cls.cpu().numpy().astype(int)
        
        # 遍历每个检测框
        for box, cls in zip(boxes, class_ids):
            if cls in car_class_id:
                x_min, y_min, x_max, y_max = box[:4]
                # 绘制检测框（可选）
                cv2.rectangle(image_bgr, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 2)
                # 检查点是否在框内
                if (x_min <= point[0] <= x_max) and (y_min <= point[1] <= y_max):
                    cv2.imwrite("yolo_res.jpg", image_bgr)
                    return False
    cv2.imwrite("yolo_res.jpg", image_bgr)
    print(f"检测结果已保存至 yolo_res.jpg")
    return True


def show_mask(mask, ax, image_path,random_color=False, borders=True, image=None, save_path=None):
    """
    根据mask区域随机选择两个对角点并在原始图像上绘制矩形框。
    
    参数：
    - `mask`: 掩码区域
    - `ax`: 用于绘制的matplotlib轴
    - `random_color`: 是否使用随机颜色
    - `borders`: 是否显示边界
    - `image`: 原始图像，用于绘制矩形框
    - `save_path`: 保存结果图像的路径
    """
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])

    h, w = mask.shape[-2:]
    mask = mask.astype(np.uint8)
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    cv2.imwrite("binary_mask.png", (mask * 255).astype(np.uint8))
    print("原始二值掩码已保存为 binary_mask.png")
    if borders:
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=5)
    # print(f"Mask unique values: {np.unique(mask)}")
    # print(f"Max value in mask: {mask.max()}, Min value in mask: {mask.min()}")
    # 如果提供了原始图像，绘制矩形框
    # size = 100
    colors = [
        (255, 0, 0),   # 红色
        (0, 255, 0),   # 绿色
        (0, 0, 255),   # 蓝色
        (255, 255, 0), # 黄色
        (255, 0, 255), # 品红色
        (0, 255, 255), # 青色
        (255, 128, 0), # 橙色
        (128, 0, 255), # 紫色
        (128, 128, 128), # 灰色
        (0, 128, 0)    # 深绿色
    ]

    for idx, contour in enumerate(contours):
        x, y, w, h = cv2.boundingRect(contour)
        print(f"轮廓{idx}: x={x}, y={y}, w={w}, h={h}")
        color = colors[idx % len(colors)]
        cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
    middle_save_path = "contours_colored_result.png"
    cv2.imwrite(middle_save_path, image)
    print(f"带颜色的轮廓结果已保存至 {middle_save_path}")
    if image is not None:
        # 找到掩码的边界
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            # print(x, y, w, h)
            if w > 50 and h > 50:
                for size in range(90,40,-5):
                    for _ in range(100):
                        random_x1 = random.randint(x, x + w - 50)
                        random_y1 = random.randint(y, y + h - 50)
                        random_x2 = random_x1 - size
                        random_y2 = random_y1 - size
                        # print(random_x1, random_y1,random_x2,random_y2)
                        # 在原图上绘制矩形框
                        # 保存结果图像
                        try:
                            if save_path and mask[random_y1, random_x1] == 1 and mask[random_y2, random_x2] == 1:
                                cv2.rectangle(image,(random_x2, random_y2), (random_x1, random_y1), (0, 255, 0), 2)
                                cv2.imwrite(save_path, image)
                                # generate_gt_mask_from_intersection([(random_x1, random_y1),(random_x2, random_y2)], yolo_boxes, image, sam2_model, threshold_iou=0.01)
                                print(f"Image with rectangle saved at {save_path}")
                                return (random_x1,random_y1),(random_x2,random_y2)
                        except:
                            pass
                            # cv2.rectangle(image,(random_x2, random_y2), (random_x1, random_y1), (0, 255, 0), 2)
                            # cv2.imwrite(save_path, image)
                            # print(f"Image with rectangle saved at {save_path}")
                            # break
                    for _ in range(100):
                        random_x1 = random.randint(x, x + w - 50)
                        random_y1 = random.randint(y, y + h - 50)
                        random_x2 = random_x1 + size
                        random_y2 = random_y1 + size
                        # print(mask[random_y1, random_x1] == 1,mask[random_y2, random_x2] == 1)
                        # 在原图上绘制矩形框
                        # 保存结果图像
                        try:
                            if save_path and mask[random_y1, random_x1] == 1 and mask[random_y2, random_x2] == 1:
                                cv2.rectangle(image,(random_x2, random_y2), (random_x1, random_y1), (0, 255, 0), 2)
                                cv2.imwrite(save_path, image)
                                print(f"Image with rectangle saved at {save_path}")
                                # generate_gt_mask_from_intersection([(random_x1, random_y1),(random_x2, random_y2)], yolo_boxes, image, sam2_model, threshold_iou=0.01)
                                return (random_x1,random_y1),(random_x2,random_y2)
                        except:
                            pass

    ax.imshow(mask_image)
    plt.axis('off')
    
    

def attention_mask(mask, ax, image_path,strategy="LOA",random_color=False, borders=True, image=None, save_path=None):
    """
    根据mask区域随机选择两个对角点并在原始图像上绘制矩形框。
    
    参数：
    - `mask`: 掩码区域
    - `ax`: 用于绘制的matplotlib轴
    - `random_color`: 是否使用随机颜色
    - `borders`: 是否显示边界
    - `image`: 原始图像，用于绘制矩形框
    - `save_path`: 保存结果图像的路径
    """
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    orig_w, orig_h = image.shape[1],image.shape[0]
    # print(image.shape)
    h, w = mask.shape[-2:]
    mask = mask.astype(np.uint8)
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    cv2.imwrite("binary_mask.png", (mask * 255).astype(np.uint8))
    print("原始二值掩码已保存为 binary_mask.png")
    # if borders:
    #     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    #     contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
    #     mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
    # colors = [
    #     (255, 0, 0),   # 红色
    #     (0, 255, 0),   # 绿色
    #     (0, 0, 255),   # 蓝色
    #     (255, 255, 0), # 黄色
    #     (255, 0, 255), # 品红色
    #     (0, 255, 255), # 青色
    #     (255, 128, 0), # 橙色
    #     (128, 0, 255), # 紫色
    #     (128, 128, 128), # 灰色
    #     (0, 128, 0)    # 深绿色
    # ]
    # # print(mask.shape)
    # for idx, contour in enumerate(contours):
    #     x, y, w, h = cv2.boundingRect(contour)
    #     print(f"轮廓{idx}: x={x}, y={y}, w={w}, h={h}")
    #     color = colors[idx % len(colors)]
    #     cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
    # middle_save_path = "contours_colored_result.png"
    # cv2.imwrite(middle_save_path, image)
    # print(f"带颜色的轮廓结果已保存至 {middle_save_path}")
    candidates = []
    path = image_path
    cls_heatmap = heatmap_dict[path]['cls_heatmap']
    reg_heatmap = heatmap_dict[path]['reg_heatmap']
    font = cv2.FONT_HERSHEY_SIMPLEX
    if strategy == "LDA":
       combined = cls_heatmap.astype(np.float32)
    if strategy == "LOA" or strategy == "LRA":
       combined = reg_heatmap.astype(np.float32)
    print(mask.shape)
    mask = cv2.resize(mask, (combined.shape[1], combined.shape[0]), interpolation=cv2.INTER_NEAREST)
    mask = (mask > 0.5).astype(np.uint8)
    cv2.imwrite("crop_binary_mask.png", (mask * 255).astype(np.uint8))
    print("处理后的裁剪二值掩码已保存为 crop_binary_mask.png")
    print(combined.shape)
    vis_image = cv2.imread(image_path)
    vis_image = cv2.resize(vis_image,(combined.shape[1],combined.shape[0]))
    mask_image =  cv2.resize(mask_image,(combined.shape[1],combined.shape[0]))
    image =  cv2.resize(image,(combined.shape[1],combined.shape[0]))
    if borders:
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
    colors = [
        (255, 0, 0),   # 红色
        (0, 255, 0),   # 绿色
        (0, 0, 255),   # 蓝色
        (255, 255, 0), # 黄色
        (255, 0, 255), # 品红色
        (0, 255, 255), # 青色
        (255, 128, 0), # 橙色
        (128, 0, 255), # 紫色
        (128, 128, 128), # 灰色
        (0, 128, 0)    # 深绿色
    ]
    # print(mask.shape)
    for idx, contour in enumerate(contours):
        x, y, w, h = cv2.boundingRect(contour)
        print(f"轮廓{idx}: x={x}, y={y}, w={w}, h={h}")
        color = colors[idx % len(colors)]
        cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
    middle_save_path = "contours_colored_result.png"
    cv2.imwrite(middle_save_path, image)
    print(f"带颜色的轮廓结果已保存至 {middle_save_path}")
    if image is not None:
        # 找到掩码的边界
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # print(contours)
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            print("the contour is:",x, y, w, h)
            if w > 50 and h > 50:
                for size in range(50,40,-5):
                    for y_step in range(y, y+h - size,5):
                        for x_step in range(x, x+w - size,5):
                            x1, y1, x2, y2 = x_step, y_step, x_step + size, y_step + size
                            # print(mask[y1:y2, x1:x2].sum())
                            if mask[y1:y2, x1:x2].sum() >= size * size:  # 掩码区域必须都在内部
                                heat_value = combined[y1:y2, x1:x2].mean()
                                # print("the heat_value is:",heat_value,y1,y2, x1,x2,combined.shape)
                                if not math.isnan(heat_value):
                                    candidates.append(((x1, y1, x2, y2), heat_value))
                                    cv2.rectangle(vis_image, (x1, y1), (x2, y2), (0, 255, 0), 1)
                                    cv2.putText(vis_image, f'{heat_value:.1f}', (x1, y1 - 2), font, 0.4, (0, 0, 255), 1)
                    if not candidates:
                        print("⚠️ 没有找到满足掩码内区域的候选框")
                    else:
                        break
                cv2.imwrite("attention_vis.jpg", vis_image)
                print(f"Attention 候选框可视化已保存 attention_vis.jpg")
                # 从高到低排序，选择热值最高的
                candidates.sort(key=lambda x: x[1], reverse=True)
                print(save_path,candidates[0],candidates[-1])
                for (x1, y1, x2, y2), _ in candidates:
                    try:
                        if mask[y1, x1] == 1 and mask[y2, x2] == 1:
                            # 可视化 + 保存
                            if save_path:
                                image = cv2.imread(image_path)
                                image = cv2.resize(image,(combined.shape[1],combined.shape[0]))
                                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                                # os.makedirs(os.path.dirname(save_path), exist_ok=True)
                                cv2.imwrite(save_path, image)
                                print(f"Image with rectangle saved at {save_path}")
                            resize_w, resize_h = combined.shape[1],combined.shape[0]
                            scale_x = orig_w / resize_w
                            scale_y = orig_h / resize_h
                            x1_orig = int(x1 * scale_x)
                            x2_orig = int(x2 * scale_x)
                            y1_orig = int(y1 * scale_y)
                            y2_orig = int(y2 * scale_y)
                            cx = (x1_orig + x2_orig) // 2
                            cy = (y1_orig + y2_orig) // 2
                            target_size = 90
                            half = target_size // 2
                            x1_exp = max(0, cx - half)
                            y1_exp = max(0, cy - half)
                            x2_exp = min(orig_w - 1, cx + half)
                            y2_exp = min(orig_h - 1, cy + half)
                            print(f"扩展后的原图坐标: ({x1_exp}, {y1_exp}), ({x2_exp}, {y2_exp})")
                            image_full = cv2.imread(image_path)  # 原图大小读取
                            cv2.rectangle(image_full, (x1_exp, y1_exp), (x2_exp, y2_exp), (0, 0, 255), 2)
                            cv2.imwrite("expanded_bbox_on_original.jpg", image_full)
                            print("📌 扩大后的候选框已绘制到原图并保存为 expanded_bbox_on_original.jpg")
                            return (x1_exp, y1_exp), (x2_exp, y2_exp)
                    except Exception as e:
                        print("the error is:",e)
                        pass  # 若越界等问题，继续下一个
                    # for _ in range(100):
                    #     random_x1 = random.randint(x, x + w - 50)
                    #     random_y1 = random.randint(y, y + h - 50)
                    #     random_x2 = random_x1 + size
                    #     random_y2 = random_y1 + size
                    #     # print(mask[random_y1, random_x1] == 1,mask[random_y2, random_x2] == 1)
                    #     try:
                    #         if save_path and mask[random_y1, random_x1] == 1 and mask[random_y2, random_x2] == 1:
                    #             cv2.rectangle(image,(random_x2, random_y2), (random_x1, random_y1), (0, 255, 0), 2)
                    #             cv2.imwrite(save_path, image)
                    #             print(f"Image with rectangle saved at {save_path}")  
                    #             return (random_x1,random_y1),(random_x2,random_y2)
                    #     except:
                    #         pass

    ax.imshow(mask_image)
    plt.axis('off')

def generate_gt_mask_from_intersection(random_rectangle, yolo_boxes, image, mask_img,sam2_model, threshold_iou):
    """
    判断随机生成的矩形与YOLO的框是否足够接近，
    若满足条件则调用SAM获取精准掩码作为GT。
    """
    image_np = np.array(image)  
    x1_rect, y1_rect = random_rectangle[0]
    x2_rect, y2_rect = random_rectangle[1]
    rect_mask = np.zeros(image_np.shape[:2], dtype=np.uint8)
    cv2.rectangle(rect_mask, (x1_rect, y1_rect), (x2_rect, y2_rect), color=255, thickness=-1)

    rect_box = [min(x1_rect, x2_rect), min(y1_rect, y2_rect), max(x1_rect, x2_rect), max(y1_rect, y2_rect)]

    for box in yolo_boxes:
        iou = calculate_iou(rect_box, box)
        print(f"与YOLO box的IoU为: {iou}, 阈值: {threshold_iou}")
        
        if iou >= threshold_iou:
            # 在YOLO框内随机取两个点
            x_min, y_min, x_max, y_max = box
            input_point1 = (np.random.randint(x_min, x_max), np.random.randint(y_min, y_max))
            input_point2 = (np.random.randint(x_min, x_max), np.random.randint(y_min, y_max))
            input_point3 = (np.random.randint(x_min, x_max), np.random.randint(y_min, y_max))

            # 使用SAM生成精准掩码
            gt_mask = get_gt_mask_from_sam(image, sam2_model, [input_point1, input_point2,input_point3], rect_mask)
            mask_img[gt_mask > 0] = 0
            # 保存gt掩码
            cv2.imwrite('gt_mask_from_sam.png', gt_mask)
            print(f"SAM生成的GT掩码已保存至 gt_mask_from_sam.png")

            return gt_mask,mask_img
    h, w = image_np.shape[:2]
    black_mask = np.zeros((h, w), dtype=np.uint8)
    no_match_save_path = 'gt_mask_from_sam.png'
    cv2.imwrite(no_match_save_path, black_mask)
    print("未找到满足阈值条件的YOLO box。")
    print(f"未匹配成功，保存空掩码图至 {no_match_save_path}")
    return None,mask_img

def calculate_iou(boxA, boxB):
    """计算两个box的IoU."""
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])

    inter_area = max(0, xB - xA + 1) * max(0, yB - yA + 1)

    boxA_area = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxB_area = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)

    iou = inter_area / float(boxA_area + boxB_area - inter_area)
    return iou

def get_gt_mask_from_sam(image, sam2_model, input_points, rect_mask):
    """使用SAM根据两个点生成掩码，并保存选取点和掩码图"""
    predictor = SAM2ImagePredictor(sam2_model)
    print("load sam2")
    predictor.set_image(image)

    input_point_np = np.array(input_points)
    input_label = np.array([1, 1,1])

    masks, _, _ = predictor.predict(
        point_coords=input_point_np,
        point_labels=input_label,
        multimask_output=False,
    )

    mask_img = masks[0].astype(np.uint8) * 255
    # mask_img[rect_mask == 255] = 0  # 将 `random_rectangle` 区域设为黑色

    # 保存SAM生成的掩码图
    mask_save_path = 'sam_gt_mask.jpg'
    cv2.imwrite(mask_save_path, mask_img)
    print(f"SAM生成的掩码已保存至 {mask_save_path}")

    # 把选取的两个点画在原图上
    image_with_points = np.array(image).copy()
    for point in input_points:
        cv2.circle(image_with_points, point, radius=5, color=(255, 0, 0), thickness=-1)

    # 保存带有标记点的原图
    point_marked_save_path = 'image_with_points.jpg'
    image_bgr = cv2.cvtColor(image_with_points, cv2.COLOR_RGB2BGR)
    cv2.imwrite(point_marked_save_path, image_bgr)
    print(f"带点标记的原图已保存至 {point_marked_save_path}")

    return mask_img

def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   

def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))    

def display_mask(mask, ax, random_color=False, borders = True):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask = mask.astype(np.uint8)
    mask_image =  mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    if borders:
        import cv2
        contours, _ = cv2.findContours(mask,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 
        # Try to smooth contours
        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2) 
        cv2.imwrite("check.jpg", mask_image)
    ax.imshow(mask_image)
    
def random_points_below(point, radius, min_distance, model, image, max_attempts=100):
    """
    在给定的point偏下方50像素的区域内，随机选择两个点直到满足条件。
    
    参数：
    - point: (x, y) 格式的坐标
    - radius: 随机点的最大半径
    - min_distance: 两个随机点之间的最小距离
    - max_attempts: 最大尝试次数，避免死循环
    
    返回：
    - 两个随机点的坐标，如果没有找到合适的点则返回None
    """
    for _ in range(max_attempts):
        # 在点的偏下方50像素区域内随机选择两个点
        x1 = random.randint(point[0] - radius, point[0] + radius)
        y1 = random.randint(point[1] + 50, point[1] + 50 + radius)  # 偏下50像素
        
        x2 = random.randint(point[0] - radius, point[0] + radius)
        y2 = random.randint(point[1] + 50, point[1] + 50 + radius)  # 偏下50像素
        
        # 计算两个点之间的欧几里得距离
        distance = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
        
        # 检查距离条件
        if distance >= min_distance and is_point_in_car_area((x1, y1), model, image) and is_point_in_car_area((x2, y2), model, image) :
            return [(x1, y1), (x2, y2)]
    
    # 如果超过最大尝试次数还没有找到合适的点，返回None
    return None
    

def show_masks(image, masks, scores, image_path, strategy,point_coords=None, box_coords=None, input_labels=None, borders=True, save_path=None):
    for i, (mask, score) in enumerate(zip(masks, scores)):
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        display_mask(mask, plt.gca(), borders=borders)
        if point_coords is not None:
            assert input_labels is not None
            show_points(point_coords, input_labels, plt.gca())
        if box_coords is not None:
            # boxes
            show_box(box_coords, plt.gca())
        plt.axis('off')
        plt.savefig('check.jpg', bbox_inches='tight', pad_inches=0)  # 保存图像
        point1,point2 = attention_mask(mask, plt.gca(), image_path,strategy,borders=borders, image=image, save_path=save_path)
        return point1,point2 

def random_crop(image, target_width, target_height, mask_point1, mask_point2):
    # global global_mask_point1_relative, global_mask_point2_relative
    """从两个对角点的中点裁剪指定宽度和高度的区域，避免超出图像边界"""
    width, height = image.size
    # 计算两个对角点的中点
    center_x = (mask_point1[0] + mask_point2[0]) // 2
    center_y = (mask_point1[1] + mask_point2[1]) // 2
    
    # 计算裁剪区域的左上角和右下角
    left = center_x - target_width // 2
    top = center_y - target_height // 2
    right = left + target_width
    bottom = top + target_height
    
    # 确保裁剪区域不会超出图像边界
    if left < 0:
        left = 0
        right = target_width
    if top < 0:
        top = 0
        bottom = target_height
    if right > width:
        right = width
        left = width - target_width
    if bottom > height:
        bottom = height
        top = height - target_height
    
    # 计算 padding
    top_padding = max(0, top)
    left_padding = max(0, left)
    
    # 裁剪图像
    cropped_image = image.crop((left, top, right, bottom))
    
    global_mask_point1_relative = (mask_point1[0] - left, mask_point1[1] - top)
    global_mask_point2_relative = (mask_point2[0] - left, mask_point2[1] - top)
    print("裁剪后点的相对位置为:")
    print("mask_point1:", global_mask_point1_relative)
    print("mask_point2:", global_mask_point2_relative)
    return cropped_image, top_padding, left_padding,global_mask_point1_relative,global_mask_point2_relative

def get_left_right_points(lane_data,image_path):
    lanes = lane_data["lanes"]
    h_samples = lane_data["h_samples"]
    model = load_yolov5_model()
    # 找到h_samples的中间索引
    mid_idx = len(h_samples) // 2
    image = cv2.imread(image_path)
    # 存储最左和最右的点
    left_point = None
    right_point = None
    points = []
    # 遍历每条车道线
    for lane in lanes:
        # 去掉值为-2的无效点
        valid_points = [(x, y) for x, y in zip(lane, h_samples) if x != -2]

        if valid_points:
            if lane[mid_idx] != -2:
                for i in range(mid_idx-2,0,-1):
                    left_point = lane[i] 
                    print(left_point)
                    if lane[i] != -2:
                        point = (left_point,h_samples[i])
                        FLAG = is_point_in_car_area(point, model, image)
                        print(point,FLAG)
                        if FLAG:
                            points.append((left_point,h_samples[i]))
                            break
            else:
                point = (1540/2, 590/2+30)  # 初始点坐标
                radius = 50         # 随机点的最大半径
                min_distance = 40   # 两个点之间的最小距离
                points = random_points_below(point, radius, min_distance,model,image)
                # first_non_minus_two = next((x for x in lane if x != -2), None)
                # if first_non_minus_two:
                #     idx = lane.index(first_non_minus_two)
                #     for i in range(idx+5,idx,-1):
                #         left_point = lane[i] 
                #         if lane[i] != -2:
                #             point = (left_point,h_samples[i])
                #             FLAG = is_point_in_car_area(point, model, image)
                #             if FLAG:
                #                 points.append((left_point,h_samples[i]))
                #                 break
                        
    # return left_point, right_point
    return points

def sam2segment(image_path,points,strategy):
    # print(points)
    image = Image.open(image_path)
    image = np.array(image.convert("RGB"))
    predictor.set_image(image)
    # print([points[0][0], points[0][1]])
    input_point = np.array([(points[0][0], points[0][1])])
    input_label = np.array([1])
    masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
    )
    sorted_ind = np.argsort(scores)[::-1]
    masks = masks[sorted_ind]
    scores = scores[sorted_ind]
    logits = logits[sorted_ind]
    #mask 
    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask
    points_set = []
    for point in points:
        points_set.append((point[0], point[1]))
    # print(points_set)
    input_point = np.array(points_set)
    input_label = np.array([1]*len(points_set))
    masks, scores, _ = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    mask_input=mask_input[None, :, :],
    multimask_output=False,
    )
    # random_mask_selection(image, masks, mask_index=0,output_path="cropped_image.jpg")
    point1,point2 = show_masks(image, masks, scores, image_path, strategy,save_path="masked_image.jpg")
    return point1,point2

def draw_point(image_path,points):
    image = cv2.imread(image_path)
    if image is not None:
        # 绘制点
        for point in points:
            cv2.circle(image, point, radius=5, color=(0, 255, 0), thickness=-1)  # 绿色点

        # 保存图像
        output_path = "output_image_with_points.jpg"
        cv2.imwrite(output_path, image)
        print(f"Image saved with points at {output_path}")
    else:
        print("Error: Image could not be loaded.")
        
def generate_mask(original_img_path, point1, point2):
    """根据坐标生成掩码图像"""
    # 读取原图
    original_img = cv2.imread(original_img_path)
    
    # 获取原图的尺寸
    height, width, _ = original_img.shape

    # 创建一个黑色的 mask 图像，尺寸与原图相同
    mask = np.zeros((height, width), dtype=np.uint8)

    # 计算3/4点
    three_quarter_point = (
        int(point1[0] + 0.95 * (point2[0] - point1[0])),  # 计算 x 坐标
        int(point1[1] + 0.95 * (point2[1] - point1[1]))   # 计算 y 坐标
    )

    # 画出一个白色的矩形（将该区域填充为白色）
    cv2.rectangle(mask, point1, three_quarter_point, color=255, thickness=-1)

    # 保存生成的mask图像
    mask_path = original_img_path.replace('test.jpg', 'mask_test.jpg')
    cv2.imwrite(mask_path, mask)
    print(mask_path)
    return mask_path, point1, three_quarter_point

    def extract_lanes_in_crop(lane_data, crop_x_min, crop_x_max, crop_y_min, crop_y_max):
        """
        过滤 TuSimple `lanes`，只保留 `crop` 内的部分
        """
        cropped_lanes = []
        for lane in lane_data["lanes"]:
            cropped_lane = []
            for x, y in zip(lane, lane_data["h_samples"]):
                if x != -2 and crop_x_min <= x <= crop_x_max and crop_y_min <= y <= crop_y_max:
                    cropped_lane.append((x, y))
                    # new_x = x - crop_x_min
                    # new_y = y - crop_y_min
                    # cropped_lane.append((new_x, new_y))
            if cropped_lane:
                cropped_lanes.append(cropped_lane)

        return cropped_lanes


def generate_trigger_crop(image_path: str, lane_data: dict):
    """
    输入一张图像路径，返回处理后的 crop 图像和 crop mask 图像路径。
    """
    # 1. 获取触发点
    points = get_left_right_points(lane_data, image_path)
    print(f"[INFO] 获取 trigger 点: {points}")
    draw_point(image_path, points)

    # 2. 使用 SAM2 获取 mask 点
    image = load_image(image_path)
    mask_point1, mask_point2 = sam2segment(image_path, points, "LDA")

    # 3. Crop 原图
    input_image, *_ = random_crop(image, 512, 512, mask_point1, mask_point2)
    input_crop_path = "crop.jpg"
    input_image.save(input_crop_path)

    # 4. 生成 trigger mask
    mask_path, point1, point2 = generate_mask(image_path, mask_point1, mask_point2)
    mask_img = load_image(mask_path)
    mask_img, *_ = random_crop(mask_img, 512, 512, mask_point1, mask_point2)
    crop_mask_path = "crop_mask.jpg"
    cv2.imwrite(crop_mask_path, np.array(mask_img))

    return input_crop_path, crop_mask_path

if __name__ == "__main__":
    lane_data = {"lanes": [[-2, -2, -2, -2, -2, -2, -2, 814, 751, 688, 625, 562, 500, 438, 373, 305, 234, 160, 88, 16, -64, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2], [-2, -2, -2, -2, -2, -2, -2, 818, 801, 784, 768, 751, 734, 717, 701, 685, 668, 651, 634, 618, 601, 585, 568, 551, 535, 518, 502, 484, 468, 451, 435, 418, 401, 385, 368, 351, 335, 318, 301, 287], [-2, -2, -2, -2, -2, -2, -2, 863, 872, 881, 890, 899, 908, 918, 927, 936, 945, 954, 964, 972, 982, 991, 1000, 1009, 1018, 1027, 1036, 1046, 1055, 1064, 1073, 1082, 1091, 1100, 1109, 1119, 1128, 1137, 1146, 1154]], "h_samples": [200, 210, 220, 230, 240, 250, 260, 270, 280, 290, 300, 310, 320, 330, 340, 350, 360, 370, 380, 390, 400, 410, 420, 430, 440, 450, 460, 470, 480, 490, 500, 510, 520, 530, 540, 550, 560, 570, 580, 590], "raw_file": "driver_182_30frame/06010513_0036.MP4/00270.jpg"}

    image_path = "driver_182_30frame/06010513_0036.MP4/00270.jpg" 
    points = get_left_right_points(lane_data,image_path)
    print(points)
    draw_point(image_path,points)
    # left_point, right_point = get_left_right_points(lane_data)
    # print(f"Left point: {left_point}, Right point: {right_point}")
    # sam2segment(image_path,left_point, right_point)
    image = load_image(image_path)
    mask_point1,mask_point2 = sam2segment(image_path,points,"LDA")
    input_image,top_padding,left_padding,global_mask_point1_relative,global_mask_point2_relative = random_crop(image, 512, 512, mask_point1, mask_point2)
    input_image.save("crop.jpg")  # 直接用 PIL 的 `save()` 方法
    print(f"Image saved with points at crop.jpg")
    mask_path, point1, point2 = generate_mask('culane_test.jpg', mask_point1, mask_point2)
    mask_img = load_image(mask_path)
    mask_img,top_padding,left_padding,global_mask_point1_relative,global_mask_point2_relative = random_crop(mask_img, 512, 512,mask_point1,mask_point2)
    
    mask_img = np.array(mask_img)
    # print(mask_img.shape)
    model = load_yolov5_model()
    yolo_results = model(input_image)
    yolo_boxes = []
    car_class_id = [2, 5, 7]  # 汽车、巴士、卡车等类别ID，根据实际情况调整

    for result in yolo_results:
        boxes = result.boxes.xyxy.cpu().numpy()
        class_ids = result.boxes.cls.cpu().numpy().astype(int)

        for box, cls in zip(boxes, class_ids):
            if cls in car_class_id:
                x_min, y_min, x_max, y_max = box[:4]
                yolo_boxes.append([int(x_min), int(y_min), int(x_max), int(y_max)])
    _,mask_img=generate_gt_mask_from_intersection([global_mask_point1_relative,global_mask_point2_relative], yolo_boxes, input_image, mask_img,sam2_model, threshold_iou=0.01)
    cv2.imwrite("crop_mask.jpg", mask_img)

    print("Mask 已成功保存至 crop_mask.jpg")
    crop_x_min = min(mask_point1[0], mask_point2[0])
    crop_x_max = max(mask_point1[0], mask_point2[0])
    crop_y_min = min(mask_point1[1], mask_point2[1])
    crop_y_max = max(mask_point1[1], mask_point2[1])


    def extract_lanes_in_crop(lane_data, crop_x_min, crop_x_max, crop_y_min, crop_y_max):
        """
        过滤 TuSimple `lanes`，只保留 `crop` 内的部分
        """
        cropped_lanes = []
        for lane in lane_data["lanes"]:
            cropped_lane = []
            for x, y in zip(lane, lane_data["h_samples"]):
                if x != -2 and crop_x_min <= x <= crop_x_max and crop_y_min <= y <= crop_y_max:
                    cropped_lane.append((x, y))
                    # new_x = x - crop_x_min
                    # new_y = y - crop_y_min
                    # cropped_lane.append((new_x, new_y))
            if cropped_lane:
                cropped_lanes.append(cropped_lane)

        return cropped_lanes

    # **获取在 crop 范围内的 lane**
    cropped_lanes = extract_lanes_in_crop(lane_data, crop_x_min, crop_x_max, crop_y_min, crop_y_max)
    # print(cropped_lanes)
    # def draw_lane_mask(image, lanes):
    #     """
    #     画出 `lane_mask` 只在 `crop` 图像中
    #     """
    #     height, width, _ = image.shape
    #     lane_mask = np.zeros((height, width), dtype=np.uint8)

    #     for lane in lanes:
    #         points = np.array(lane, dtype=np.int32)
    #         cv2.polylines(lane_mask, [points], isClosed=False, color=255, thickness=5)

    #     return lane_mask

    # crop_image = load_image("crop.jpg").convert("RGB")
    # crop_image = np.array(crop_image)
    # lane_mask = draw_lane_mask(crop_image, cropped_lanes)
    def draw_lane_mask_on_original(image, cropped_lanes):
        """
        在原图上绘制 **仅包含 cropped_lanes** 的车道线
        """
        height, width, _ = image.shape
        lane_mask = np.zeros((height, width), dtype=np.uint8)

        for lane in cropped_lanes:
            points = np.array(lane, dtype=np.int32)
            cv2.polylines(lane_mask, [points], isClosed=False, color=255, thickness=10)

        return lane_mask
    
    def random_crop_lane(image, target_width, target_height, mask_point1, mask_point2):
        """从两个对角点的中点裁剪指定宽度和高度的区域，避免超出图像边界"""
        
        # **确保 image 是 NumPy 数组**
        if isinstance(image, Image.Image):
            image = np.array(image)

        height, width = image.shape[:2]  # 获取 NumPy 数组的大小

        # 计算两个对角点的中点
        center_x = (mask_point1[0] + mask_point2[0]) // 2
        center_y = (mask_point1[1] + mask_point2[1]) // 2

        # 计算裁剪区域的左上角和右下角
        left = max(0, center_x - target_width // 2)
        top = max(0, center_y - target_height // 2)
        right = min(width, left + target_width)
        bottom = min(height, top + target_height)

        # 计算 padding（如果裁剪区域超出边界）
        top_padding = max(0, target_height - (bottom - top))
        left_padding = max(0, target_width - (right - left))

        # **使用 NumPy 进行裁剪**
        cropped_image = image[top:bottom, left:right]

        return cropped_image, top_padding, left_padding
    # **绘制 lane_mask 在原图上**
    raw_image = np.array(load_image(image_path).convert("RGB"))
    lane_mask = draw_lane_mask_on_original(raw_image, cropped_lanes)
    lane_mask_pil = Image.fromarray(lane_mask) 
    crop_image,top_padding,left_padding,global_mask_point1_relative,global_mask_point2_relative = random_crop(lane_mask_pil, 512, 512,mask_point1,mask_point2)

    # **保存 lane_mask**
    crop_image.save("lane_mask_crop.jpg")
    print("✅ 车道 Mask 已保存为 lane_mask_crop.jpg")
    
    crop_img = cv2.imread("crop.jpg")  # 读取原图（BGR格式）
    mask_img = cv2.imread("crop_mask.jpg", cv2.IMREAD_GRAYSCALE)  # 读取掩码（灰度图）
    if crop_img.shape[:2] != mask_img.shape:
        print("⚠️ Resizing mask to match crop image size...")
        mask_img = cv2.resize(mask_img, (crop_img.shape[1], crop_img.shape[0]))
    white_overlay = np.ones_like(crop_img) * 255  # 生成全白图
    masked_result = np.where(mask_img[:, :, None] == 255, white_overlay, crop_img)  # 只替换白色部分

    # **保存叠加后的图像**
    cv2.imwrite("crop_with_mask.jpg", masked_result)
    print("✅ 叠加后的 Mask 图像已保存至 crop_with_mask.jpg")