SkyLi
fd1a5f4050
- 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
180 lines
5.5 KiB
Python
180 lines
5.5 KiB
Python
#!/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()
|