SkyLi
085cbc9950
1. 输出格式优化: - faults中移除冗余的bbox和confidence字段 - summary中fault_count改名为detection_count 2. fault_status与fault_type逻辑统一: - normal状态时fault_type为'none'，不再使用'suspect' - abnormal状态时才进行故障类型分类(hotspot/local_overheating/overheat) 3. 修复usage_example.py图像通道问题: - YOLO检测使用彩色图，温度提取时再转灰度图
140 lines
4.0 KiB
Python
140 lines
4.0 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
|
|
|
|
SAMPLE_IMAGE = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg"
|
|
|
|
|
|
def example1_basic_detection():
|
|
"""示例1：基础目标检测"""
|
|
print("\n【示例1】基础目标检测")
|
|
|
|
detector = InfraredDetector("weights/best.pt", conf=0.25)
|
|
detections = detector.detect(SAMPLE_IMAGE)
|
|
|
|
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/best.pt")
|
|
detector.detect_and_draw(
|
|
SAMPLE_IMAGE,
|
|
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/best.pt")
|
|
detector = InfraredDetector("weights/best.pt")
|
|
|
|
# 读取图像
|
|
color_img = cv2.imread(SAMPLE_IMAGE, cv2.IMREAD_COLOR)
|
|
|
|
# 检测目标（彩色图）
|
|
detections = detector.detect(color_img)
|
|
|
|
# 转为灰度图用于温度提取
|
|
gray_img = cv2.cvtColor(color_img, cv2.COLOR_BGR2GRAY)
|
|
|
|
# 提取第一个检测框的温度特征
|
|
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/best.pt")
|
|
result = diagnoser.diagnose(SAMPLE_IMAGE)
|
|
|
|
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/best.pt")
|
|
diagnoser.batch_diagnose(
|
|
"data/valid/images",
|
|
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/best.pt", temp_threshold=threshold)
|
|
result = diagnoser.diagnose(SAMPLE_IMAGE)
|
|
|
|
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()
|