diff --git a/scripts/usage_example.py b/scripts/usage_example.py index 66b038c..9bc135a 100644 --- a/scripts/usage_example.py +++ b/scripts/usage_example.py @@ -14,13 +14,15 @@ 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("data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg") + detections = detector.detect(SAMPLE_IMAGE) print(f"检测到 {len(detections)} 个目标:") for i, det in enumerate(detections): @@ -36,7 +38,7 @@ def example2_detect_and_draw(): detector = InfraredDetector("weights/best.pt") detector.detect_and_draw( - "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg", + SAMPLE_IMAGE, output_path=str(output_dir / "detection_result.jpg") ) print(f"结果已保存至: {output_dir / 'detection_result.jpg'}") @@ -46,15 +48,14 @@ def example3_temperature_extraction(): """示例3:温度特征提取""" print("\n【示例3】温度特征提取") - sample_image = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg" extractor = TempFeatureExtractor("weights/best.pt") detector = InfraredDetector("weights/best.pt") # 读取图像 - gray_img = cv2.imread(sample_image, cv2.IMREAD_GRAYSCALE) + gray_img = cv2.imread(SAMPLE_IMAGE, cv2.IMREAD_GRAYSCALE) # 检测目标 - detections = detector.detect(sample_image) + detections = detector.detect(SAMPLE_IMAGE) # 提取第一个检测框的温度特征 if detections: @@ -72,9 +73,8 @@ def example4_fault_diagnosis(): """示例4:故障诊断""" print("\n【示例4】故障诊断") - sample_image = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg" diagnoser = FaultDiagnoser("weights/best.pt") - result = diagnoser.diagnose(sample_image) + result = diagnoser.diagnose(SAMPLE_IMAGE) print(f"诊断状态: {result['summary']['status']}") print(f"故障类型: {result['summary']['fault_types']}") @@ -102,7 +102,6 @@ def example6_custom_threshold(): """示例6:自定义温度阈值""" print("\n【示例6】自定义温度阈值") - sample_image = "data/valid/images/20230907095711003_jpg.rf.0e14e8b4344dce21114683dc3f3e7e2e.jpg" # 自定义温度阈值 threshold = { "max_temp": 200, # 最高温度阈值 @@ -110,7 +109,7 @@ def example6_custom_threshold(): } diagnoser = FaultDiagnoser("weights/best.pt", temp_threshold=threshold) - result = diagnoser.diagnose(sample_image) + result = diagnoser.diagnose(SAMPLE_IMAGE) print(f"使用自定义阈值: max_temp={threshold['max_temp']}, std_temp={threshold['std_temp']}") print(f"诊断结果: {result['summary']['status']}")