detect_yolo.py

import torch
# import torchvision # Keep if needed for image loading/manipulation during conversion
# import torchvision.transforms.v2 as T # Keep if needed during conversion
from torch.utils.data import DataLoader, Dataset # Keep if needed during conversion
import os
import json
from PIL import Image
import matplotlib.pyplot as plt # Keep for visualization if desired
# import matplotlib.patches as patches
import glob
import xml.etree.ElementTree as ET
import logging
import numpy as np
# from pycocotools.coco import COCO # Keep if reading COCO format
# from pycocotools.cocoeval import COCOeval # YOLOv5 handles evaluation internally
import argparse
import yaml # Keep for reading/writing YAML
import sys
from pathlib import Path
import random
import shutil # For file operations during conversion
import subprocess # To call YOLOv5 training script

# --- 命令行参数解析 ---
parser = argparse.ArgumentParser(description='目标检测训练脚本 (YOLOv5)')
parser.add_argument('--dataset_type', type=str, default='coco', choices=['coco', 'pascal', 'yolo'],
                    help='原始数据集类型：coco, pascal, 或 yolo')
parser.add_argument('--yolov5_model', type=str, default='yolov5s',
                    help='YOLOv5 模型类型 (e.g., yolov5s, yolov5m, yolov5l, yolov5x, or path to .pt)')
parser.add_argument('--img_size', type=int, default=640, help='YOLOv5 训练图像尺寸')
parser.add_argument('--batch_size', type=int, default=16, help='训练批次大小 (adjust based on GPU memory)')
# parser.add_argument('--eval_batch_size', type=int, default=16, help='评估批次大小 (YOLOv5 handles this)')
parser.add_argument('--epochs', type=int, default=50, help='训练轮数')
parser.add_argument('--lr', type=float, default=0.01, help='初始学习率 (YOLOv5 train.py default)') # Note: YOLOv5 has its own defaults
parser.add_argument('--output_dir', type=str, default='./yolov5_training_results', help='YOLOv5 训练输出项目目录')
parser.add_argument('--run_name', type=str, default='exp', help='YOLOv5 训练运行名称')
# parser.add_argument('--logging_steps', type=int, default=50, help='日志记录步数 (YOLOv5 handles this)')
# parser.add_argument('--eval_steps', type=int, default=200, help='评估步数 (YOLOv5 handles this)')
# parser.add_argument('--save_steps', type=int, default=200, help='模型保存步数 (YOLOv5 handles this)')
parser.add_argument('--subset_ratio', type=float, default=1.0, help='使用原始数据集的比例 (用于快速转换测试, 1.0 使用全部)')
parser.add_argument('--yolov5_repo_path', type=str, default='./yolov5', help='本地 YOLOv5 代码库路径') # Important: Path to cloned yolov5 repo
parser.add_argument('--converted_data_dir', type=str, default='./yolov5_data', help='转换后 YOLOv5 格式数据的存放目录')
parser.add_argument('--workers', type=int, default=8, help='YOLOv5 dataloader workers')


args = parser.parse_args()
DATASET_TYPE = args.dataset_type
# MODEL_CHECKPOINT = args.model_checkpoint # Removed
YOLOV5_MODEL = args.yolov5_model
IMG_SIZE = args.img_size
BATCH_SIZE = args.batch_size
# EVAL_BATCH_SIZE = args.eval_batch_size # Removed
NUM_EPOCHS = args.epochs
LEARNING_RATE = args.lr # Can be passed to train.py if needed, otherwise default is used
OUTPUT_DIR = Path(args.output_dir) # Use Path object
RUN_NAME = args.run_name
# LOGGING_STEPS = args.logging_steps # Removed
# EVAL_STEPS = args.eval_steps # Removed
# SAVE_STEPS = args.save_steps # Removed
SUBSET_RATIO = args.subset_ratio
YOLOV5_REPO_PATH = Path(args.yolov5_repo_path)
CONVERTED_DATA_DIR = Path(args.converted_data_dir)
WORKERS = args.workers

print(f"选择的原始数据集类型: {DATASET_TYPE}")
logging.info(f"选择的原始数据集类型: {DATASET_TYPE}")
print(f"使用的 YOLOv5 模型: {YOLOV5_MODEL}")
logging.info(f"使用的 YOLOv5 模型: {YOLOV5_MODEL}")
print(f"转换后的数据目录: {CONVERTED_DATA_DIR}")
logging.info(f"转换后的数据目录: {CONVERTED_DATA_DIR}")
print(f"YOLOv5 代码库路径: {YOLOV5_REPO_PATH}")
logging.info(f"YOLOv5 代码库路径: {YOLOV5_REPO_PATH}")


# --- 1. 配置和超参数 (类别加载部分保留) ---

# --- 数据集路径定义 (保持不变) ---
COCO_BASE_PATH = Path('/hdd_16T/Zirui/work/CNNa2/dataset/dataset_42028assg2_24902417/Object_Detection/coco')
COCO_TRAIN_IMG_DIR = COCO_BASE_PATH / 'train'
COCO_TRAIN_ANN_FILE = COCO_BASE_PATH / 'train/train_annotations.json'
COCO_VAL_IMG_DIR = COCO_BASE_PATH / 'valid'
COCO_VAL_ANN_FILE = COCO_BASE_PATH / 'valid/valid_annotations.json'
COCO_TEST_IMG_DIR = COCO_BASE_PATH / 'test'
COCO_TEST_ANN_FILE = COCO_BASE_PATH / 'test/test_annotations.json' # Optional

PASCAL_BASE_PATH = Path('/hdd_16T/Zirui/work/CNNa2/dataset/dataset_42028assg2_24902417/Object_Detection/pascal')
PASCAL_YEAR = '2007'

YOLO_BASE_PATH = Path('/hdd_16T/Zirui/work/CNNa2/dataset/dataset_42028assg2_24902417/Object_Detection/yolo')
YOLO_CLASSES_FILE = YOLO_BASE_PATH / 'data.yaml'

# --- 类别定义 (保持不变, 用于生成 data.yaml) ---
label2id = {}
id2label = {}
NUM_CLASSES = 0
class_names = [] # Need a list of names for data.yaml

# (Keep the functions load_coco_categories, setup_pascal_categories, load_yolo_classes)
# Minor modification: Store class names in a list
def load_coco_categories(ann_file):
    global label2id, id2label, NUM_CLASSES, class_names
    with open(ann_file, 'r') as f:
        data = json.load(f)
    categories = sorted(data['categories'], key=lambda x: x['id'])
    label2id = {cat['name']: idx for idx, cat in enumerate(categories)}
    id2label = {idx: cat['name'] for idx, cat in enumerate(categories)}
    class_names = [cat['name'] for cat in categories]
    NUM_CLASSES = len(categories)
    print("COCO Categories Loaded:")
    print(f"  Num Classes: {NUM_CLASSES}")
    print(f"  Class Names: {class_names}")
    return id2label, label2id, NUM_CLASSES, class_names

def setup_pascal_categories():
    global label2id, id2label, NUM_CLASSES, class_names
    PASCAL_CLASS_NAMES = ["young", "empty_pod"]
    label2id = {name: idx for idx, name in enumerate(PASCAL_CLASS_NAMES)}
    id2label = {idx: name for idx, name in enumerate(PASCAL_CLASS_NAMES)}
    class_names = PASCAL_CLASS_NAMES
    NUM_CLASSES = len(PASCAL_CLASS_NAMES)
    print("Pascal VOC Categories Setup:")
    print(f"  Num Classes: {NUM_CLASSES}")
    print(f"  Class Names: {class_names}")
    return id2label, label2id, NUM_CLASSES, class_names

def load_yolo_classes(classes_file=None):
    global label2id, id2label, NUM_CLASSES, class_names
    classes = []
    if classes_file and classes_file.exists():
        with open(classes_file, 'r') as f:
            classes = [line.strip() for line in f.readlines() if line.strip()]
    else:
        yaml_file = YOLO_BASE_PATH / 'data.yaml'
        if yaml_file.exists():
            try:
                import yaml
                with open(yaml_file, 'r') as f:
                    data = yaml.safe_load(f)
                    if 'names' in data and isinstance(data['names'], list):
                        classes = data['names']
            except Exception as e:
                print(f"从YAML加载类别失败: {e}, 使用默认值")
        if not classes:
             classes = ['Ready', 'empty_pod', 'germination', 'pod', 'young']
             print("Using default YOLO classes.")

    label2id = {name: idx for idx, name in enumerate(classes)}
    id2label = {idx: name for idx, name in enumerate(classes)}
    class_names = classes
    NUM_CLASSES = len(classes)
    print("YOLO Categories Loaded:")
    print(f"  Num Classes: {NUM_CLASSES}")
    print(f"  Class Names: {class_names}")
    return id2label, label2id, NUM_CLASSES, class_names

# 根据数据集类型加载相应的类别信息
if DATASET_TYPE == 'coco':
    ann_file_to_load = COCO_TRAIN_ANN_FILE if COCO_TRAIN_ANN_FILE.exists() else COCO_VAL_ANN_FILE
    if ann_file_to_load.exists():
        id2label, label2id, NUM_CLASSES, class_names = load_coco_categories(ann_file_to_load)
    else:
        raise FileNotFoundError("无法找到COCO注释文件来加载类别")
elif DATASET_TYPE == 'pascal':
    id2label, label2id, NUM_CLASSES, class_names = setup_pascal_categories()
elif DATASET_TYPE == 'yolo':
    id2label, label2id, NUM_CLASSES, class_names = load_yolo_classes(YOLO_CLASSES_FILE)
else:
    raise ValueError(f"不支持的数据集类型: {DATASET_TYPE}")

if NUM_CLASSES == 0:
     raise ValueError("未能加载类别信息，类别数量为0")

print(f"最终类别数量 (Num Classes): {NUM_CLASSES}")
logging.info(f"最终类别数量 (Num Classes): {NUM_CLASSES}")
print(f"ID to Label mapping: {id2label}")
logging.info(f"ID to Label mapping: {id2label}")
print(f"Class Names List: {class_names}")
logging.info(f"Class Names List: {class_names}")

# --- 训练参数 ---
DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(f"使用的设备: {DEVICE}")
logging.info(f"使用的设备: {DEVICE}")



CONVERTED_DATA_DIR.mkdir(parents=True, exist_ok=True)
(CONVERTED_DATA_DIR / 'images' / 'train').mkdir(parents=True, exist_ok=True)
(CONVERTED_DATA_DIR / 'images' / 'val').mkdir(parents=True, exist_ok=True)
(CONVERTED_DATA_DIR / 'images' / 'test').mkdir(parents=True, exist_ok=True)
(CONVERTED_DATA_DIR / 'labels' / 'train').mkdir(parents=True, exist_ok=True)
(CONVERTED_DATA_DIR / 'labels' / 'val').mkdir(parents=True, exist_ok=True)
(CONVERTED_DATA_DIR / 'labels' / 'test').mkdir(parents=True, exist_ok=True)

def convert_to_yolo_format(image_path, annotations, img_width, img_height, label_map, output_label_path):
    """ Converts annotations for a single image to a YOLO format txt file. """
    yolo_lines = []
    for ann in annotations:
        # Assuming ann is a dict with 'bbox'=[xmin, ymin, xmax, ymax] and 'category_id' (0-based)
        label_id = ann['category_id']
        box = ann['bbox']
        
        # Convert [xmin, ymin, xmax, ymax] to YOLO format [center_x, center_y, width, height] (normalized)
        x_min, y_min, x_max, y_max = box
        center_x = ((x_min + x_max) / 2) / img_width
        center_y = ((y_min + y_max) / 2) / img_height
        width = (x_max - x_min) / img_width
        height = (y_max - y_min) / img_height

        # Clamp values to [0.0, 1.0]
        center_x = max(0.0, min(1.0, center_x))
        center_y = max(0.0, min(1.0, center_y))
        width = max(0.0, min(1.0, width))
        height = max(0.0, min(1.0, height))

        # Ensure valid box dimensions after conversion
        if width > 0 and height > 0:
            yolo_lines.append(f"{label_id} {center_x:.6f} {center_y:.6f} {width:.6f} {height:.6f}")

    # Write to label file
    if yolo_lines:
        with open(output_label_path, 'w') as f:
            f.write("\n".join(yolo_lines))
    # else: # Optional: create empty file if no objects
    #     output_label_path.touch()


def process_dataset(split, use_subset=False, subset_ratio=1.0):
    """Processes a dataset split (train, val, test) and converts it."""
    print(f"\nProcessing {split} split...")
    output_img_dir = CONVERTED_DATA_DIR / 'images' / split
    output_label_dir = CONVERTED_DATA_DIR / 'labels' / split

    items_to_process = [] # List of tuples: (img_path, annotations, width, height, original_id)

    if DATASET_TYPE == 'coco':
        img_dir = COCO_TRAIN_IMG_DIR if split == 'train' else (COCO_VAL_IMG_DIR if split == 'val' else COCO_TEST_IMG_DIR)
        ann_file = COCO_TRAIN_ANN_FILE if split == 'train' else (COCO_VAL_ANN_FILE if split == 'val' else COCO_TEST_ANN_FILE)

        if not ann_file.exists():
            print(f"Warning: Annotation file {ann_file} not found for split {split}. Skipping.")
            return

        from pycocotools.coco import COCO # Import locally if needed
        coco = COCO(str(ann_file))
        image_ids = list(sorted(coco.imgs.keys()))

        if use_subset:
             num_subset = int(len(image_ids) * subset_ratio)
             image_ids = random.sample(image_ids, num_subset)
             print(f"Using a subset of {num_subset} COCO {split} images.")

        # Create category mapping (COCO ID -> 0-based ID)
        coco_cat_id_to_label_id = {
            coco_cat_id: label2id[cat['name']]
            for coco_cat_id, cat in coco.cats.items() if cat['name'] in label2id
        }

        for img_id in image_ids:
            img_info = coco.loadImgs(img_id)[0]
            original_img_path = img_dir / img_info['file_name']
            width, height = img_info['width'], img_info['height']

            ann_ids = coco.getAnnIds(imgIds=img_id)
            coco_anns = coco.loadAnns(ann_ids)

            target_annotations = []
            for ann in coco_anns:
                if ann.get('iscrowd', 0) or ann.get('ignore', 0):
                    continue
                coco_cat_id = ann['category_id']
                if coco_cat_id in coco_cat_id_to_label_id:
                    label_id = coco_cat_id_to_label_id[coco_cat_id]
                    x, y, w, h = ann['bbox']
                    if w > 0 and h > 0:
                         target_annotations.append({
                             "bbox": [x, y, x + w, y + h],
                             "category_id": label_id,
                         })
            items_to_process.append((original_img_path, target_annotations, width, height, img_id))

    elif DATASET_TYPE == 'pascal':
        # Pascal VOC has predefined splits (trainval, test). Let's map train -> train, valid -> trainval? Or just train/val.
        # Assuming 'train', 'valid', 'test' maps to corresponding directories under PASCAL_BASE_PATH

        if split == 'val':
            data_dir = PASCAL_BASE_PATH / 'valid'
        else:
            data_dir = PASCAL_BASE_PATH / split


        if not data_dir.exists():
             print(f"Warning: Pascal directory {data_dir} not found for split {split}. Skipping.")
             return

        all_files = list(data_dir.glob('*.jpg')) + list(data_dir.glob('*.png'))
        image_files = []
        for img_file in all_files:
            xml_file = data_dir / f"{img_file.stem}.xml"
            if xml_file.exists():
                image_files.append(img_file)

        if use_subset:
             num_subset = int(len(image_files) * subset_ratio)
             image_files = random.sample(image_files, num_subset)
             print(f"Using a subset of {num_subset} Pascal {split} images.")

        name_to_id = label2id
        for index, img_path in enumerate(image_files):
            xml_path = data_dir / f"{img_path.stem}.xml"
            try:
                img = Image.open(img_path)
                img_width, img_height = img.size
                tree = ET.parse(xml_path)
                root = tree.getroot()
                target_annotations = []
                for obj in root.findall('object'):
                    label_name = obj.find('name').text
                    if label_name in name_to_id:
                        label_id = name_to_id[label_name]
                        bbox = obj.find('bndbox')
                        xmin = float(bbox.find('xmin').text)
                        ymin = float(bbox.find('ymin').text)
                        xmax = float(bbox.find('xmax').text)
                        ymax = float(bbox.find('ymax').text)
                        if xmax > xmin and ymax > ymin:
                            target_annotations.append({
                                "bbox": [xmin, ymin, xmax, ymax],
                                "category_id": label_id,
                            })
                items_to_process.append((img_path, target_annotations, img_width, img_height, index))
            except Exception as e:
                print(f"Error processing Pascal file {img_path} or {xml_path}: {e}")


    elif DATASET_TYPE == 'yolo':
        # 修正验证集目录名称
        if split == 'val':
            img_dir = YOLO_BASE_PATH / 'valid' / 'images'  # 使用'valid'而不是'val'
            label_dir = YOLO_BASE_PATH / 'valid' / 'labels'
        else:
            img_dir = YOLO_BASE_PATH / split / 'images'
            label_dir = YOLO_BASE_PATH / split / 'labels'

        if not img_dir.exists() or not label_dir.exists():
            print(f"Warning: YOLO directory {img_dir} or {label_dir} not found for split {split}. Skipping.")
            return

        img_files = sorted(list(img_dir.glob('*.jpg')) + list(img_dir.glob('*.png')))
        valid_files = []
        for img_path in img_files:
            label_path = label_dir / f"{img_path.stem}.txt"
            if label_path.exists():
                valid_files.append((img_path, label_path))

        if use_subset:
             num_subset = int(len(valid_files) * subset_ratio)
             valid_files = random.sample(valid_files, num_subset)
             print(f"Using a subset of {num_subset} YOLO {split} images.")

        print(f"Copying {len(valid_files)} YOLO {split} files...")
        for img_path, label_path in valid_files:
            # Just copy files directly as they are already in YOLO format
            shutil.copy(img_path, output_img_dir / img_path.name)
            shutil.copy(label_path, output_label_dir / label_path.name)
        # No further processing needed for YOLO -> YOLO conversion
        return # Skip the conversion loop below

    else:
        raise ValueError(f"Unsupported dataset type for conversion: {DATASET_TYPE}")

    # --- Conversion Loop ---
    print(f"Converting {len(items_to_process)} items for {split} split...")
    processed_count = 0
    for img_path, annotations, width, height, original_id in items_to_process:
        if not img_path.exists():
            print(f"Warning: Image file not found: {img_path}. Skipping.")
            continue
            
        # Define output paths
        output_img_path = output_img_dir / img_path.name
        output_label_path = output_label_dir / f"{img_path.stem}.txt"

        # 1. Copy image file
        shutil.copy(img_path, output_img_path)

        # 2. Convert annotations and write label file
        convert_to_yolo_format(img_path, annotations, width, height, label2id, output_label_path)
        processed_count += 1

    print(f"Finished converting {processed_count} items for {split} split.")


# --- Run Conversion for Train/Val/Test splits ---
use_subset_flag = SUBSET_RATIO < 1.0
process_dataset('train', use_subset=use_subset_flag, subset_ratio=SUBSET_RATIO)
process_dataset('val', use_subset=use_subset_flag, subset_ratio=SUBSET_RATIO) # Use subset for val too if desired
process_dataset('test', use_subset=False) # Usually convert full test set

# --- 3. Create data.yaml file ---
data_yaml_path = CONVERTED_DATA_DIR / 'data.yaml'
data_yaml_content = {
    'train': str((CONVERTED_DATA_DIR / 'images' / 'train').resolve()), # Use absolute paths
    'val': str((CONVERTED_DATA_DIR / 'images' / 'val').resolve()),
    'test': str((CONVERTED_DATA_DIR / 'images' / 'test').resolve()) if (CONVERTED_DATA_DIR / 'images' / 'test').iterdir() else '', # Optional test path
    'nc': NUM_CLASSES,
    'names': class_names
}

with open(data_yaml_path, 'w') as f:
    yaml.dump(data_yaml_content, f, default_flow_style=None, sort_keys=False)

print(f"\nCreated data configuration file: {data_yaml_path}")
print("Dataset conversion to YOLOv5 format complete.")


# --- 4. Execute YOLOv5 Training ---

print("\nStarting YOLOv5 training...")

# Ensure YOLOv5 repository exists
if not (YOLOV5_REPO_PATH / 'train.py').exists():
    print(f"Error: YOLOv5 train.py not found in {YOLOV5_REPO_PATH}")
    print("Please clone the YOLOv5 repository: git clone https://github.com/ultralytics/yolov5.git")
    print("And ensure --yolov5_repo_path points to it.")
    exit(1)

# Construct the training command
# Note: We use --project for the main output folder and --name for the specific run subfolder.
# YOLOv5 typically handles learning rates and schedulers internally based on defaults/hyp.yaml,
# but you can override with arguments like --lr0 if needed.
train_cmd = [
    sys.executable, # Use the current Python interpreter
    str(YOLOV5_REPO_PATH / 'train.py'),
    '--img', str(IMG_SIZE),
    '--batch', str(BATCH_SIZE),
    '--epochs', str(NUM_EPOCHS),
    '--data', str(data_yaml_path.resolve()), # Absolute path to data.yaml
    '--weights', f"{YOLOV5_MODEL}.pt" if not Path(YOLOV5_MODEL).suffix == '.pt' else YOLOV5_MODEL, # Handle model name vs path
    '--project', str(OUTPUT_DIR.resolve()), # Main output directory
    '--name', RUN_NAME, # Specific experiment name under project dir
    '--workers', str(WORKERS),
    # '--device', '0' # Specify GPU ID if needed, otherwise YOLOv5 auto-detects
    # '--lr0', str(LEARNING_RATE) # Optionally override initial LR
    '--exist-ok' # Allow reusing existing project/name directory
]

# Add device argument if using GPU
if DEVICE.type == 'cuda':
    # YOLOv5 usually expects device index, e.g., '0' or '0,1'
    gpu_index = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
    train_cmd.extend(['--device', str(gpu_index)])

print("\nExecuting command:")
print(" ".join(train_cmd))

# Execute the training command
try:
    subprocess.run(train_cmd, check=True)
    print("\nYOLOv5 training finished successfully.")
    print(f"Results saved in: {OUTPUT_DIR / RUN_NAME}")

    # --- Optional: Run validation explicitly on the best weights ---
    best_weights_path = OUTPUT_DIR / RUN_NAME / 'weights' / 'best.pt'
    if best_weights_path.exists():
        print("\nRunning validation on best weights...")
        val_cmd = [
            sys.executable,
            str(YOLOV5_REPO_PATH / 'val.py'),
            '--img', str(IMG_SIZE),
            '--batch', str(BATCH_SIZE * 2), # Can often use larger batch for validation
            '--data', str(data_yaml_path.resolve()),
            '--weights', str(best_weights_path.resolve()),
            '--project', str(OUTPUT_DIR.resolve()),
            '--name', f'{RUN_NAME}_val_best',
            '--task', 'val',
             '--workers', str(WORKERS),
             '--device', str(gpu_index) if DEVICE.type == 'cuda' else 'cpu',
            '--exist-ok'
        ]
        print("\nExecuting command:")
        print(" ".join(val_cmd))
        subprocess.run(val_cmd, check=True)
        print(f"Validation results saved in: {OUTPUT_DIR / f'{RUN_NAME}_val_best'}")
    else:
        print("Could not find best.pt for validation.")

except subprocess.CalledProcessError as e:
    print(f"\nYOLOv5 training/validation failed with error code {e.returncode}")
    print(e)
except FileNotFoundError:
     print(f"Error: Could not find Python interpreter or train.py/val.py script.")
     print("Ensure Python is in your PATH and --yolov5_repo_path is correct.")
     
     
     
     
# #python your_modified_script_name.py \
#     --dataset_type coco \ # 或 pascal, yolo
#     --yolov5_model yolov5s \
#     --img_size 640 \
#     --batch_size 16 \
#     --epochs 50 \
#     --output_dir ./yolov5_finetuned_coco \
#     --run_name coco_run1 \
#     --yolov5_repo_path ./yolov5 \ # 指向克隆的 yolov5 文件夹
#     --converted_data_dir ./coco_yolov5_format \
#     --subset_ratio 0.1 # 可选，使用10%数据快速测试