refactor: decouple TempFeatureExtractor from detection logic
- Remove detect_and_extract() method from TempFeatureExtractor (moved to FaultDiagnoser pipeline) - Remove YOLO model loading from TempFeatureExtractor.__init__ - Remove duplicate train_yolo_model/validate_model from detect.py (already exist in scripts/train.py) - Fix usage_example.py to use new import path and TempFeatureExtractor API
This commit is contained in:
@@ -8,9 +8,7 @@ from pathlib import Path
|
|||||||
# 添加项目根目录到路径
|
# 添加项目根目录到路径
|
||||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||||
|
|
||||||
from src.detect import InfraredDetector
|
from src.infrared_data import InfraredDetector, TempFeatureExtractor, FaultDiagnoser
|
||||||
from src.extract_temp import TempFeatureExtractor
|
|
||||||
from src.diagnose import FaultDiagnoser
|
|
||||||
import cv2
|
import cv2
|
||||||
import json
|
import json
|
||||||
|
|
||||||
@@ -48,7 +46,7 @@ def example3_temperature_extraction():
|
|||||||
"""示例3:温度特征提取"""
|
"""示例3:温度特征提取"""
|
||||||
print("\n【示例3】温度特征提取")
|
print("\n【示例3】温度特征提取")
|
||||||
|
|
||||||
extractor = TempFeatureExtractor("weights/best.pt")
|
extractor = TempFeatureExtractor()
|
||||||
detector = InfraredDetector("weights/best.pt")
|
detector = InfraredDetector("weights/best.pt")
|
||||||
|
|
||||||
# 读取图像
|
# 读取图像
|
||||||
|
|||||||
@@ -3,7 +3,6 @@ YOLO目标检测脚本
|
|||||||
用于红外图像中的目标检测
|
用于红外图像中的目标检测
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Optional, List, Dict, Union
|
from typing import Optional, List, Dict, Union
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -117,66 +116,6 @@ class InfraredDetector:
|
|||||||
return img, len(detections)
|
return img, len(detections)
|
||||||
|
|
||||||
|
|
||||||
def train_yolo_model(data_yaml: str,
|
|
||||||
model_size: str = 'n',
|
|
||||||
epochs: int = 100,
|
|
||||||
imgsz: int = 640,
|
|
||||||
batch: int = 16,
|
|
||||||
output_dir: str = 'models/train') -> str:
|
|
||||||
"""
|
|
||||||
训练YOLO模型
|
|
||||||
|
|
||||||
Args:
|
|
||||||
data_yaml: 数据集配置文件路径
|
|
||||||
model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
|
|
||||||
epochs: 训练轮数
|
|
||||||
imgsz: 输入图像大小
|
|
||||||
batch: 批量大小
|
|
||||||
output_dir: 输出目录
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
最佳模型路径
|
|
||||||
"""
|
|
||||||
model = YOLO(f'yolov8{model_size}.pt')
|
|
||||||
|
|
||||||
results = model.train(
|
|
||||||
data=data_yaml,
|
|
||||||
epochs=epochs,
|
|
||||||
imgsz=imgsz,
|
|
||||||
batch=batch,
|
|
||||||
project=output_dir,
|
|
||||||
name='infrared_fault',
|
|
||||||
exist_ok=True,
|
|
||||||
patience=20,
|
|
||||||
save=True,
|
|
||||||
plots=True
|
|
||||||
)
|
|
||||||
|
|
||||||
best_model_path = f"{output_dir}/infrared_fault/weights/best.pt"
|
|
||||||
return best_model_path
|
|
||||||
|
|
||||||
|
|
||||||
def validate_model(model_path: str, data_yaml: str) -> Dict:
|
|
||||||
"""
|
|
||||||
验证模型性能
|
|
||||||
|
|
||||||
Args:
|
|
||||||
model_path: 模型路径
|
|
||||||
data_yaml: 数据集配置路径
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
验证指标
|
|
||||||
"""
|
|
||||||
model = YOLO(model_path)
|
|
||||||
metrics = model.val(data=data_yaml)
|
|
||||||
return {
|
|
||||||
'mAP50': float(metrics.box.map50),
|
|
||||||
'mAP50-95': float(metrics.box.map),
|
|
||||||
'precision': float(metrics.box.mp),
|
|
||||||
'recall': float(metrics.box.mr)
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
import argparse
|
import argparse
|
||||||
|
|
||||||
@@ -191,23 +130,9 @@ if __name__ == '__main__':
|
|||||||
detect_parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
|
detect_parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
|
||||||
detect_parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
|
detect_parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
|
||||||
|
|
||||||
# Train command
|
|
||||||
train_parser = subparsers.add_parser('train', help='Train YOLO model')
|
|
||||||
train_parser.add_argument('--data', type=str, required=True, help='Data yaml path')
|
|
||||||
train_parser.add_argument('--model-size', type=str, default='n', choices=['n', 's', 'm', 'l', 'x'])
|
|
||||||
train_parser.add_argument('--epochs', type=int, default=100)
|
|
||||||
train_parser.add_argument('--imgsz', type=int, default=640)
|
|
||||||
train_parser.add_argument('--batch', type=int, default=16)
|
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if args.command == 'detect':
|
if args.command == 'detect':
|
||||||
detector = InfraredDetector(args.model, args.conf)
|
detector = InfraredDetector(args.model, args.conf)
|
||||||
img, n_detections = detector.detect_and_draw(args.image, args.output, device=args.device)
|
img, n_detections = detector.detect_and_draw(args.image, args.output, device=args.device)
|
||||||
print(f"Detected {n_detections} objects")
|
print(f"Detected {n_detections} objects")
|
||||||
|
|
||||||
elif args.command == 'train':
|
|
||||||
best_path = train_yolo_model(
|
|
||||||
args.data, args.model_size, args.epochs, args.imgsz, args.batch
|
|
||||||
)
|
|
||||||
print(f"Training complete. Best model: {best_path}")
|
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ class FaultDiagnoser:
|
|||||||
temp_threshold: 温度阈值配置
|
temp_threshold: 温度阈值配置
|
||||||
"""
|
"""
|
||||||
self.detector = InfraredDetector(model_path) if model_path else None
|
self.detector = InfraredDetector(model_path) if model_path else None
|
||||||
self.extractor = TempFeatureExtractor(model_path)
|
self.extractor = TempFeatureExtractor()
|
||||||
self.temp_threshold = temp_threshold or {}
|
self.temp_threshold = temp_threshold or {}
|
||||||
|
|
||||||
def diagnose(self, image: Union[str, np.ndarray],
|
def diagnose(self, image: Union[str, np.ndarray],
|
||||||
|
|||||||
@@ -3,83 +3,26 @@ ROI温度特征提取模块
|
|||||||
基于YOLO检测结果,提取ROI区域的温度特征
|
基于YOLO检测结果,提取ROI区域的温度特征
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import cv2
|
|
||||||
import json
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Dict, List, Optional, Union
|
from typing import Dict, List, Union
|
||||||
from ultralytics import YOLO
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
from .utils.temp_utils import (
|
from .utils.temp_utils import (
|
||||||
extract_temp_stats,
|
|
||||||
extract_temp_histogram,
|
|
||||||
calculate_temp_gradient,
|
calculate_temp_gradient,
|
||||||
normalize_roi
|
extract_temp_histogram,
|
||||||
|
extract_temp_stats,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class TempFeatureExtractor:
|
class TempFeatureExtractor:
|
||||||
"""红外图像ROI温度特征提取器"""
|
"""红外图像ROI温度特征提取器"""
|
||||||
|
|
||||||
def __init__(self, model_path: Optional[str] = None):
|
def extract_from_roi(
|
||||||
"""
|
self, roi: np.ndarray, include_hist: bool = True, include_gradient: bool = True
|
||||||
初始化温度特征提取器
|
) -> Dict:
|
||||||
|
|
||||||
Args:
|
|
||||||
model_path: YOLO模型路径,如果为None则只提取温度特征
|
|
||||||
"""
|
|
||||||
self.model = None
|
|
||||||
if model_path and Path(model_path).exists():
|
|
||||||
self.model = YOLO(model_path)
|
|
||||||
|
|
||||||
def detect_and_extract(self, image: np.ndarray,
|
|
||||||
conf: float = 0.25,
|
|
||||||
device: Optional[str] = None) -> List[Dict]:
|
|
||||||
"""
|
|
||||||
对图像进行检测并提取ROI温度特征
|
|
||||||
|
|
||||||
Args:
|
|
||||||
image: 输入红外图像
|
|
||||||
conf: 检测置信度阈值
|
|
||||||
device: 运行设备,例如 '0' 表示GPU 0,'cpu' 表示CPU
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
检测结果列表,每个元素包含检测框和温度特征
|
|
||||||
"""
|
|
||||||
if self.model is None:
|
|
||||||
raise ValueError("Model not loaded. Please provide a valid model path.")
|
|
||||||
|
|
||||||
results = self.model(image, conf=conf, device=device)
|
|
||||||
detections = []
|
|
||||||
|
|
||||||
for result in results:
|
|
||||||
boxes = result.boxes
|
|
||||||
if boxes is None:
|
|
||||||
continue
|
|
||||||
|
|
||||||
for box in boxes:
|
|
||||||
xyxy = box.xyxy[0].cpu().numpy()
|
|
||||||
conf_score = float(box.conf[0])
|
|
||||||
cls_id = int(box.cls[0])
|
|
||||||
|
|
||||||
x1, y1, x2, y2 = map(int, xyxy)
|
|
||||||
roi = image[y1:y2, x1:x2]
|
|
||||||
|
|
||||||
temp_features = extract_temp_stats(roi)
|
|
||||||
temp_features.update({
|
|
||||||
'bbox': [int(x1), int(y1), int(x2), int(y2)],
|
|
||||||
'confidence': conf_score,
|
|
||||||
'class_id': cls_id,
|
|
||||||
'class_name': result.names[cls_id]
|
|
||||||
})
|
|
||||||
|
|
||||||
detections.append(temp_features)
|
|
||||||
|
|
||||||
return detections
|
|
||||||
|
|
||||||
def extract_from_roi(self, roi: np.ndarray,
|
|
||||||
include_hist: bool = True,
|
|
||||||
include_gradient: bool = True) -> Dict:
|
|
||||||
"""
|
"""
|
||||||
从指定ROI提取温度特征
|
从指定ROI提取温度特征
|
||||||
|
|
||||||
@@ -94,15 +37,14 @@ class TempFeatureExtractor:
|
|||||||
features = extract_temp_stats(roi)
|
features = extract_temp_stats(roi)
|
||||||
|
|
||||||
if include_hist:
|
if include_hist:
|
||||||
features['histogram'] = extract_temp_histogram(roi).tolist()
|
features["histogram"] = extract_temp_histogram(roi).tolist()
|
||||||
|
|
||||||
if include_gradient:
|
if include_gradient:
|
||||||
features.update(calculate_temp_gradient(roi))
|
features.update(calculate_temp_gradient(roi))
|
||||||
|
|
||||||
return features
|
return features
|
||||||
|
|
||||||
def extract_from_bbox(self, image: np.ndarray,
|
def extract_from_bbox(self, image: np.ndarray, bbox: List[int]) -> Dict:
|
||||||
bbox: List[int]) -> Dict:
|
|
||||||
"""
|
"""
|
||||||
从给定边界框提取温度特征
|
从给定边界框提取温度特征
|
||||||
|
|
||||||
@@ -134,56 +76,53 @@ def load_infrared_image(image_path: str) -> np.ndarray:
|
|||||||
return image
|
return image
|
||||||
|
|
||||||
|
|
||||||
def batch_extract_features(image_dir: Union[str, Path],
|
def batch_extract_from_detections(
|
||||||
output_path: Union[str, Path],
|
image_dir: Union[str, Path],
|
||||||
model_path: Optional[str] = None,
|
detections_list: List[Dict],
|
||||||
conf: float = 0.25,
|
output_path: Union[str, Path],
|
||||||
device: Optional[str] = None) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
批量处理图像并提取温度特征
|
批量处理已检测的结果,提取温度特征
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
image_dir: 图像目录路径
|
image_dir: 图像目录路径
|
||||||
|
detections_list: 检测结果列表
|
||||||
output_path: 输出JSON文件路径
|
output_path: 输出JSON文件路径
|
||||||
model_path: YOLO模型路径
|
|
||||||
conf: 检测置信度阈值
|
|
||||||
device: 运行设备,例如 '0' 表示GPU 0,'cpu' 表示CPU
|
|
||||||
"""
|
"""
|
||||||
image_dir = Path(image_dir)
|
image_dir = Path(image_dir)
|
||||||
extractor = TempFeatureExtractor(model_path)
|
extractor = TempFeatureExtractor()
|
||||||
|
|
||||||
all_results = []
|
all_results = []
|
||||||
for img_path in sorted(image_dir.glob('*')):
|
for det_result in detections_list:
|
||||||
if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
|
img_name = det_result.get("image", "")
|
||||||
|
img_path = image_dir / img_name
|
||||||
|
if not img_path.exists():
|
||||||
continue
|
continue
|
||||||
|
|
||||||
try:
|
try:
|
||||||
image = load_infrared_image(str(img_path))
|
image = load_infrared_image(str(img_path))
|
||||||
detections = extractor.detect_and_extract(image, conf=conf, device=device)
|
for det in det_result.get("detections", []):
|
||||||
|
bbox = det.get("bbox")
|
||||||
all_results.append({
|
if bbox is None:
|
||||||
'image': img_path.name,
|
continue
|
||||||
'detections': detections
|
x1, y1, x2, y2 = bbox
|
||||||
})
|
roi = image[y1:y2, x1:x2]
|
||||||
|
temp_features = extractor.extract_from_roi(roi)
|
||||||
|
temp_features.update(
|
||||||
|
{
|
||||||
|
"bbox": bbox,
|
||||||
|
"class_name": det.get("class_name"),
|
||||||
|
"confidence": det.get("confidence"),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
all_results.append(
|
||||||
|
{"image": img_name, "detections": det_result.get("detections", [])}
|
||||||
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error processing {img_path.name}: {e}")
|
print(f"Error processing {img_name}: {e}")
|
||||||
continue
|
continue
|
||||||
|
|
||||||
with open(output_path, 'w', encoding='utf-8') as f:
|
with open(output_path, "w", encoding="utf-8") as f:
|
||||||
json.dump(all_results, f, indent=2, ensure_ascii=False)
|
json.dump(all_results, f, indent=2, ensure_ascii=False)
|
||||||
|
|
||||||
print(f"Processed {len(all_results)} images, results saved to {output_path}")
|
print(f"Processed {len(all_results)} images, results saved to {output_path}")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
import argparse
|
|
||||||
|
|
||||||
parser = argparse.ArgumentParser(description='Extract temperature features from infrared images')
|
|
||||||
parser.add_argument('--image-dir', type=str, required=True, help='Input image directory')
|
|
||||||
parser.add_argument('--output', type=str, required=True, help='Output JSON path')
|
|
||||||
parser.add_argument('--model', type=str, default=None, help='YOLO model path')
|
|
||||||
parser.add_argument('--conf', type=float, default=0.25, help='Detection confidence threshold')
|
|
||||||
parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
|
|
||||||
|
|
||||||
args = parser.parse_args()
|
|
||||||
batch_extract_features(args.image_dir, args.output, args.model, args.conf, args.device)
|
|
||||||
|
|||||||
Reference in New Issue
Block a user