diff --git a/scripts/convert_to_yolo.py b/scripts/convert_to_yolo.py deleted file mode 100644 index 53ef36a..0000000 --- a/scripts/convert_to_yolo.py +++ /dev/null @@ -1,179 +0,0 @@ -#!/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()