fix: restructure src layout to proper package directory
Move src/__init__.py, detect.py, diagnose.py, extract_temp.py, utils/ into src/infrared_data/ to fix wheel package structure. Previously the whl placed files at root instead of infrared_data/ directory, causing ModuleNotFoundError on install. Also update pyproject.toml packages.find config to use standard src-layout discovery.
This commit is contained in:
SkyLi
@@ -0,0 +1,6 @@
|
||||
# Infrared Fault Detection Package
|
||||
from .diagnose import FaultDiagnoser
|
||||
from .detect import InfraredDetector
|
||||
from .extract_temp import TempFeatureExtractor
|
||||
|
||||
__all__ = ['FaultDiagnoser', 'InfraredDetector', 'TempFeatureExtractor']
|
||||
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
YOLO目标检测脚本
|
||||
用于红外图像中的目标检测
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Union
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
import cv2
|
||||
|
||||
|
||||
class InfraredDetector:
|
||||
"""红外图像目标检测器"""
|
||||
|
||||
def __init__(self, model_path: Optional[str] = None, conf: float = 0.25):
|
||||
"""
|
||||
初始化检测器
|
||||
|
||||
Args:
|
||||
model_path: YOLO模型路径
|
||||
conf: 置信度阈值
|
||||
"""
|
||||
self.model = None
|
||||
self.conf = conf
|
||||
|
||||
if model_path and Path(model_path).exists():
|
||||
self.load_model(model_path)
|
||||
|
||||
def load_model(self, model_path: str) -> None:
|
||||
"""加载YOLO模型"""
|
||||
self.model = YOLO(model_path)
|
||||
print(f"Model loaded from: {model_path}")
|
||||
|
||||
def detect(self, image: Union[str, np.ndarray],
|
||||
conf: Optional[float] = None,
|
||||
device: Optional[str] = None) -> List[Dict]:
|
||||
"""
|
||||
对图像进行目标检测
|
||||
|
||||
Args:
|
||||
image: 图像路径或图像数组
|
||||
conf: 置信度阈值（覆盖默认值）
|
||||
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
|
||||
|
||||
Returns:
|
||||
检测结果列表
|
||||
"""
|
||||
if self.model is None:
|
||||
raise ValueError("Model not loaded. Please load a model first.")
|
||||
|
||||
threshold = conf if conf is not None else self.conf
|
||||
results = self.model(image, conf=threshold, 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()
|
||||
xywh = box.xywh[0].cpu().numpy()
|
||||
conf_score = float(box.conf[0])
|
||||
cls_id = int(box.cls[0])
|
||||
|
||||
detections.append({
|
||||
'bbox': xyxy.tolist(), # [x1, y1, x2, y2]
|
||||
'bbox_xywh': xywh.tolist(), # [x_center, y_center, w, h]
|
||||
'confidence': conf_score,
|
||||
'class_id': cls_id,
|
||||
'class_name': result.names[cls_id]
|
||||
})
|
||||
|
||||
return detections
|
||||
|
||||
def detect_and_draw(self, image: Union[str, np.ndarray],
|
||||
output_path: Optional[str] = None,
|
||||
conf: Optional[float] = None,
|
||||
device: Optional[str] = None) -> np.ndarray:
|
||||
"""
|
||||
检测并绘制结果
|
||||
|
||||
Args:
|
||||
image: 图像路径或图像数组
|
||||
output_path: 输出图像路径
|
||||
conf: 置信度阈值
|
||||
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
|
||||
|
||||
Returns:
|
||||
绘制了检测结果的图像
|
||||
"""
|
||||
if isinstance(image, str):
|
||||
img = cv2.imread(image)
|
||||
else:
|
||||
img = image.copy()
|
||||
|
||||
detections = self.detect(img, conf, device=device)
|
||||
|
||||
for det in detections:
|
||||
x1, y1, x2, y2 = map(int, det['bbox'])
|
||||
conf_score = det['confidence']
|
||||
class_name = det['class_name']
|
||||
|
||||
color = (0, 255, 0)
|
||||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||||
|
||||
label = f"{class_name} {conf_score:.2f}"
|
||||
cv2.putText(img, label, (x1, y1 - 10),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
||||
|
||||
if output_path:
|
||||
cv2.imwrite(output_path, img)
|
||||
print(f"Result saved to: {output_path}")
|
||||
|
||||
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__':
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description='YOLO Infrared Detection')
|
||||
subparsers = parser.add_subparsers(dest='command', help='Commands')
|
||||
|
||||
# Detect command
|
||||
detect_parser = subparsers.add_parser('detect', help='Detect objects in image')
|
||||
detect_parser.add_argument('--model', type=str, required=True, help='Model path')
|
||||
detect_parser.add_argument('--image', type=str, required=True, help='Image path')
|
||||
detect_parser.add_argument('--output', type=str, help='Output image path')
|
||||
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)')
|
||||
|
||||
# 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()
|
||||
|
||||
if args.command == 'detect':
|
||||
detector = InfraredDetector(args.model, args.conf)
|
||||
img, n_detections = detector.detect_and_draw(args.image, args.output, device=args.device)
|
||||
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}")
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
故障诊断模块
|
||||
基于检测结果和温度特征进行故障诊断
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Dict, List, Optional, Union
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
|
||||
from .extract_temp import TempFeatureExtractor
|
||||
from .detect import InfraredDetector
|
||||
|
||||
|
||||
class FaultDiagnoser:
|
||||
"""故障诊断器"""
|
||||
|
||||
def __init__(self,
|
||||
model_path: Optional[str] = None,
|
||||
temp_threshold: Optional[Dict] = None):
|
||||
"""
|
||||
初始化诊断器
|
||||
|
||||
Args:
|
||||
model_path: YOLO模型路径
|
||||
temp_threshold: 温度阈值配置
|
||||
"""
|
||||
self.detector = InfraredDetector(model_path) if model_path else None
|
||||
self.extractor = TempFeatureExtractor(model_path)
|
||||
self.temp_threshold = temp_threshold or {}
|
||||
|
||||
def diagnose(self, image: Union[str, np.ndarray],
|
||||
device: Optional[str] = None) -> Dict:
|
||||
"""
|
||||
对图像进行故障诊断
|
||||
|
||||
Args:
|
||||
image: 图像路径或数组
|
||||
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
|
||||
|
||||
Returns:
|
||||
诊断结果
|
||||
"""
|
||||
result = {
|
||||
'image': str(image) if isinstance(image, str) else 'array',
|
||||
'detections': [],
|
||||
'faults': [],
|
||||
'summary': {}
|
||||
}
|
||||
|
||||
if self.detector and self.detector.model:
|
||||
detections = self.detector.detect(image, device=device)
|
||||
result['detections'] = detections
|
||||
|
||||
import cv2
|
||||
if isinstance(image, str):
|
||||
img = cv2.imread(image, cv2.IMREAD_GRAYSCALE)
|
||||
else:
|
||||
img = image
|
||||
|
||||
for det in detections:
|
||||
bbox = det['bbox']
|
||||
x1, y1, x2, y2 = map(int, bbox)
|
||||
|
||||
roi = img[y1:y2, x1:x2]
|
||||
temp_features = self.extractor.extract_from_roi(roi)
|
||||
|
||||
fault_status, fault_type = self._analyze_fault(temp_features)
|
||||
temp_features['fault_status'] = fault_status
|
||||
temp_features['fault_type'] = fault_type
|
||||
|
||||
result['faults'].append(temp_features)
|
||||
else:
|
||||
raise ValueError("Detector not initialized with model")
|
||||
|
||||
result['summary'] = self._generate_summary(result['faults'])
|
||||
return result
|
||||
|
||||
def _analyze_fault(self, temp_features: Dict) -> tuple:
|
||||
"""
|
||||
分析温度特征，判断故障状态和类型
|
||||
|
||||
Args:
|
||||
temp_features: 温度特征
|
||||
|
||||
Returns:
|
||||
(fault_status, fault_type) 元组
|
||||
"""
|
||||
max_temp = temp_features.get('temp_max', 0)
|
||||
std_temp = temp_features.get('temp_std', 0)
|
||||
temp_range = max_temp - temp_features.get('temp_mean', 0)
|
||||
|
||||
max_threshold = self.temp_threshold.get('max_temp', 200)
|
||||
std_threshold = self.temp_threshold.get('std_temp', 50)
|
||||
|
||||
# 判断状态：两项都不超标才为 normal
|
||||
if max_temp <= max_threshold and std_temp <= std_threshold:
|
||||
return ('normal', 'none')
|
||||
|
||||
# 分类故障类型
|
||||
if std_temp > 30:
|
||||
return ('abnormal', 'hotspot')
|
||||
elif temp_range > 50:
|
||||
return ('abnormal', 'local_overheating')
|
||||
elif max_temp > 180:
|
||||
return ('abnormal', 'overheat')
|
||||
else:
|
||||
return ('abnormal', 'suspect')
|
||||
|
||||
def _generate_summary(self, faults: List[Dict]) -> Dict:
|
||||
"""
|
||||
生成诊断摘要
|
||||
|
||||
Args:
|
||||
faults: 故障列表
|
||||
|
||||
Returns:
|
||||
摘要信息
|
||||
"""
|
||||
if not faults:
|
||||
return {'status': 'no_detection', 'detection_count': 0}
|
||||
|
||||
abnormal_count = sum(1 for f in faults if f.get('fault_status') == 'abnormal')
|
||||
fault_types = [f.get('fault_type', 'unknown') for f in faults]
|
||||
|
||||
return {
|
||||
'status': 'fault_detected' if abnormal_count > 0 else 'normal',
|
||||
'detection_count': len(faults),
|
||||
'abnormal_count': abnormal_count,
|
||||
'fault_types': list(set(fault_types))
|
||||
}
|
||||
|
||||
def batch_diagnose(self, image_dir: Union[str, Path],
|
||||
output_path: Union[str, Path],
|
||||
device: Optional[str] = None) -> None:
|
||||
"""
|
||||
批量诊断
|
||||
|
||||
Args:
|
||||
image_dir: 图像目录
|
||||
output_path: 输出JSON路径
|
||||
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
|
||||
"""
|
||||
image_dir = Path(image_dir)
|
||||
results = []
|
||||
|
||||
for img_path in sorted(image_dir.glob('*')):
|
||||
if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
|
||||
continue
|
||||
|
||||
try:
|
||||
result = self.diagnose(str(img_path), device=device)
|
||||
result['image'] = img_path.name
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
print(f"Error processing {img_path.name}: {e}")
|
||||
continue
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(results, f, indent=2, ensure_ascii=False)
|
||||
|
||||
print(f"Batch diagnosis complete. {len(results)} images processed.")
|
||||
print(f"Results saved to: {output_path}")
|
||||
|
||||
|
||||
def load_threshold_config(config_path: str) -> Dict:
|
||||
"""加载温度阈值配置"""
|
||||
with open(config_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description='Fault Diagnosis')
|
||||
parser.add_argument('--model', type=str, required=True, help='YOLO model path')
|
||||
parser.add_argument('--image-dir', type=str, required=True, help='Image directory')
|
||||
parser.add_argument('--output', type=str, required=True, help='Output JSON path')
|
||||
parser.add_argument('--threshold', type=str, help='Threshold config JSON')
|
||||
parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
threshold = None
|
||||
if args.threshold:
|
||||
threshold = load_threshold_config(args.threshold)
|
||||
|
||||
diagnoser = FaultDiagnoser(args.model, threshold)
|
||||
diagnoser.batch_diagnose(args.image_dir, args.output, device=args.device)
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
ROI温度特征提取模块
|
||||
基于YOLO检测结果，提取ROI区域的温度特征
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Union
|
||||
from ultralytics import YOLO
|
||||
|
||||
from .utils.temp_utils import (
|
||||
extract_temp_stats,
|
||||
extract_temp_histogram,
|
||||
calculate_temp_gradient,
|
||||
normalize_roi
|
||||
)
|
||||
|
||||
|
||||
class TempFeatureExtractor:
|
||||
"""红外图像ROI温度特征提取器"""
|
||||
|
||||
def __init__(self, model_path: Optional[str] = None):
|
||||
"""
|
||||
初始化温度特征提取器
|
||||
|
||||
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提取温度特征
|
||||
|
||||
Args:
|
||||
roi: ROI区域图像
|
||||
include_hist: 是否包含直方图特征
|
||||
include_gradient: 是否包含梯度特征
|
||||
|
||||
Returns:
|
||||
温度特征字典
|
||||
"""
|
||||
features = extract_temp_stats(roi)
|
||||
|
||||
if include_hist:
|
||||
features['histogram'] = extract_temp_histogram(roi).tolist()
|
||||
|
||||
if include_gradient:
|
||||
features.update(calculate_temp_gradient(roi))
|
||||
|
||||
return features
|
||||
|
||||
def extract_from_bbox(self, image: np.ndarray,
|
||||
bbox: List[int]) -> Dict:
|
||||
"""
|
||||
从给定边界框提取温度特征
|
||||
|
||||
Args:
|
||||
image: 输入图像
|
||||
bbox: [x1, y1, x2, y2] 边界框坐标
|
||||
|
||||
Returns:
|
||||
温度特征字典
|
||||
"""
|
||||
x1, y1, x2, y2 = bbox
|
||||
roi = image[y1:y2, x1:x2]
|
||||
return self.extract_from_roi(roi)
|
||||
|
||||
|
||||
def load_infrared_image(image_path: str) -> np.ndarray:
|
||||
"""
|
||||
加载红外图像
|
||||
|
||||
Args:
|
||||
image_path: 图像路径
|
||||
|
||||
Returns:
|
||||
图像数组
|
||||
"""
|
||||
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
||||
if image is None:
|
||||
raise ValueError(f"Failed to load image: {image_path}")
|
||||
return image
|
||||
|
||||
|
||||
def batch_extract_features(image_dir: Union[str, Path],
|
||||
output_path: Union[str, Path],
|
||||
model_path: Optional[str] = None,
|
||||
conf: float = 0.25,
|
||||
device: Optional[str] = None) -> None:
|
||||
"""
|
||||
批量处理图像并提取温度特征
|
||||
|
||||
Args:
|
||||
image_dir: 图像目录路径
|
||||
output_path: 输出JSON文件路径
|
||||
model_path: YOLO模型路径
|
||||
conf: 检测置信度阈值
|
||||
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
|
||||
"""
|
||||
image_dir = Path(image_dir)
|
||||
extractor = TempFeatureExtractor(model_path)
|
||||
|
||||
all_results = []
|
||||
for img_path in sorted(image_dir.glob('*')):
|
||||
if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
|
||||
continue
|
||||
|
||||
try:
|
||||
image = load_infrared_image(str(img_path))
|
||||
detections = extractor.detect_and_extract(image, conf=conf, device=device)
|
||||
|
||||
all_results.append({
|
||||
'image': img_path.name,
|
||||
'detections': detections
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"Error processing {img_path.name}: {e}")
|
||||
continue
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(all_results, f, indent=2, ensure_ascii=False)
|
||||
|
||||
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)
|
||||
@@ -0,0 +1 @@
|
||||
# Utils Package
|
||||
@@ -0,0 +1,127 @@
|
||||
"""
|
||||
温度特征提取工具函数
|
||||
从红外图像ROI区域提取温度统计特征
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
from typing import Dict, List, Tuple, Optional
|
||||
|
||||
|
||||
def extract_temp_stats(roi: np.ndarray) -> Dict[str, float]:
|
||||
"""
|
||||
从ROI区域提取温度统计特征
|
||||
|
||||
Args:
|
||||
roi: ROI区域图像数组
|
||||
|
||||
Returns:
|
||||
温度统计特征字典
|
||||
"""
|
||||
if roi.size == 0:
|
||||
return {
|
||||
'temp_max': 0.0,
|
||||
'temp_min': 0.0,
|
||||
'temp_mean': 0.0,
|
||||
'temp_std': 0.0,
|
||||
'temp_median': 0.0,
|
||||
}
|
||||
|
||||
return {
|
||||
'temp_max': float(np.max(roi)),
|
||||
'temp_min': float(np.min(roi)),
|
||||
'temp_mean': float(np.mean(roi)),
|
||||
'temp_std': float(np.std(roi)),
|
||||
'temp_median': float(np.median(roi)),
|
||||
}
|
||||
|
||||
|
||||
def extract_temp_histogram(roi: np.ndarray, bins: int = 32) -> np.ndarray:
|
||||
"""
|
||||
提取温度分布直方图特征
|
||||
|
||||
Args:
|
||||
roi: ROI区域图像数组
|
||||
bins: 直方图bin数量
|
||||
|
||||
Returns:
|
||||
归一化直方图向量
|
||||
"""
|
||||
if roi.size == 0:
|
||||
return np.zeros(bins)
|
||||
|
||||
hist, _ = np.histogram(roi.flatten(), bins=bins, range=(0, 255))
|
||||
hist = hist.astype(np.float32) / (hist.sum() + 1e-7)
|
||||
return hist
|
||||
|
||||
|
||||
def calculate_temp_gradient(roi: np.ndarray) -> Dict[str, float]:
|
||||
"""
|
||||
计算温度梯度特征（反映温度变化剧烈程度）
|
||||
|
||||
Args:
|
||||
roi: ROI区域图像数组
|
||||
|
||||
Returns:
|
||||
温度梯度统计
|
||||
"""
|
||||
if roi.size == 0:
|
||||
return {'gradient_mean': 0.0, 'gradient_std': 0.0}
|
||||
|
||||
grad_x = cv2.Sobel(roi, cv2.CV_32F, 1, 0, ksize=3)
|
||||
grad_y = cv2.Sobel(roi, cv2.CV_32F, 0, 1, ksize=3)
|
||||
gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)
|
||||
|
||||
return {
|
||||
'gradient_mean': float(np.mean(gradient_magnitude)),
|
||||
'gradient_std': float(np.std(gradient_magnitude)),
|
||||
}
|
||||
|
||||
|
||||
def extract_temp_features(roi: np.ndarray, include_hist: bool = True,
|
||||
include_gradient: bool = True) -> Dict:
|
||||
"""
|
||||
从ROI区域提取完整温度特征
|
||||
|
||||
Args:
|
||||
roi: ROI区域图像数组
|
||||
include_hist: 是否包含直方图特征
|
||||
include_gradient: 是否包含梯度特征
|
||||
|
||||
Returns:
|
||||
完整温度特征字典
|
||||
"""
|
||||
features = extract_temp_stats(roi)
|
||||
|
||||
if include_hist:
|
||||
features['temp_histogram'] = extract_temp_histogram(roi)
|
||||
|
||||
if include_gradient:
|
||||
features.update(calculate_temp_gradient(roi))
|
||||
|
||||
return features
|
||||
|
||||
|
||||
def normalize_roi(roi: np.ndarray, method: str = 'minmax') -> np.ndarray:
|
||||
"""
|
||||
对ROI进行归一化处理
|
||||
|
||||
Args:
|
||||
roi: ROI区域图像数组
|
||||
method: 归一化方法 ('minmax' 或 'zscore')
|
||||
|
||||
Returns:
|
||||
归一化后的图像
|
||||
"""
|
||||
if method == 'minmax':
|
||||
min_val, max_val = roi.min(), roi.max()
|
||||
if max_val - min_val < 1e-7:
|
||||
return np.zeros_like(roi)
|
||||
return (roi - min_val) / (max_val - min_val)
|
||||
elif method == 'zscore':
|
||||
mean_val, std_val = np.mean(roi), np.std(roi)
|
||||
if std_val < 1e-7:
|
||||
return np.zeros_like(roi)
|
||||
return (roi - mean_val) / std_val
|
||||
else:
|
||||
raise ValueError(f"Unknown normalization method: {method}")
|
||||
Reference in New Issue
Block a user