from ultralytics import YOLO
 
# 加载预训练的YOLOv11n模型
model = YOLO(r"runs/segment/train2/weights/best.pt")
 
# 对指定的图像文件夹进行推理，并设置各种参数
results = model.predict(
    source="datasets/seg_point_offer_20240930_num30/images/",   # 数据来源，可以是文件夹、图片路径、视频、URL，或设备ID（如摄像头）
    conf=0.45,                      # 置信度阈值
    iou=0.6,                        # IoU 阈值
    imgsz=640,                      # 图像大小
    half=False,                     # 使用半精度推理
    device=None,                    # 使用设备，None 表示自动选择，比如'cpu','0'
    max_det=300,                    # 最大检测数量
    vid_stride=1,                   # 视频帧跳跃设置
    stream_buffer=False,            # 视频流缓冲
    visualize=False,                # 可视化模型特征
    augment=False,                  # 启用推理时增强
    agnostic_nms=False,             # 启用类无关的NMS
    classes=None,                   # 指定要检测的类别
    retina_masks=False,             # 使用高分辨率分割掩码
    embed=None,                     # 提取特征向量层
    show=False,                     # 是否显示推理图像
    save=True,                      # 保存推理结果
    save_frames=False,              # 保存视频的帧作为图像
    save_txt=True,                  # 保存检测结果到文本文件
    save_conf=False,                # 保存置信度到文本文件
    save_crop=False,                # 保存裁剪的检测对象图像
    show_labels=True,               # 显示检测的标签
    show_conf=True,                 # 显示检测置信度
    show_boxes=True,                # 显示检测框
    line_width=None                 # 设置边界框的线条宽度，比如2，4
)

# 加载预训练的YOLO模型
model = YOLO(r'D:\zlc\drone\drone-train\runs\segment\train\weights\best.pt')  # 替换为你的模型路径
# 设置要预测的图像文件夹路径
image_folder = r"D:\zlc\drone\drone-train\data\land\test\images"  # 替换为你的图像文件夹路径
# 设置保存结果的文件夹路径
output_folder = r"D:\zlc\drone\drone-train\data\land\test\output"  # 替换为你要保存结果的文件夹路径
os.makedirs(output_folder, exist_ok=True)
names = ['background','land']
# 支持的图像文件扩展名
valid_image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']

# 创建一个颜色表，为每个类别分配固定颜色
def create_color_map(num_classes):
    color_map = {}
    for class_id in range(num_classes):
        # 生成随机但固定的RGB颜色值
        color = tuple(np.random.randint(0, 255, size=3).tolist())
        color_map[class_id] = color
    return color_map

# 假设 names 是从模型结果中获得的类别名称列表
num_classes = len(names)
color_map = create_color_map(num_classes)

def process_images_in_folder(folder_path):
    for filename in os.listdir(folder_path):
        # 检查文件是否为图像
        ext = os.path.splitext(filename)[1].lower()
        if ext not in valid_image_extensions:
            continue

        # 构建完整的文件路径
        image_path = os.path.join(folder_path, filename)

        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图像: {image_path}")
            continue

        # 进行预测
        results = model.predict(source=image_path, conf=0.25, iou=0.45, imgsz=640,show_labels=False,show_boxes=False,show_conf=False)

        # 将图像从 BGR 转换为 RGB 格式（OpenCV 默认是 BGR）
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # 创建一个与输入图像相同大小的全透明图像，用于绘制掩码
        mask_overlay = np.zeros_like(image_rgb)

        for result in results:  # 遍历所有检测结果
            boxes = result.boxes.cpu().numpy() if hasattr(result.boxes, 'cpu') else result.boxes
            masks = result.masks.data.cpu().numpy() if hasattr(result, 'masks') and hasattr(result.masks, 'data') else None
            names = result.names

            for i, (box, cls) in enumerate(zip(boxes.xyxy, boxes.cls)):
                if masks is not None and i < len(masks):  # 确保有对应的掩码
                    mask = masks[i]

                    # 确保掩码是一个 NumPy 数组
                    if not isinstance(mask, np.ndarray):
                        print("掩码不是 NumPy 数组")
                        continue

                    # 调整掩码大小以匹配图像
                    mask_resized = cv2.resize(mask, (image.shape[1], image.shape[0]))

                    # 使用预定义的颜色映射获取当前类别的颜色
                    class_id = int(cls)
                    color = color_map.get(class_id, (0, 0, 0))  # 默认黑色

                    # 在 mask_overlay 上绘制掩码
                    mask_overlay[mask_resized > 0.5] = color

                    # 绘制边界框和标签
                    # x1, y1, x2, y2 = map(int, box)
                    # cv2.rectangle(image_rgb, (x1, y1), (x2, y2), color, 2)
                    # label = f"{names[int(cls)]}"
                    # cv2.putText(image_rgb, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

        # 将掩码与原始图像合并
        alpha = 0.5  # 掩码透明度
        output_image = cv2.addWeighted(image_rgb, 1, mask_overlay, alpha, 0)

        # 构建输出文件路径
        output_path = os.path.join(output_folder, filename)
        # 保存带有标注的图像到输出文件夹
        cv2.imwrite(output_path, output_image)
        print(f"已处理并保存: {output_path}")