import os
from ultralytics import YOLO
from pathlib import Path

# === 设置路径 ===
weights_path = './runs/segment/train6/weights/best.pt'  # 模型权重
source_dir = './Data/images/train'                      # 输入图片目录
output_dir = './Data/result/train'                     # 推理结果输出目录
os.makedirs(output_dir, exist_ok=True)

# === 加载模型 ===
model = YOLO(weights_path)

# === 推理参数 ===
results = model.predict(
    source=source_dir,        # 可为单图像路径、目录、视频、摄像头索引等
    save=True,                # 是否保存图像
    save_txt=False,           # 是否保存标签（分割任务不常用）
    save_conf=True,          # 是否保存置信度
    save_crop=False,          # 是否保存目标裁剪图
    project=output_dir,       # 输出根路径
    name='',                  # 子文件夹名（为空即直接输出到 output_dir）
    exist_ok=True,            # 允许覆盖已有文件夹
    imgsz=640,                # 推理分辨率（默认640）
    conf=0.02,                # 置信度阈值 # 设置的越高，分类越少
    iou=0.45,                 # NMS IOU 阈值
    # device='cuda:0'           # 改为 'cpu' 如果没有GPU
    device='cpu'   # ✅ 修改为 CPU 模式
)

print(f"✅ 推理完成，结果保存在：{output_dir}")
