将 tensorflow keras 训练数据集转换为 Yolo 训练数据集

以 https://www.kaggle.com/datasets/vipoooool/new-plant-diseases-dataset 为例

1. 图像分类数据集文件结构 (例如用于 `yolov11n-cls.pt` 训练)

import os
import csv
import random
from PIL import Image
from sklearn.model_selection import train_test_split
import shutil# ====================== 配置参数 ======================
# 从 Kaggle Hub 下载植物病害数据集
# https://www.kaggle.com/datasets/vipoooool/new-plant-diseases-dataset
import kagglehub
tf_download_path = kagglehub.dataset_download("vipoooool/new-plant-diseases-dataset")
print("Path to dataset files:", tf_download_path)
# 定义数据集路径
tf_dataset_path = f"{tf_download_path}/New Plant Diseases Dataset(Augmented)/New Plant Diseases Dataset(Augmented)"INPUT_DATA_DIR = tf_dataset_path  # 输入数据集路径（解压后的根目录）
OUTPUT_YOLO_DIR = "./runs/traindata/yolo/yolo_plant_diseases_classify"        # 输出YOLO数据集路径
if os.path.exists(OUTPUT_YOLO_DIR):shutil.rmtree(OUTPUT_YOLO_DIR)
os.makedirs(OUTPUT_YOLO_DIR, exist_ok=True)TRAIN_SIZE = 0.8                                 # 训练集比例
IMAGE_EXTENSIONS = [".JPG", ".jpg", ".jpeg", ".png"]     # 支持的图像扩展名# ====================== 类别映射（需根据实际数据集调整） ======================
# 从原数据集的类别名称生成映射（示例：假设病害类别为文件夹名）
def get_class_mapping(data_dir):class_names = []for folder in os.listdir(data_dir):folder_path = os.path.join(data_dir, folder)if os.path.isdir(folder_path) and not folder.startswith("."):class_names.append(folder)class_names.sort()  # 按字母序排序，确保类别编号固定return {cls: idx for idx, cls in enumerate(class_names)}# ====================== 划分数据集并保存 ======================
def save_dataset(annotations, class_map, output_dir, train_size=0.8):# 划分训练集和验证集random.shuffle(annotations)split_idx = int(len(annotations) * train_size)train_data = annotations[:split_idx]val_data = annotations[split_idx:]# 创建目录结构os.makedirs(os.path.join(output_dir, "train"), exist_ok=True)os.makedirs(os.path.join(output_dir, "val"), exist_ok=True)for cls in class_map.keys():os.makedirs(os.path.join(output_dir, "train", cls), exist_ok=True)os.makedirs(os.path.join(output_dir, "val", cls), exist_ok=True)# 保存训练集for data in train_data:img_path = data["image_path"]cls = data["class_name"]try:shutil.copy2(img_path, os.path.join(output_dir, "train", cls))print(f"图像 {img_path} 复制到训练集 {cls} 类成功")except Exception as e:print(f"图像 {img_path} 复制到训练集 {cls} 类失败，错误信息: {e}")# 保存验证集for data in val_data:img_path = data["image_path"]cls = data["class_name"]try:shutil.copy2(img_path, os.path.join(output_dir, "val", cls))print(f"图像 {img_path} 复制到验证集 {cls} 类成功")except Exception as e:print(f"图像 {img_path} 复制到验证集 {cls} 类失败，错误信息: {e}")# 生成类别名文件（classes.names）with open(os.path.join(output_dir, "classes.names"), "w") as f:for cls in class_map.keys():f.write(f"{cls}\n")# 生成数据集配置文件（dataset.yaml）yaml_path = os.path.join(output_dir, "dataset.yaml")with open(yaml_path, "w") as f:f.write(f"path: {output_dir}\n")  # 数据集根路径f.write(f"train: train\n")  # 训练集路径（相对于path）f.write(f"val: val\n")      # 验证集路径# f.write(f"test: images/test\n")   # 测试集路径（如果有）f.write(f"nc: {len(class_map)}\n")  # 类别数# 修改 names 字段输出格式class_names = list(class_map.keys())f.write(f"names: {class_names}\n")return train_data, val_data# ====================== 主函数 ======================
if __name__ == "__main__":# 1. 检查输入路径是否存在if not os.path.exists(INPUT_DATA_DIR):raise FileNotFoundError(f"请先下载数据集并解压到路径：{INPUT_DATA_DIR}")# 2. 获取类别映射（假设图像按类别存放在子文件夹中）class_map = get_class_mapping(os.path.join(INPUT_DATA_DIR, "train"))  # 假设训练集图像在train子文件夹中，每个子文件夹为一个类别# 3. 解析标注（仅按文件夹分类）annotations = []for cls, idx in class_map.items():cls_dir = os.path.join(INPUT_DATA_DIR, "train", cls)  # 假设类别文件夹路径为train/类别名for img_file in os.listdir(cls_dir):if any(img_file.lower().endswith(ext) for ext in IMAGE_EXTENSIONS):img_path = os.path.join(cls_dir, img_file)annotations.append({"image_path": img_path,"class_name": cls})# 4. 保存为YOLO格式train_data, val_data = save_dataset(annotations, class_map, OUTPUT_YOLO_DIR, train_size=TRAIN_SIZE)print(f"✅ 转换完成！YOLO数据集已保存至：{OUTPUT_YOLO_DIR}")print(f"类别数：{len(class_map)}，训练集样本数：{len(train_data)}，验证集样本数：{len(val_data)}")

train的时候，使用的文件夹

2. 目标检测数据集文件结构 (例如用于 `yolo11n.pt` 训练)

import os
import csv
import random
from PIL import Image
from sklearn.model_selection import train_test_split
import shutil# ====================== 配置参数 ======================
# 从 Kaggle Hub 下载植物病害数据集
# https://www.kaggle.com/datasets/vipoooool/new-plant-diseases-dataset
import kagglehub
tf_download_path = kagglehub.dataset_download("vipoooool/new-plant-diseases-dataset")
print("Path to dataset files:", tf_download_path)
# 定义数据集路径
tf_dataset_path = f"{tf_download_path}/New Plant Diseases Dataset(Augmented)/New Plant Diseases Dataset(Augmented)"INPUT_DATA_DIR = tf_dataset_path  # 输入数据集路径（解压后的根目录）
OUTPUT_YOLO_DIR = "./traindata/yolo/yolo_plant_diseases"        # 输出YOLO数据集路径
if os.path.exists(OUTPUT_YOLO_DIR):shutil.rmtree(OUTPUT_YOLO_DIR)
os.makedirs(OUTPUT_YOLO_DIR, exist_ok=True)TRAIN_SIZE = 0.8                                 # 训练集比例
IMAGE_EXTENSIONS = [".JPG", ".jpg", ".jpeg", ".png"]     # 支持的图像扩展名# ====================== 类别映射（需根据实际数据集调整） ======================
# 从原数据集的类别名称生成映射（示例：假设病害类别为文件夹名）
def get_class_mapping(data_dir):class_names = []for folder in os.listdir(data_dir):folder_path = os.path.join(data_dir, folder)if os.path.isdir(folder_path) and not folder.startswith("."):class_names.append(folder)class_names.sort()  # 按字母序排序，确保类别编号固定return {cls: idx for idx, cls in enumerate(class_names)}# ====================== 解析CSV标注（假设标注在CSV中） ======================
def parse_csv_annotations(csv_path, class_map, image_dir):annotations = []with open(csv_path, "r", encoding="utf-8") as f:reader = csv.DictReader(f)for row in reader:image_name = row["image_path"]class_name = row["disease_class"]  # 需与CSV中的类别列名一致x_min = float(row["x_min"])y_min = float(row["y_min"])x_max = float(row["x_max"])y_max = float(row["y_max"])# 检查图像是否存在image_path = os.path.join(image_dir, image_name)if not os.path.exists(image_path):continue# 获取图像尺寸with Image.open(image_path) as img:img_width, img_height = img.size# 转换为YOLO坐标center_x = (x_min + x_max) / 2 / img_widthcenter_y = (y_min + y_max) / 2 / img_heightwidth = (x_max - x_min) / img_widthheight = (y_max - y_min) / img_heightannotations.append({"image_path": image_path,"class_id": class_map[class_name],"bbox": (center_x, center_y, width, height)})return annotations# ====================== 划分数据集并保存 ======================
def save_dataset(annotations, class_map, output_dir, train_size=0.8):# 划分训练集和验证集random.shuffle(annotations)split_idx = int(len(annotations) * train_size)train_data = annotations[:split_idx]val_data = annotations[split_idx:]# 创建目录结构os.makedirs(os.path.join(output_dir, "images/train"), exist_ok=True)os.makedirs(os.path.join(output_dir, "images/val"), exist_ok=True)os.makedirs(os.path.join(output_dir, "labels/train"), exist_ok=True)os.makedirs(os.path.join(output_dir, "labels/val"), exist_ok=True)# 保存训练集for data in train_data:img_path = data["image_path"]lbl_path = os.path.join(output_dir, "labels/train",os.path.splitext(os.path.basename(img_path))[0] + ".txt")# 复制图像try:shutil.copy2(img_path, os.path.join(output_dir, 'images/train'))print(f"图像 {img_path} 复制到训练集成功")except Exception as e:print(f"图像 {img_path} 复制到训练集失败，错误信息: {e}")# 保存标注with open(lbl_path, "w") as f:f.write(f"{data['class_id']} {' '.join(map(str, data['bbox']))}\n")# 保存验证集for data in val_data:img_path = data["image_path"]lbl_path = os.path.join(output_dir, "labels/val",os.path.splitext(os.path.basename(img_path))[0] + ".txt")# 复制图像try:shutil.copy2(img_path, os.path.join(output_dir, 'images/val'))print(f"图像 {img_path} 复制到验证集成功")except Exception as e:print(f"图像 {img_path} 复制到验证集失败，错误信息: {e}")# 保存标注with open(lbl_path, "w") as f:f.write(f"{data['class_id']} {' '.join(map(str, data['bbox']))}\n")# 生成类别名文件（classes.names）with open(os.path.join(output_dir, "classes.names"), "w") as f:for cls in class_map.keys():f.write(f"{cls}\n")# 生成数据集配置文件（dataset.yaml）yaml_path = os.path.join(output_dir, "dataset.yaml")with open(yaml_path, "w") as f:f.write(f"path: {output_dir}\n")  # 数据集根路径f.write(f"train: images/train\n")  # 训练集路径（相对于path）f.write(f"val: images/val\n")      # 验证集路径# f.write(f"test: images/test\n")   # 测试集路径（如果有）f.write(f"nc: {len(class_map)}\n")  # 类别数f.write("names:\n")for idx, cls in enumerate(class_map.keys()):f.write(f"  {idx}: {cls}\n")return train_data, val_data# ====================== 主函数 ======================
if __name__ == "__main__":# 1. 检查输入路径是否存在if not os.path.exists(INPUT_DATA_DIR):raise FileNotFoundError(f"请先下载数据集并解压到路径：{INPUT_DATA_DIR}")# 2. 获取类别映射（假设图像按类别存放在子文件夹中，无CSV标注时使用此方法）# 若有CSV标注，需手动指定CSV路径和列名，注释掉下方代码并取消注释parse_csv_annotations部分class_map = get_class_mapping(os.path.join(INPUT_DATA_DIR, "train"))  # 假设训练集图像在train子文件夹中，每个子文件夹为一个类别# 3. 解析标注（根据实际情况选择CSV或文件夹分类）# 情况A：无标注，仅按文件夹分类（弱监督，边界框为图像全尺寸）annotations = []for cls, idx in class_map.items():cls_dir = os.path.join(INPUT_DATA_DIR, "train", cls)  # 假设类别文件夹路径为train/类别名for img_file in os.listdir(cls_dir):if any(img_file.lower().endswith(ext) for ext in IMAGE_EXTENSIONS):img_path = os.path.join(cls_dir, img_file)with Image.open(img_path) as img:img_width, img_height = img.size# 边界框为全图（弱监督场景，仅用于分类任务，非检测）annotations.append({"image_path": img_path,"class_id": idx,"bbox": (0.5, 0.5, 1.0, 1.0)  # 全图边界框})# # 情况B：有CSV标注（需取消注释以下代码并调整参数）# CSV_PATH = os.path.join(INPUT_DATA_DIR, "labels.csv")  # CSV标注文件路径# IMAGE_DIR = os.path.join(INPUT_DATA_DIR, "images")     # 图像根目录# class_map = {"Apple Scab": 0, "Black Rot": 1, ...}    # 手动定义类别映射# annotations = parse_csv_annotations(CSV_PATH, class_map, IMAGE_DIR)# 4. 保存为YOLO格式train_data, val_data = save_dataset(annotations, class_map, OUTPUT_YOLO_DIR, train_size=TRAIN_SIZE)print(f"✅ 转换完成！YOLO数据集已保存至：{OUTPUT_YOLO_DIR}")print(f"类别数：{len(class_map)}，训练集样本数：{len(train_data)}，验证集样本数：{len(val_data)}")

train的时候，使用的yaml文件路径