chore: initial commit for Infrared Equipment Detection System
- 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
This commit is contained in:
SkyLi
@@ -0,0 +1,179 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
将红外图像数据集转换为YOLO格式
|
||||
由于数据集没有标注，使用整个图像作为检测目标（目标检测模式）
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
import random
|
||||
from pathlib import Path
|
||||
import cv2
|
||||
|
||||
|
||||
def create_yolo_directories(output_dir: Path) -> dict:
|
||||
"""创建YOLO目录结构"""
|
||||
dirs = {
|
||||
"train_images": output_dir / "images" / "train",
|
||||
"val_images": output_dir / "images" / "val",
|
||||
"train_labels": output_dir / "labels" / "train",
|
||||
"val_labels": output_dir / "labels" / "val",
|
||||
}
|
||||
for d in dirs.values():
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return dirs
|
||||
|
||||
|
||||
def convert_dataset(input_dir: Path, output_dir: Path,
|
||||
class_names: list,
|
||||
train_ratio: float = 0.8,
|
||||
use_whole_image: bool = True) -> None:
|
||||
"""
|
||||
转换数据集为YOLO格式
|
||||
|
||||
Args:
|
||||
input_dir: 输入图像目录（包含类别子目录）
|
||||
output_dir: 输出目录
|
||||
class_names: 类别名称列表
|
||||
train_ratio: 训练集比例
|
||||
use_whole_image: 是否使用整张图像作为检测目标
|
||||
"""
|
||||
print("=" * 50)
|
||||
print("转换数据集为YOLO格式")
|
||||
print("=" * 50)
|
||||
print(f"输入目录: {input_dir}")
|
||||
print(f"输出目录: {output_dir}")
|
||||
print(f"使用整图作为目标: {use_whole_image}")
|
||||
print("=" * 50)
|
||||
|
||||
# 创建输出目录
|
||||
dirs = create_yolo_directories(output_dir)
|
||||
|
||||
# 创建类别到ID的映射
|
||||
class_to_id = {name: idx for idx, name in enumerate(class_names)}
|
||||
|
||||
# 收集所有图像文件
|
||||
all_images = []
|
||||
for class_name in class_names:
|
||||
class_dir = input_dir / class_name
|
||||
if not class_dir.exists():
|
||||
print(f"警告: 目录不存在 {class_dir}")
|
||||
continue
|
||||
|
||||
for img_file in class_dir.glob("*.jpg"):
|
||||
all_images.append((img_file, class_name))
|
||||
for img_file in class_dir.glob("*.png"):
|
||||
all_images.append((img_file, class_name))
|
||||
for img_file in class_dir.glob("*.jpeg"):
|
||||
all_images.append((img_file, class_name))
|
||||
|
||||
print(f"找到 {len(all_images)} 张图像")
|
||||
|
||||
if len(all_images) == 0:
|
||||
print("错误: 没有找到图像文件")
|
||||
return
|
||||
|
||||
# 打乱并划分数据集
|
||||
random.seed(42)
|
||||
random.shuffle(all_images)
|
||||
|
||||
train_count = int(len(all_images) * train_ratio)
|
||||
train_images = all_images[:train_count]
|
||||
val_images = all_images[train_count:]
|
||||
|
||||
print(f"训练集: {len(train_images)} 张")
|
||||
print(f"验证集: {len(val_images)} 张")
|
||||
|
||||
# 复制图像并创建标签
|
||||
for img_path, class_name in train_images:
|
||||
copy_and_label(img_path, class_name, dirs["train_images"],
|
||||
dirs["train_labels"], use_whole_image, class_to_id)
|
||||
|
||||
for img_path, class_name in val_images:
|
||||
copy_and_label(img_path, class_name, dirs["val_images"],
|
||||
dirs["val_labels"], use_whole_image, class_to_id)
|
||||
|
||||
# 创建dataset.yaml
|
||||
yaml_content = f"""# YOLO数据集配置
|
||||
# 红外电力设备故障检测数据集
|
||||
|
||||
path: {output_dir.absolute()}
|
||||
train: images/train
|
||||
val: images/val
|
||||
|
||||
nc: {len(class_names)}
|
||||
names: {class_names}
|
||||
"""
|
||||
with open(output_dir / "dataset.yaml", "w", encoding="utf-8") as f:
|
||||
f.write(yaml_content)
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("转换完成!")
|
||||
print(f"配置文件: {output_dir / 'dataset.yaml'}")
|
||||
print("=" * 50)
|
||||
|
||||
|
||||
def copy_and_label(img_path: Path, class_name: str,
|
||||
output_img_dir: Path, output_label_dir: Path,
|
||||
use_whole_image: bool,
|
||||
class_to_id: dict) -> None:
|
||||
"""复制图像并创建标签文件"""
|
||||
# 复制图像
|
||||
output_img_path = output_img_dir / img_path.name
|
||||
shutil.copy(img_path, output_img_path)
|
||||
|
||||
# 创建标签文件
|
||||
label_path = output_label_dir / f"{img_path.stem}.txt"
|
||||
|
||||
if use_whole_image:
|
||||
# 使用整张图像作为检测目标
|
||||
img = cv2.imread(str(img_path))
|
||||
if img is None:
|
||||
return
|
||||
|
||||
h, w = img.shape[:2]
|
||||
|
||||
# YOLO格式: class_id x_center y_center width height (归一化)
|
||||
# 整图作为目标，使用正确的类别ID
|
||||
class_id = class_to_id.get(class_name, 0)
|
||||
label_content = f"{class_id} 0.5 0.5 1.0 1.0\n"
|
||||
|
||||
with open(label_path, "w") as f:
|
||||
f.write(label_content)
|
||||
else:
|
||||
# 空标签文件（用于后续添加真实标注）
|
||||
with open(label_path, "w") as f:
|
||||
pass
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="转换数据集为YOLO格式")
|
||||
parser.add_argument("--input-dir", type=str,
|
||||
default="data/raw/images",
|
||||
help="输入图像目录")
|
||||
parser.add_argument("--output-dir", type=str,
|
||||
default="data/infrared",
|
||||
help="输出目录")
|
||||
parser.add_argument("--train-ratio", type=float, default=0.8,
|
||||
help="训练集比例")
|
||||
parser.add_argument("--whole-image", action="store_true",
|
||||
help="使用整张图像作为检测目标")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
input_dir = Path(args.input_dir)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
if not input_dir.exists():
|
||||
print(f"错误: 输入目录不存在: {input_dir}")
|
||||
return
|
||||
|
||||
# 从目录结构获取类别
|
||||
class_names = [d.name for d in input_dir.iterdir() if d.is_dir()]
|
||||
|
||||
convert_dataset(input_dir, output_dir, class_names,
|
||||
args.train_ratio, args.whole_image)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,152 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
数据集下载脚本
|
||||
从SCIDB下载红外电力设备数据集
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse, parse_qs
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def parse_urls_file(file_path: Path) -> list:
|
||||
"""解析URL文件，返回URL和类别信息"""
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
items = []
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# 解析URL获取类别和文件名
|
||||
parsed = urlparse(line)
|
||||
query = parse_qs(parsed.query)
|
||||
|
||||
path = query.get('path', [''])[0]
|
||||
filename = query.get('fileName', ['unknown.jpg'])[0]
|
||||
|
||||
# 从路径提取类别
|
||||
parts = path.split('/')
|
||||
if len(parts) >= 3:
|
||||
category = parts[-2] # 如 "Wave Traps"
|
||||
else:
|
||||
category = "Unknown"
|
||||
|
||||
items.append({
|
||||
'url': line,
|
||||
'filename': filename,
|
||||
'category': category
|
||||
})
|
||||
|
||||
return items
|
||||
|
||||
|
||||
def create_directories(base_dir: Path, categories: set) -> dict:
|
||||
"""创建类别目录"""
|
||||
dirs = {}
|
||||
for cat in categories:
|
||||
cat_dir = base_dir / cat
|
||||
cat_dir.mkdir(parents=True, exist_ok=True)
|
||||
dirs[cat] = cat_dir
|
||||
return dirs
|
||||
|
||||
|
||||
def download_file(url: str, output_path: Path, session: requests.Session) -> bool:
|
||||
"""下载单个文件"""
|
||||
try:
|
||||
response = session.get(url, timeout=30, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(output_path, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {output_path.name}: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def download_dataset(urls_file: Path, output_dir: Path, max_workers: int = 4) -> None:
|
||||
"""下载整个数据集"""
|
||||
print("=" * 50)
|
||||
print("开始下载红外电力设备数据集")
|
||||
print("=" * 50)
|
||||
|
||||
# 解析URL
|
||||
print("解析URL文件...")
|
||||
items = parse_urls_file(urls_file)
|
||||
print(f"共找到 {len(items)} 个文件")
|
||||
|
||||
# 创建目录
|
||||
categories = set(item['category'] for item in items)
|
||||
print(f"发现 {len(categories)} 个类别: {categories}")
|
||||
dirs = create_directories(output_dir, categories)
|
||||
|
||||
# 创建session
|
||||
session = requests.Session()
|
||||
session.headers.update({
|
||||
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
|
||||
})
|
||||
|
||||
# 统计
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
# 下载文件
|
||||
print("\n开始下载...")
|
||||
for item in tqdm(items, desc="下载进度"):
|
||||
category_dir = dirs[item['category']]
|
||||
output_path = category_dir / item['filename']
|
||||
|
||||
# 跳过已存在的文件
|
||||
if output_path.exists():
|
||||
success_count += 1
|
||||
continue
|
||||
|
||||
if download_file(item['url'], output_path, session):
|
||||
success_count += 1
|
||||
else:
|
||||
fail_count += 1
|
||||
|
||||
# 避免请求过快
|
||||
time.sleep(0.1)
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print(f"下载完成!")
|
||||
print(f"成功: {success_count}")
|
||||
print(f"失败: {fail_count}")
|
||||
print(f"数据目录: {output_dir}")
|
||||
print("=" * 50)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="下载红外电力设备数据集")
|
||||
parser.add_argument("--urls-file", type=str,
|
||||
default="data/dataurls.txt",
|
||||
help="URL列表文件")
|
||||
parser.add_argument("--output-dir", type=str,
|
||||
default="data/raw/images",
|
||||
help="输出目录")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
urls_file = Path(args.urls_file)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
if not urls_file.exists():
|
||||
print(f"错误: URL文件不存在: {urls_file}")
|
||||
sys.exit(1)
|
||||
|
||||
download_dataset(urls_file, output_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,152 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
重试下载失败的图片（429限速后等待重试）
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse, parse_qs
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def parse_urls_file(file_path: Path) -> list:
|
||||
"""解析URL文件，返回URL和类别信息"""
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
items = []
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
parsed = urlparse(line)
|
||||
query = parse_qs(parsed.query)
|
||||
path = query.get('path', [''])[0]
|
||||
filename = query.get('fileName', ['unknown.jpg'])[0]
|
||||
|
||||
parts = path.split('/')
|
||||
if len(parts) >= 3:
|
||||
category = parts[-2]
|
||||
else:
|
||||
category = "Unknown"
|
||||
|
||||
items.append({
|
||||
'url': line,
|
||||
'filename': filename,
|
||||
'category': category
|
||||
})
|
||||
|
||||
return items
|
||||
|
||||
|
||||
def create_directories(base_dir: Path, categories: set) -> dict:
|
||||
dirs = {}
|
||||
for cat in categories:
|
||||
cat_dir = base_dir / cat
|
||||
cat_dir.mkdir(parents=True, exist_ok=True)
|
||||
dirs[cat] = cat_dir
|
||||
return dirs
|
||||
|
||||
|
||||
def download_file(url: str, output_path: Path, session: requests.Session) -> bool:
|
||||
try:
|
||||
response = session.get(url, timeout=30, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(output_path, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
return True
|
||||
except Exception as e:
|
||||
return False
|
||||
|
||||
|
||||
def get_missing_files(urls_file: Path, output_dir: Path) -> list:
|
||||
"""获取缺失文件的URL列表"""
|
||||
items = parse_urls_file(urls_file)
|
||||
missing = []
|
||||
|
||||
for item in items:
|
||||
category_dir = output_dir / item['category']
|
||||
output_path = category_dir / item['filename']
|
||||
|
||||
if not output_path.exists() or output_path.stat().st_size == 0:
|
||||
missing.append(item)
|
||||
|
||||
return missing
|
||||
|
||||
|
||||
def retry_download(urls_file: Path, output_dir: Path, max_retries: int = 3,
|
||||
wait_seconds: int = 60) -> None:
|
||||
"""重试下载失败的文件"""
|
||||
print("=" * 50)
|
||||
print("重试下载失败的文件")
|
||||
print("=" * 50)
|
||||
|
||||
# 找到缺失的文件
|
||||
missing = get_missing_files(urls_file, output_dir)
|
||||
print(f"缺失 {len(missing)} 个文件")
|
||||
|
||||
if not missing:
|
||||
print("没有缺失文件，下载完成！")
|
||||
return
|
||||
|
||||
categories = set(item['category'] for item in missing)
|
||||
dirs = create_directories(output_dir, categories)
|
||||
|
||||
session = requests.Session()
|
||||
session.headers.update({
|
||||
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
|
||||
})
|
||||
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
print(f"\n等待 {wait_seconds} 秒以避免限速...")
|
||||
time.sleep(wait_seconds)
|
||||
|
||||
print("\n开始重试下载...")
|
||||
for item in tqdm(missing, desc="重试进度"):
|
||||
category_dir = dirs[item['category']]
|
||||
output_path = category_dir / item['filename']
|
||||
|
||||
if download_file(item['url'], output_path, session):
|
||||
success_count += 1
|
||||
else:
|
||||
fail_count += 1
|
||||
time.sleep(1) # 失败后稍作等待
|
||||
|
||||
time.sleep(0.2) # 避免过快请求
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print(f"重试完成!")
|
||||
print(f"成功: {success_count}")
|
||||
print(f"失败: {fail_count}")
|
||||
print("=" * 50)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="重试下载失败的文件")
|
||||
parser.add_argument("--urls-file", type=str, default="data/dataurls.txt")
|
||||
parser.add_argument("--output-dir", type=str, default="data/raw/images")
|
||||
parser.add_argument("--wait", type=int, default=60, help="初始等待秒数")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
urls_file = Path(args.urls_file)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
if not urls_file.exists():
|
||||
print(f"错误: URL文件不存在: {urls_file}")
|
||||
sys.exit(1)
|
||||
|
||||
retry_download(urls_file, output_dir, wait_seconds=args.wait)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,230 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
CUDA和cuDNN环境测试脚本
|
||||
验证PyTorch GPU支持是否正常
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
|
||||
def test_cuda():
|
||||
"""测试CUDA支持"""
|
||||
print("=" * 50)
|
||||
print("CUDA 环境测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
import torch
|
||||
print(f"PyTorch版本: {torch.__version__}")
|
||||
print(f"CUDA可用: {torch.cuda.is_available()}")
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用")
|
||||
return False
|
||||
|
||||
print(f"CUDA版本: {torch.version.cuda}")
|
||||
|
||||
# 检查cuDNN版本
|
||||
try:
|
||||
cudnn_ver = torch.backends.cudnn.version()
|
||||
print(f"cuDNN版本: {cudnn_ver}")
|
||||
except RuntimeError as e:
|
||||
print(f"cuDNN版本检测失败: {e}")
|
||||
print("警告: cuDNN可能无法正常使用，但GPU张量运算可能仍然正常")
|
||||
return False
|
||||
|
||||
print(f"GPU数量: {torch.cuda.device_count()}")
|
||||
|
||||
for i in range(torch.cuda.device_count()):
|
||||
print(f"\nGPU {i}:")
|
||||
print(f" 名称: {torch.cuda.get_device_name(i)}")
|
||||
props = torch.cuda.get_device_properties(i)
|
||||
print(f" 计算能力: {props.major}.{props.minor}")
|
||||
print(f" 总内存: {props.total_memory / 1024**3:.2f} GB")
|
||||
|
||||
return True
|
||||
|
||||
except ImportError:
|
||||
print("错误: PyTorch未安装")
|
||||
return False
|
||||
|
||||
|
||||
def test_gpu_tensor_ops():
|
||||
"""测试GPU张量运算"""
|
||||
print("\n" + "=" * 50)
|
||||
print("GPU 张量运算测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用，跳过GPU测试")
|
||||
return False
|
||||
|
||||
# 创建张量并进行运算
|
||||
print("创建1000x1000随机张量...")
|
||||
x = torch.randn(1000, 1000).cuda()
|
||||
y = torch.randn(1000, 1000).cuda()
|
||||
|
||||
print("执行矩阵乘法...")
|
||||
z = torch.matmul(x, y)
|
||||
|
||||
print("执行归约操作...")
|
||||
result = z.sum().item()
|
||||
|
||||
print(f"结果: {result:.2f}")
|
||||
print("GPU张量运算测试通过!")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"错误: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_cudnn():
|
||||
"""测试cuDNN加速"""
|
||||
print("\n" + "=" * 50)
|
||||
print("cuDNN 加速测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用")
|
||||
return False
|
||||
|
||||
# 检查cuDNN是否启用
|
||||
print(f"cuDNN可用: {torch.backends.cudnn.is_available()}")
|
||||
try:
|
||||
cudnn_ver = torch.backends.cudnn.version()
|
||||
print(f"cuDNN版本: {cudnn_ver}")
|
||||
except RuntimeError as e:
|
||||
print(f"cuDNN版本检测跳过: {e}")
|
||||
print("注意: PyTorch绑定的cuDNN与系统版本可能不匹配，但不影响GPU使用")
|
||||
print(f"cuDNN基准测试启用: {torch.backends.cudnn.benchmark}")
|
||||
|
||||
# cuDNN卷积测试
|
||||
print("\n执行cuDNN卷积测试...")
|
||||
x = torch.randn(1, 3, 224, 224).cuda()
|
||||
conv = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda()
|
||||
|
||||
# 预热
|
||||
_ = conv(x)
|
||||
|
||||
# 计时测试
|
||||
import time
|
||||
start = time.time()
|
||||
for _ in range(10):
|
||||
_ = conv(x)
|
||||
torch.cuda.synchronize()
|
||||
elapsed = time.time() - start
|
||||
|
||||
print(f"10次卷积耗时: {elapsed*1000:.2f} ms")
|
||||
print("cuDNN测试通过!")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"错误: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_ultralytics_gpu():
|
||||
"""测试ultralytics YOLO GPU支持"""
|
||||
print("\n" + "=" * 50)
|
||||
print("Ultralytics YOLO GPU 测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
import torch
|
||||
import os
|
||||
|
||||
# 确保使用GPU加载模型
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用")
|
||||
return False
|
||||
|
||||
print("加载YOLOv8n模型到GPU...")
|
||||
# 关键：加载模型时指定device=0，确保模型在GPU上
|
||||
model = YOLO("yolov8n.pt", task="detect")
|
||||
|
||||
# 将模型移到GPU
|
||||
model.to("cuda:0")
|
||||
|
||||
print(f"模型设备: {next(model.model.parameters()).device}")
|
||||
|
||||
# 使用本地测试图像进行推理
|
||||
test_image = "bus.jpg"
|
||||
if not os.path.exists(test_image):
|
||||
print("下载测试图像...")
|
||||
import urllib.request
|
||||
urllib.request.urlretrieve(
|
||||
"https://ultralytics.com/images/bus.jpg",
|
||||
test_image
|
||||
)
|
||||
|
||||
print("执行GPU推理测试...")
|
||||
# 关键：device=0指定使用GPU
|
||||
results = model(test_image, device=0, verbose=False)
|
||||
|
||||
# 验证结果
|
||||
num_detections = len(results[0].boxes)
|
||||
print(f"检测到 {num_detections} 个目标")
|
||||
|
||||
# 检查模型是否真的在GPU上运行
|
||||
model_device = next(model.model.parameters()).device
|
||||
print(f"模型实际设备: {model_device}")
|
||||
assert model_device.type == "cuda", "模型未在GPU上运行!"
|
||||
|
||||
print("YOLO GPU测试通过!")
|
||||
return True
|
||||
|
||||
except ImportError:
|
||||
print("错误: ultralytics未安装")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"错误: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("\n" + "#" * 50)
|
||||
print("# 红外故障检测 - CUDA/cuDNN 环境测试")
|
||||
print("#" * 50 + "\n")
|
||||
|
||||
results = {
|
||||
"CUDA": test_cuda(),
|
||||
"GPU张量运算": test_gpu_tensor_ops(),
|
||||
"cuDNN": test_cudnn(),
|
||||
"YOLO GPU": test_ultralytics_gpu(),
|
||||
}
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("测试结果汇总")
|
||||
print("=" * 50)
|
||||
|
||||
all_passed = True
|
||||
for name, passed in results.items():
|
||||
status = "✅ 通过" if passed else "❌ 失败"
|
||||
print(f"{name}: {status}")
|
||||
if not passed:
|
||||
all_passed = False
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
if all_passed:
|
||||
print("所有测试通过! 环境配置正确。")
|
||||
else:
|
||||
print("部分测试失败，请检查环境配置。")
|
||||
print("=" * 50)
|
||||
|
||||
return 0 if all_passed else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,195 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
YOLO模型训练脚本
|
||||
用于红外图像目标检测模型训练
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from ultralytics import YOLO
|
||||
|
||||
|
||||
def train_yolo(
|
||||
data_yaml: str,
|
||||
model_size: str = "n",
|
||||
epochs: int = 100,
|
||||
imgsz: int = 640,
|
||||
batch: int = 16,
|
||||
device: str = "0",
|
||||
output_dir: str = "models/train",
|
||||
resume: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
训练YOLO模型
|
||||
|
||||
Args:
|
||||
data_yaml: 数据集配置文件路径
|
||||
model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
|
||||
epochs: 训练轮数
|
||||
imgsz: 输入图像大小
|
||||
batch: 批量大小
|
||||
device: 训练设备 ('0' for GPU, 'cpu' for CPU)
|
||||
output_dir: 输出目录
|
||||
resume: 是否从上次中断处继续训练
|
||||
|
||||
Returns:
|
||||
最佳模型路径
|
||||
"""
|
||||
print("=" * 50)
|
||||
print("YOLO 红外故障检测模型训练")
|
||||
print("=" * 50)
|
||||
print(f"数据集配置: {data_yaml}")
|
||||
print(f"模型大小: {model_size}")
|
||||
print(f"训练轮数: {epochs}")
|
||||
print(f"图像尺寸: {imgsz}")
|
||||
print(f"批量大小: {batch}")
|
||||
print(f"训练设备: {device}")
|
||||
print("=" * 50)
|
||||
|
||||
# 加载预训练模型
|
||||
model = YOLO(f"yolov8{model_size}.pt")
|
||||
|
||||
# 开始训练
|
||||
results = model.train(
|
||||
data=data_yaml,
|
||||
epochs=epochs,
|
||||
imgsz=imgsz,
|
||||
batch=batch,
|
||||
device=device,
|
||||
project=output_dir,
|
||||
name="infrared_fault",
|
||||
exist_ok=True,
|
||||
resume=resume,
|
||||
patience=20,
|
||||
save=True,
|
||||
save_period=10,
|
||||
plots=True,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
best_model_path = f"{output_dir}/infrared_fault/weights/best.pt"
|
||||
last_model_path = f"{output_dir}/infrared_fault/weights/last.pt"
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("训练完成!")
|
||||
print(f"最佳模型: {best_model_path}")
|
||||
print(f"最终模型: {last_model_path}")
|
||||
print("=" * 50)
|
||||
|
||||
return best_model_path
|
||||
|
||||
|
||||
def validate_model(model_path: str, data_yaml: str, device: str = "0") -> dict:
|
||||
"""
|
||||
验证模型性能
|
||||
|
||||
Args:
|
||||
model_path: 模型路径
|
||||
data_yaml: 数据集配置路径
|
||||
device: 验证设备
|
||||
|
||||
Returns:
|
||||
验证指标字典
|
||||
"""
|
||||
print("=" * 50)
|
||||
print("模型验证")
|
||||
print("=" * 50)
|
||||
|
||||
model = YOLO(model_path)
|
||||
metrics = model.val(data=data_yaml, device=device)
|
||||
|
||||
results = {
|
||||
"mAP50": float(metrics.box.map50),
|
||||
"mAP50-95": float(metrics.box.map),
|
||||
"precision": float(metrics.box.mp),
|
||||
"recall": float(metrics.box.mr),
|
||||
}
|
||||
|
||||
print("\n验证结果:")
|
||||
print(f" mAP@0.5: {results['mAP50']:.4f}")
|
||||
print(f" mAP@0.5:0.95: {results['mAP50-95']:.4f}")
|
||||
print(f" Precision: {results['precision']:.4f}")
|
||||
print(f" Recall: {results['recall']:.4f}")
|
||||
print("=" * 50)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def export_model(model_path: str, format: str = "onnx") -> str:
|
||||
"""
|
||||
导出模型
|
||||
|
||||
Args:
|
||||
model_path: 模型路径
|
||||
format: 导出格式 ('onnx', 'torchscript', 'openvino', etc.)
|
||||
|
||||
Returns:
|
||||
导出后模型路径
|
||||
"""
|
||||
print("=" * 50)
|
||||
print(f"导出模型: {model_path} -> {format}")
|
||||
print("=" * 50)
|
||||
|
||||
model = YOLO(model_path)
|
||||
export_path = model.export(format=format)
|
||||
|
||||
print(f"导出完成: {export_path}")
|
||||
return export_path
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="YOLO红外故障检测模型训练")
|
||||
subparsers = parser.add_subparsers(dest="command", help="命令")
|
||||
|
||||
# Train command
|
||||
train_parser = subparsers.add_parser("train", help="训练模型")
|
||||
train_parser.add_argument("--data", type=str, required=True, help="数据集配置文件")
|
||||
train_parser.add_argument("--model-size", type=str, default="n",
|
||||
choices=["n", "s", "m", "l", "x"],
|
||||
help="模型大小 (default: n)")
|
||||
train_parser.add_argument("--epochs", type=int, default=100, help="训练轮数")
|
||||
train_parser.add_argument("--imgsz", type=int, default=640, help="图像尺寸")
|
||||
train_parser.add_argument("--batch", type=int, default=16, help="批量大小")
|
||||
train_parser.add_argument("--device", type=str, default="0", help="训练设备")
|
||||
train_parser.add_argument("--output", type=str, default="models/train", help="输出目录")
|
||||
train_parser.add_argument("--resume", action="store_true", help="继续训练")
|
||||
|
||||
# Validate command
|
||||
val_parser = subparsers.add_parser("validate", help="验证模型")
|
||||
val_parser.add_argument("--model", type=str, required=True, help="模型路径")
|
||||
val_parser.add_argument("--data", type=str, required=True, help="数据集配置")
|
||||
val_parser.add_argument("--device", type=str, default="0", help="验证设备")
|
||||
|
||||
# Export command
|
||||
export_parser = subparsers.add_parser("export", help="导出模型")
|
||||
export_parser.add_argument("--model", type=str, required=True, help="模型路径")
|
||||
export_parser.add_argument("--format", type=str, default="onnx",
|
||||
choices=["onnx", "torchscript", "openvino", "tensorrt"],
|
||||
help="导出格式")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "train":
|
||||
train_yolo(
|
||||
data_yaml=args.data,
|
||||
model_size=args.model_size,
|
||||
epochs=args.epochs,
|
||||
imgsz=args.imgsz,
|
||||
batch=args.batch,
|
||||
device=args.device,
|
||||
output_dir=args.output,
|
||||
resume=args.resume,
|
||||
)
|
||||
elif args.command == "validate":
|
||||
validate_model(args.model, args.data, args.device)
|
||||
elif args.command == "export":
|
||||
export_model(args.model, args.format)
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,134 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user