123 lines
4.9 KiB
Python
123 lines
4.9 KiB
Python
import yaml
|
|
from ultralytics import YOLO
|
|
import os
|
|
import glob
|
|
|
|
|
|
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
def find_latest_run_dir(project_path='runs/detect/default_project'):
|
|
"""
|
|
在指定的项目路径下，根据文件夹的修改时间找到最新的 'train' 目录。
|
|
"""
|
|
if not os.path.exists(project_path):
|
|
return None
|
|
train_dirs = [d for d in glob.glob(os.path.join(project_path, 'train*')) if os.path.isdir(d)]
|
|
if not train_dirs:
|
|
return None
|
|
latest_dir = max(train_dirs, key=os.path.getmtime)
|
|
return latest_dir
|
|
|
|
def run_pipeline(config_path='config.yaml'):
|
|
"""
|
|
读取配置文件并执行YOLO训练和导出流程。
|
|
"""
|
|
# 1. --- 读取配置文件 ---
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = yaml.safe_load(f)
|
|
print("✅ 配置文件加载成功！")
|
|
print(f"项目名称: {config['project_name']}")
|
|
except Exception as e:
|
|
print(f"❌ 加载配置文件时发生错误: {e}")
|
|
return
|
|
|
|
# 提取核心配置
|
|
base_model = config['base_model']
|
|
data_yaml_path = config['data_yaml']
|
|
project_name = config['project_name']
|
|
|
|
model = YOLO(base_model)
|
|
print(f"✅ 模型 '{base_model}' 初始化成功。")
|
|
|
|
best_model_path = None
|
|
|
|
# 2. --- 执行训练 ---
|
|
if config.get('run_training', False):
|
|
print("\n🚀 开始YOLO训练...")
|
|
try:
|
|
results = model.train(
|
|
data=data_yaml_path,
|
|
epochs=config['epochs'],
|
|
imgsz=config['imgsz'],
|
|
batch=config['batch_size'],
|
|
workers=config['workers'],
|
|
project=project_name,
|
|
name='train'
|
|
)
|
|
|
|
final_results = results[0] if isinstance(results, (list, tuple)) else results
|
|
|
|
assert final_results is not None and hasattr(final_results, 'save_dir'), \
|
|
"训练未产生一个包含保存目录的有效结果对象。"
|
|
|
|
best_model_path = os.path.join(final_results.save_dir, 'weights/best.pt')
|
|
|
|
print(f"✅ 训练完成！最佳模型已保存在: {best_model_path}")
|
|
except Exception as e:
|
|
print(f"❌ 训练过程中发生错误: {e}")
|
|
return
|
|
else:
|
|
print("\n⏩ 根据配置，跳过训练步骤。")
|
|
if config.get('run_export', False):
|
|
print(" 正在查找最新的训练结果...")
|
|
project_path = os.path.join(ROOT_DIR, project_name)
|
|
latest_run_dir = find_latest_run_dir(project_path)
|
|
|
|
if latest_run_dir:
|
|
potential_model_path = os.path.join(latest_run_dir, 'weights', 'best.pt')
|
|
if os.path.exists(potential_model_path):
|
|
best_model_path = potential_model_path
|
|
print(f"✅ 成功找到最新的模型: {best_model_path}")
|
|
else:
|
|
print(f"❌ 在最新的训练目录 '{latest_run_dir}' 中未找到 'best.pt' 文件。")
|
|
else:
|
|
print(f"❌ 在项目 '{project_name}' 中未找到任何过往的训练结果。")
|
|
|
|
# 3. --- 执行导出 ---
|
|
print(f"root_dir:{ROOT_DIR}")
|
|
# print(f"run_export:{config.get('run_export', False)}")
|
|
print(f"best model path:{best_model_path}")
|
|
if config.get('run_export', False) and best_model_path:
|
|
|
|
# --- 修改开始：增加最终的安全检查 ---
|
|
# 确保 best_model_path 是一个字符串，而不是元组或列表
|
|
if isinstance(best_model_path, (list, tuple)):
|
|
print(f"⚠️ 检测到模型路径为序列类型，自动提取第一个元素。原始值: {best_model_path}")
|
|
best_model_path = best_model_path[0]
|
|
# --- 修改结束 ---
|
|
|
|
print(f"\n🚀 开始将模型 '{best_model_path}' 导出为 {config['export_format']} 格式...")
|
|
try:
|
|
model_to_export = YOLO(best_model_path)
|
|
|
|
model_to_export.export(
|
|
format=config['export_format'],
|
|
imgsz=config['imgsz'],
|
|
half=config.get('half_precision', False)
|
|
)
|
|
|
|
exported_file_name = os.path.basename(best_model_path).replace('.pt', f".{config['export_format']}")
|
|
exported_file_path = os.path.join(os.path.dirname(best_model_path), exported_file_name)
|
|
print(f"✅ 导出成功！文件已保存在: {exported_file_path}")
|
|
except Exception as e:
|
|
print(f"❌ 导出过程中发生错误: {e}")
|
|
elif config.get('run_export', False):
|
|
print("\n⏩ 跳过导出步骤，因为未找到有效的模型路径。")
|
|
else:
|
|
print("\n⏩ 根据配置，跳过导出步骤。")
|
|
|
|
print("\n🎉 自动化流程执行完毕！")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run_pipeline() |