將 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文件路徑