Initial commit
This commit is contained in:
commit
eae2961d5b
|
|
@ -0,0 +1,209 @@
|
|||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[codz]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py.cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# UV
|
||||
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
#uv.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
#poetry.toml
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
|
||||
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
|
||||
#pdm.lock
|
||||
#pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# pixi
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
|
||||
#pixi.lock
|
||||
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
|
||||
# in the .venv directory. It is recommended not to include this directory in version control.
|
||||
.pixi
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.envrc
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
# Abstra
|
||||
# Abstra is an AI-powered process automation framework.
|
||||
# Ignore directories containing user credentials, local state, and settings.
|
||||
# Learn more at https://abstra.io/docs
|
||||
.abstra/
|
||||
|
||||
# Visual Studio Code
|
||||
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
||||
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
||||
# you could uncomment the following to ignore the entire vscode folder
|
||||
# .vscode/
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# Marimo
|
||||
marimo/_static/
|
||||
marimo/_lsp/
|
||||
__marimo__/
|
||||
|
||||
# Streamlit
|
||||
.streamlit/secrets.toml
|
||||
|
||||
#项目文件
|
||||
Safety_Detection_Project
|
||||
|
||||
#模型
|
||||
*.pt
|
||||
|
|
@ -0,0 +1,123 @@
|
|||
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()
|
||||
|
|
@ -0,0 +1,32 @@
|
|||
# ===================================================================
|
||||
# 自动化训练与导出配置文件
|
||||
# ===================================================================
|
||||
|
||||
# --- 项目与模型设置 ---
|
||||
# 项目名称，所有训练结果将保存在 runs/detect/{project_name} 文件夹下
|
||||
project_name: 'Safety_Detection_Project'
|
||||
|
||||
# 基础模型：可以是官方的预训练模型 (如 yolov8n.pt), 也可以是你自己的 .pt 文件路径
|
||||
base_model: 'yolo11n.pt'
|
||||
|
||||
# 数据集配置文件：【非常重要】请务必修改为你自己的 data .yaml 文件的绝对或相对路径
|
||||
data_yaml: 'coco128.yaml'
|
||||
|
||||
# --- 流程控制 ---
|
||||
# 是否执行训练步骤
|
||||
run_training: false
|
||||
# 是否执行导出步骤
|
||||
run_export: true
|
||||
|
||||
# --- 训练参数 ---
|
||||
epochs: 100
|
||||
imgsz: 640
|
||||
batch_size: 16
|
||||
# 使用多少个CPU核心进行数据加载，0表示只用主进程
|
||||
workers: 8
|
||||
|
||||
# --- 导出参数 ---
|
||||
# 目标格式 (e.g., onnx, tensorrt, coreml)
|
||||
export_format: 'onnx'
|
||||
# 是否使用半精度(FP16)导出，可以提速并减小文件大小
|
||||
half_precision: true
|
||||
Reference in New Issue