SkyLi
fd1a5f4050
- 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
135 lines
4.1 KiB
Python
135 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
红外故障检测模块使用示例
|
|
"""
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# 添加项目根目录到路径
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from src.detect import InfraredDetector
|
|
from src.extract_temp import TempFeatureExtractor
|
|
from src.diagnose import FaultDiagnoser
|
|
import cv2
|
|
import json
|
|
|
|
|
|
def example1_basic_detection():
|
|
"""示例1：基础目标检测"""
|
|
print("\n【示例1】基础目标检测")
|
|
|
|
detector = InfraredDetector("weights/infrared_yolov8n.pt", conf=0.25)
|
|
detections = detector.detect("data/infrared/images/val/FLIR0166.jpg")
|
|
|
|
print(f"检测到 {len(detections)} 个目标:")
|
|
for i, det in enumerate(detections):
|
|
print(f" [{i+1}] {det['class_name']}: conf={det['confidence']:.3f}")
|
|
|
|
|
|
def example2_detect_and_draw():
|
|
"""示例2：检测并绘制结果"""
|
|
print("\n【示例2】检测并绘制结果")
|
|
|
|
output_dir = Path(__file__).parent.parent / "output"
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
detector = InfraredDetector("weights/infrared_yolov8n.pt")
|
|
detector.detect_and_draw(
|
|
"data/infrared/images/val/FLIR0166.jpg",
|
|
output_path=str(output_dir / "detection_result.jpg")
|
|
)
|
|
print(f"结果已保存至: {output_dir / 'detection_result.jpg'}")
|
|
|
|
|
|
def example3_temperature_extraction():
|
|
"""示例3：温度特征提取"""
|
|
print("\n【示例3】温度特征提取")
|
|
|
|
extractor = TempFeatureExtractor("weights/infrared_yolov8n.pt")
|
|
detector = InfraredDetector("weights/infrared_yolov8n.pt")
|
|
|
|
# 读取图像
|
|
gray_img = cv2.imread("data/infrared/images/val/FLIR0166.jpg", cv2.IMREAD_GRAYSCALE)
|
|
|
|
# 检测目标
|
|
detections = detector.detect("data/infrared/images/val/FLIR0166.jpg")
|
|
|
|
# 提取第一个检测框的温度特征
|
|
if detections:
|
|
bbox = detections[0]['bbox']
|
|
x1, y1, x2, y2 = map(int, bbox)
|
|
roi = gray_img[y1:y2, x1:x2]
|
|
|
|
features = extractor.extract_from_roi(roi)
|
|
print(f" - 最高温度: {features['temp_max']:.2f}")
|
|
print(f" - 平均温度: {features['temp_mean']:.2f}")
|
|
print(f" - 温度标准差: {features['temp_std']:.2f}")
|
|
|
|
|
|
def example4_fault_diagnosis():
|
|
"""示例4：故障诊断"""
|
|
print("\n【示例4】故障诊断")
|
|
|
|
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt")
|
|
result = diagnoser.diagnose("data/infrared/images/val/FLIR0166.jpg")
|
|
|
|
print(f"诊断状态: {result['summary']['status']}")
|
|
print(f"故障类型: {result['summary']['fault_types']}")
|
|
|
|
for fault in result['faults']:
|
|
print(f" - {fault['fault_type']}: temp_max={fault['temp_max']:.2f}")
|
|
|
|
|
|
def example5_batch_diagnosis():
|
|
"""示例5：批量诊断"""
|
|
print("\n【示例5】批量诊断")
|
|
|
|
output_dir = Path(__file__).parent.parent / "output"
|
|
output_dir.mkdir(exist_ok=True)
|
|
|
|
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt")
|
|
diagnoser.batch_diagnose(
|
|
"data/infrared/images/val",
|
|
str(output_dir / "batch_results.json")
|
|
)
|
|
print(f"批量诊断完成，结果已保存至: {output_dir / 'batch_results.json'}")
|
|
|
|
|
|
def example6_custom_threshold():
|
|
"""示例6：自定义温度阈值"""
|
|
print("\n【示例6】自定义温度阈值")
|
|
|
|
# 自定义温度阈值
|
|
threshold = {
|
|
"max_temp": 200, # 最高温度阈值
|
|
"std_temp": 40 # 标准差阈值
|
|
}
|
|
|
|
diagnoser = FaultDiagnoser("weights/infrared_yolov8n.pt", temp_threshold=threshold)
|
|
result = diagnoser.diagnose("data/infrared/images/val/FLIR0166.jpg")
|
|
|
|
print(f"使用自定义阈值: max_temp={threshold['max_temp']}, std_temp={threshold['std_temp']}")
|
|
print(f"诊断结果: {result['summary']['status']}")
|
|
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("红外故障检测模块使用示例")
|
|
print("=" * 60)
|
|
|
|
example1_basic_detection()
|
|
example2_detect_and_draw()
|
|
example3_temperature_extraction()
|
|
example4_fault_diagnosis()
|
|
example5_batch_diagnosis()
|
|
example6_custom_threshold()
|
|
|
|
print("\n" + "=" * 60)
|
|
print("所有示例执行完成!")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|