chore: initial commit for Infrared Equipment Detection System
- Add YOLO-based infrared thermal image detection system - Add dataset configuration (dataset.yaml) with 12 equipment classes - Add training, validation and detection scripts - Add temperature feature extraction and fault diagnosis modules - Update README and CLAUDE.md with correct dataset paths - Optimize .gitignore to track dataset.yaml while ignoring large data files
This commit is contained in:
SkyLi
@@ -0,0 +1,195 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
YOLO模型训练脚本
|
||||
用于红外图像目标检测模型训练
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from ultralytics import YOLO
|
||||
|
||||
|
||||
def train_yolo(
|
||||
data_yaml: str,
|
||||
model_size: str = "n",
|
||||
epochs: int = 100,
|
||||
imgsz: int = 640,
|
||||
batch: int = 16,
|
||||
device: str = "0",
|
||||
output_dir: str = "models/train",
|
||||
resume: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
训练YOLO模型
|
||||
|
||||
Args:
|
||||
data_yaml: 数据集配置文件路径
|
||||
model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
|
||||
epochs: 训练轮数
|
||||
imgsz: 输入图像大小
|
||||
batch: 批量大小
|
||||
device: 训练设备 ('0' for GPU, 'cpu' for CPU)
|
||||
output_dir: 输出目录
|
||||
resume: 是否从上次中断处继续训练
|
||||
|
||||
Returns:
|
||||
最佳模型路径
|
||||
"""
|
||||
print("=" * 50)
|
||||
print("YOLO 红外故障检测模型训练")
|
||||
print("=" * 50)
|
||||
print(f"数据集配置: {data_yaml}")
|
||||
print(f"模型大小: {model_size}")
|
||||
print(f"训练轮数: {epochs}")
|
||||
print(f"图像尺寸: {imgsz}")
|
||||
print(f"批量大小: {batch}")
|
||||
print(f"训练设备: {device}")
|
||||
print("=" * 50)
|
||||
|
||||
# 加载预训练模型
|
||||
model = YOLO(f"yolov8{model_size}.pt")
|
||||
|
||||
# 开始训练
|
||||
results = model.train(
|
||||
data=data_yaml,
|
||||
epochs=epochs,
|
||||
imgsz=imgsz,
|
||||
batch=batch,
|
||||
device=device,
|
||||
project=output_dir,
|
||||
name="infrared_fault",
|
||||
exist_ok=True,
|
||||
resume=resume,
|
||||
patience=20,
|
||||
save=True,
|
||||
save_period=10,
|
||||
plots=True,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
best_model_path = f"{output_dir}/infrared_fault/weights/best.pt"
|
||||
last_model_path = f"{output_dir}/infrared_fault/weights/last.pt"
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("训练完成!")
|
||||
print(f"最佳模型: {best_model_path}")
|
||||
print(f"最终模型: {last_model_path}")
|
||||
print("=" * 50)
|
||||
|
||||
return best_model_path
|
||||
|
||||
|
||||
def validate_model(model_path: str, data_yaml: str, device: str = "0") -> dict:
|
||||
"""
|
||||
验证模型性能
|
||||
|
||||
Args:
|
||||
model_path: 模型路径
|
||||
data_yaml: 数据集配置路径
|
||||
device: 验证设备
|
||||
|
||||
Returns:
|
||||
验证指标字典
|
||||
"""
|
||||
print("=" * 50)
|
||||
print("模型验证")
|
||||
print("=" * 50)
|
||||
|
||||
model = YOLO(model_path)
|
||||
metrics = model.val(data=data_yaml, device=device)
|
||||
|
||||
results = {
|
||||
"mAP50": float(metrics.box.map50),
|
||||
"mAP50-95": float(metrics.box.map),
|
||||
"precision": float(metrics.box.mp),
|
||||
"recall": float(metrics.box.mr),
|
||||
}
|
||||
|
||||
print("\n验证结果:")
|
||||
print(f" mAP@0.5: {results['mAP50']:.4f}")
|
||||
print(f" mAP@0.5:0.95: {results['mAP50-95']:.4f}")
|
||||
print(f" Precision: {results['precision']:.4f}")
|
||||
print(f" Recall: {results['recall']:.4f}")
|
||||
print("=" * 50)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def export_model(model_path: str, format: str = "onnx") -> str:
|
||||
"""
|
||||
导出模型
|
||||
|
||||
Args:
|
||||
model_path: 模型路径
|
||||
format: 导出格式 ('onnx', 'torchscript', 'openvino', etc.)
|
||||
|
||||
Returns:
|
||||
导出后模型路径
|
||||
"""
|
||||
print("=" * 50)
|
||||
print(f"导出模型: {model_path} -> {format}")
|
||||
print("=" * 50)
|
||||
|
||||
model = YOLO(model_path)
|
||||
export_path = model.export(format=format)
|
||||
|
||||
print(f"导出完成: {export_path}")
|
||||
return export_path
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="YOLO红外故障检测模型训练")
|
||||
subparsers = parser.add_subparsers(dest="command", help="命令")
|
||||
|
||||
# Train command
|
||||
train_parser = subparsers.add_parser("train", help="训练模型")
|
||||
train_parser.add_argument("--data", type=str, required=True, help="数据集配置文件")
|
||||
train_parser.add_argument("--model-size", type=str, default="n",
|
||||
choices=["n", "s", "m", "l", "x"],
|
||||
help="模型大小 (default: n)")
|
||||
train_parser.add_argument("--epochs", type=int, default=100, help="训练轮数")
|
||||
train_parser.add_argument("--imgsz", type=int, default=640, help="图像尺寸")
|
||||
train_parser.add_argument("--batch", type=int, default=16, help="批量大小")
|
||||
train_parser.add_argument("--device", type=str, default="0", help="训练设备")
|
||||
train_parser.add_argument("--output", type=str, default="models/train", help="输出目录")
|
||||
train_parser.add_argument("--resume", action="store_true", help="继续训练")
|
||||
|
||||
# Validate command
|
||||
val_parser = subparsers.add_parser("validate", help="验证模型")
|
||||
val_parser.add_argument("--model", type=str, required=True, help="模型路径")
|
||||
val_parser.add_argument("--data", type=str, required=True, help="数据集配置")
|
||||
val_parser.add_argument("--device", type=str, default="0", help="验证设备")
|
||||
|
||||
# Export command
|
||||
export_parser = subparsers.add_parser("export", help="导出模型")
|
||||
export_parser.add_argument("--model", type=str, required=True, help="模型路径")
|
||||
export_parser.add_argument("--format", type=str, default="onnx",
|
||||
choices=["onnx", "torchscript", "openvino", "tensorrt"],
|
||||
help="导出格式")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "train":
|
||||
train_yolo(
|
||||
data_yaml=args.data,
|
||||
model_size=args.model_size,
|
||||
epochs=args.epochs,
|
||||
imgsz=args.imgsz,
|
||||
batch=args.batch,
|
||||
device=args.device,
|
||||
output_dir=args.output,
|
||||
resume=args.resume,
|
||||
)
|
||||
elif args.command == "validate":
|
||||
validate_model(args.model, args.data, args.device)
|
||||
elif args.command == "export":
|
||||
export_model(args.model, args.format)
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user