32 lines
1.4 KiB
Python
32 lines
1.4 KiB
Python
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}")
|