实验2：YOLOv8 一阶段水平框检测

# pathlib：跨平台处理项目路径；json：保存实验清单；os：读取环境变量；
# sys：把项目源码加入导入路径；time：测速；platform：记录运行平台。
from pathlib import Path
import json, os, sys, time, platform

ROOT = Path.cwd()
while not (ROOT / "pyproject.toml").exists() and ROOT != ROOT.parent:
    ROOT = ROOT.parent
sys.path.insert(0, str(ROOT / "src"))
from rsdet.notebook_utils import (
    set_seed,
    build_coco_teaching_subset,
    build_yolo_teaching_subset,
    build_dota_teaching_subset,
    read_yolo_names,
    coco_summary,
    print_manifest,
    find_first_image,
)

SEED = int(os.getenv("SEED", "42"))
set_seed(SEED)
RUN_TRAINING = bool(int(os.getenv("RUN_TRAINING", "0")))
# Notebook 读取完整数据源，再写出可追溯的教学子集；原始目录不被覆盖。
SUBSET_MAX_IMAGES = int(os.getenv("SUBSET_MAX_IMAGES", "3000"))
SUBSET_MAX_CLASSES = int(os.getenv("SUBSET_MAX_CLASSES", "6"))
DATA_ROOT = Path(os.getenv("DATA_ROOT", str(ROOT / "data" / "processed" / "dior")))
RUN_ROOT = ROOT / "runs" / "notebook_outputs"
RUN_ROOT.mkdir(parents=True, exist_ok=True)
ARTIFACTS = {"seed": SEED, "platform": platform.platform(), "python": platform.python_version()}
print({"root": str(ROOT), "run_training": RUN_TRAINING, "subset_max_images": SUBSET_MAX_IMAGES, "subset_max_classes": SUBSET_MAX_CLASSES, "seed": SEED})

def save_json(name, payload):
    """将阶段性结果写入 runs/notebook_outputs，方便报告引用和复核。"""
    path = RUN_ROOT / name
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
    ARTIFACTS[name] = str(path)
    return path

def plot_history(history, name="training_curves.png"):
    """读取训练历史并绘制 loss、学习率曲线；无历史时只给出提示。"""
    if not history:
        print("暂无训练历史；设置 RUN_TRAINING=1 后会保存曲线。")
        return None
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("matplotlib 未安装，跳过曲线绘制。")
        return None
    epochs = [row.get("epoch", i + 1) for i, row in enumerate(history)]
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    axes[0].plot(epochs, [row.get("loss", float("nan")) for row in history], marker="o")
    axes[0].set_title("loss"); axes[0].set_xlabel("epoch"); axes[0].grid(alpha=.3)
    axes[1].plot(epochs, [row.get("lr", float("nan")) for row in history], marker="o", color="darkorange")
    axes[1].set_title("learning rate"); axes[1].set_xlabel("epoch"); axes[1].grid(alpha=.3)
    fig.tight_layout()
    path = RUN_ROOT / name; fig.savefig(path, dpi=150); plt.close(fig)
    ARTIFACTS[name] = str(path)
    print("曲线已保存:", path)
    return path

# Ultralytics 提供 YOLO 模型、训练、验证、推理和导出 API。
print("Ultralytics 训练建议: ultralytics>=8,<9；记录输入尺寸、batch、epoch、seed 和 GPU。")

# 从完整 YOLO 目录中选择不超过 SUBSET_MAX_IMAGES 个图像—标签对。
SOURCE_ROOT = Path(os.getenv("DATASET_ROOT", str(DATA_ROOT)))
SOURCE_IMAGE_DIR = Path(os.getenv("IMAGE_DIR", str(SOURCE_ROOT / "images")))
SOURCE_LABEL_DIR = Path(os.getenv("LABEL_DIR", str(SOURCE_ROOT / "labels")))
SOURCE_DATA_YAML = Path(os.getenv("DATA_YAML", str(SOURCE_ROOT / "data.yaml")))
subset_info = build_yolo_teaching_subset(
    SOURCE_IMAGE_DIR, SOURCE_LABEL_DIR, RUN_ROOT / "yolov8_subset",
    max_images=SUBSET_MAX_IMAGES, seed=SEED,
)
DATASET_ROOT = Path(subset_info["root"])
DATA_YAML = DATASET_ROOT / "data.yaml"
names = [x.strip() for x in os.getenv("YOLO_NAMES", "").split(",") if x.strip()] or read_yolo_names(SOURCE_DATA_YAML)
DATA_YAML.write_text(
    f"path: {DATASET_ROOT}\ntrain: images/train\nval: images/val\n"
    f"names: {names!r}\nnc: {len(names)}\n", encoding="utf-8"
)
ARTIFACTS["subset_manifest"] = str(subset_info["manifest"])
print("完整 YOLO 数据已处理为教学子集:", subset_info, "配置:", DATA_YAML)

# 标签统计用于发现空标签、类别数量异常和数据路径错误，不改变原始数据。
label_files = list((DATASET_ROOT / "labels").rglob("*.txt"))
label_count = sum(len(p.read_text(encoding="utf-8").splitlines()) for p in label_files)
print({"label_files": len(label_files), "objects": label_count, "yaml": str(DATA_YAML)})

dataset_root = DATASET_ROOT
def pair_count(image_dir, label_dir):
    """按相对路径检查图像与 YOLO 标签是否一一对应。"""
    image_dir, label_dir = Path(image_dir), Path(label_dir)
    image_stems = {p.relative_to(image_dir).with_suffix("") for p in image_dir.rglob("*") if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".tif", ".tiff"}}
    label_stems = {p.relative_to(label_dir).with_suffix("") for p in label_dir.rglob("*.txt")}
    # 若数据按 train/val/test 分目录，只比较对应 split；避免遗留的根目录标签被误报。
    split_names = {"train", "val", "test", "trainval"}
    image_splits = {key.parts[0] for key in image_stems if key.parts and key.parts[0] in split_names}
    if image_splits:
        label_stems = {key for key in label_stems if key.parts and key.parts[0] in image_splits}
    return {"images": len(image_stems), "labels": len(label_stems), "missing_labels": sorted(image_stems - label_stems)[:10], "orphan_labels": sorted(label_stems - image_stems)[:10]}
pair_summary = pair_count(dataset_root / "images", dataset_root / "labels")
print("图像—标签配对检查:", pair_summary)
if pair_summary["missing_labels"] or pair_summary["orphan_labels"]:
    raise ValueError("发现图像与标签文件名不匹配，请先修复数据集目录。")

try:
    # YOLO 类封装模型图结构、预训练权重以及 train/val/predict 方法。
    from ultralytics import YOLO
    model = YOLO(os.getenv("YOLO_WEIGHTS", "yolov8n.pt")); print("模型构建完成:", model.model.__class__.__name__ if hasattr(model, "model") else "YOLO")
except ImportError as exc:
    YOLO = None; model = None; print("Ultralytics 未安装:", exc)

# 将训练超参数集中保存，便于复现实验和将命令写入报告。
TRAIN_KWARGS = {"data": str(DATA_YAML), "epochs": int(os.getenv("EPOCHS", "1")), "imgsz": int(os.getenv("IMGSZ", "640")), "batch": int(os.getenv("BATCH", "2")), "workers": int(os.getenv("WORKERS", "0")), "seed": SEED, "project": str(RUN_ROOT), "name": "yolov8_train", "exist_ok": True, "plots": True}
print("训练参数:", TRAIN_KWARGS)

# train_result 保存 Ultralytics 的运行目录；results.csv 用于后续曲线绘制。
train_result = None; metrics = None; history = []
if model is not None and RUN_TRAINING:
    train_result = model.train(**TRAIN_KWARGS)
    train_dir = Path(getattr(train_result, "save_dir", RUN_ROOT / "yolov8_train")); ARTIFACTS["train_dir"] = str(train_dir)
    result_csv = train_dir / "results.csv"
    if result_csv.exists():
        import csv
        with result_csv.open(encoding="utf-8") as f: history = list(csv.DictReader(f))
    save_json("yolov8_history.json", history)
else:
    print("RUN_TRAINING=0 或 Ultralytics 未安装：跳过训练。")

# val() 返回 Precision、Recall、mAP 等验证指标；不把训练 loss 当作检测精度。
validation = {}
if model is not None and RUN_TRAINING:
    metrics = model.val(data=str(DATA_YAML), imgsz=TRAIN_KWARGS["imgsz"], plots=True)
    validation = getattr(metrics, "results_dict", {})
    print("验证指标:", validation)
save_json("yolov8_metrics.json", validation)

# 推理使用真实数据目录中的首张图像，仅作为流程示例，不代表模型总体性能。
first_image = find_first_image(DATASET_ROOT / "images")
inference = {"status": "skipped"}
if model is not None and RUN_TRAINING and first_image:
    prediction = model.predict(source=str(first_image), conf=.25, save=True, project=str(RUN_ROOT), name="yolov8_predict", exist_ok=True)
    inference = {"image": str(first_image), "prediction_count": len(prediction), "output": str(RUN_ROOT / "yolov8_predict")}
    ARTIFACTS["prediction_dir"] = inference["output"]
print("推理记录:", inference); save_json("yolov8_inference.json", inference)

if history:
    try:
        import matplotlib.pyplot as plt
        import pandas as pd
        frame = pd.DataFrame(history); numeric = frame.select_dtypes("number")
        ax = numeric.plot(subplots=True, figsize=(10, max(4, len(numeric.columns) * 1.8)), grid=True, title="YOLOv8 training history")
        fig = ax[0].figure if hasattr(ax, "__len__") else ax.figure; fig.tight_layout(); curve = RUN_ROOT / "yolov8_curves.png"; fig.savefig(curve, dpi=150); plt.close(fig); ARTIFACTS["curve"] = str(curve)
    except ImportError: print("pandas/matplotlib 未安装，跳过曲线。")
else: print("暂无训练历史；真实训练后将读取 results.csv 绘制曲线。")

print("建议对照：保持权重、epoch、batch、数据划分和 seed 不变，仅比较 imgsz=640 与 imgsz=1024。")
print("最终产物:", ARTIFACTS); print_manifest(RUN_ROOT, **ARTIFACTS)