chore: remove legacy convert_to_yolo.py script
新项目数据集已是完整标注的YOLO格式，无需旧的数据转换工具
This commit is contained in:
SkyLi
@@ -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()
|
||||
Reference in New Issue
Block a user