{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实验5：轻量化 YOLO 与 ONNX 推理\n",
    "\n",
    "流程：**读取完整数据 → 构建轻量化教学子集 → 数据检查 → 构建数据集 → 构建模型变体 → 构建训练器 → 训练 → 验证与指标 → 测试集推理 → ONNX 导出 → 延迟/FPS 测试 → loss 与精度速度曲线**。当前方案以 YOLOv8n 为轻量基线，并对比输入尺寸或 YOLOv8s；原始数据只读，训练使用可追溯的小规模子集。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 参数与环境准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pathlib：跨平台处理项目路径；json：保存实验清单；os：读取环境变量；\n",
    "# sys：把项目源码加入导入路径；time：测速；platform：记录运行平台。\n",
    "from pathlib import Path\n",
    "import json, os, sys, time, platform\n",
    "\n",
    "ROOT = Path.cwd()\n",
    "while not (ROOT / \"pyproject.toml\").exists() and ROOT != ROOT.parent:\n",
    "    ROOT = ROOT.parent\n",
    "sys.path.insert(0, str(ROOT / \"src\"))\n",
    "from rsdet.notebook_utils import (\n",
    "    set_seed,\n",
    "    build_coco_teaching_subset,\n",
    "    build_yolo_teaching_subset,\n",
    "    build_dota_teaching_subset,\n",
    "    read_yolo_names,\n",
    "    coco_summary,\n",
    "    print_manifest,\n",
    "    find_first_image,\n",
    ")\n",
    "\n",
    "SEED = int(os.getenv(\"SEED\", \"42\"))\n",
    "set_seed(SEED)\n",
    "RUN_TRAINING = bool(int(os.getenv(\"RUN_TRAINING\", \"0\")))\n",
    "# Notebook 读取完整数据源，再写出可追溯的教学子集；原始目录不被覆盖。\n",
    "SUBSET_MAX_IMAGES = int(os.getenv(\"SUBSET_MAX_IMAGES\", \"3000\"))\n",
    "SUBSET_MAX_CLASSES = int(os.getenv(\"SUBSET_MAX_CLASSES\", \"6\"))\n",
    "DATA_ROOT = Path(os.getenv(\"DATA_ROOT\", str(ROOT / \"data\" / \"processed\" / \"dior\")))\n",
    "RUN_ROOT = ROOT / \"runs\" / \"notebook_outputs\"\n",
    "RUN_ROOT.mkdir(parents=True, exist_ok=True)\n",
    "ARTIFACTS = {\"seed\": SEED, \"platform\": platform.platform(), \"python\": platform.python_version()}\n",
    "print({\"root\": str(ROOT), \"run_training\": RUN_TRAINING, \"subset_max_images\": SUBSET_MAX_IMAGES, \"subset_max_classes\": SUBSET_MAX_CLASSES, \"seed\": SEED})\n",
    "\n",
    "def save_json(name, payload):\n",
    "    \"\"\"将阶段性结果写入 runs/notebook_outputs，方便报告引用和复核。\"\"\"\n",
    "    path = RUN_ROOT / name\n",
    "    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding=\"utf-8\")\n",
    "    ARTIFACTS[name] = str(path)\n",
    "    return path\n",
    "\n",
    "def plot_history(history, name=\"training_curves.png\"):\n",
    "    \"\"\"读取训练历史并绘制 loss、学习率曲线；无历史时只给出提示。\"\"\"\n",
    "    if not history:\n",
    "        print(\"暂无训练历史；设置 RUN_TRAINING=1 后会保存曲线。\")\n",
    "        return None\n",
    "    try:\n",
    "        import matplotlib.pyplot as plt\n",
    "    except ImportError:\n",
    "        print(\"matplotlib 未安装，跳过曲线绘制。\")\n",
    "        return None\n",
    "    epochs = [row.get(\"epoch\", i + 1) for i, row in enumerate(history)]\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "    axes[0].plot(epochs, [row.get(\"loss\", float(\"nan\")) for row in history], marker=\"o\")\n",
    "    axes[0].set_title(\"loss\"); axes[0].set_xlabel(\"epoch\"); axes[0].grid(alpha=.3)\n",
    "    axes[1].plot(epochs, [row.get(\"lr\", float(\"nan\")) for row in history], marker=\"o\", color=\"darkorange\")\n",
    "    axes[1].set_title(\"learning rate\"); axes[1].set_xlabel(\"epoch\"); axes[1].grid(alpha=.3)\n",
    "    fig.tight_layout()\n",
    "    path = RUN_ROOT / name; fig.savefig(path, dpi=150); plt.close(fig)\n",
    "    ARTIFACTS[name] = str(path)\n",
    "    print(\"曲线已保存:\", path)\n",
    "    return path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 DIOR 数据准备：DIOR → YOLO 格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 轻量化实验仍需使用真实图像和标签，只减少教学运行的图像对数量。\n",
    "SOURCE_ROOT = Path(os.getenv(\"DATASET_ROOT\", str(DATA_ROOT)))\n",
    "SOURCE_IMAGE_DIR = Path(os.getenv(\"IMAGE_DIR\", str(SOURCE_ROOT / \"images\")))\n",
    "SOURCE_LABEL_DIR = Path(os.getenv(\"LABEL_DIR\", str(SOURCE_ROOT / \"labels\")))\n",
    "SOURCE_DATA_YAML = Path(os.getenv(\"DATA_YAML\", str(SOURCE_ROOT / \"data.yaml\")))\n",
    "subset_info = build_yolo_teaching_subset(\n",
    "    SOURCE_IMAGE_DIR, SOURCE_LABEL_DIR, RUN_ROOT / \"lightweight_subset\",\n",
    "    max_images=SUBSET_MAX_IMAGES, seed=SEED,\n",
    ")\n",
    "DATASET_ROOT = Path(subset_info[\"root\"])\n",
    "DATA_YAML = DATASET_ROOT / \"data.yaml\"\n",
    "names = [x.strip() for x in os.getenv(\"YOLO_NAMES\", \"\").split(\",\") if x.strip()] or read_yolo_names(SOURCE_DATA_YAML)\n",
    "DATA_YAML.write_text(\n",
    "    f\"path: {DATASET_ROOT}\\ntrain: images/train\\nval: images/val\\n\"\n",
    "    f\"names: {names!r}\\nnc: {len(names)}\\n\", encoding=\"utf-8\"\n",
    ")\n",
    "ARTIFACTS[\"subset_manifest\"] = str(subset_info[\"manifest\"])\n",
    "MODEL_VARIANTS = [x.strip() for x in os.getenv(\"MODEL_VARIANTS\", \"yolov8n.pt,yolov8s.pt\").split(\",\") if x.strip()]\n",
    "IMGSZ_LIST = [int(x) for x in os.getenv(\"IMGSZ_LIST\", \"320,640\").split(\",\")]\n",
    "print(\"完整数据已处理为轻量化教学子集:\", subset_info, {\"data\": str(DATA_YAML), \"models\": MODEL_VARIANTS, \"imgsz\": IMGSZ_LIST})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 数据检查与标签统计"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 统计标签文件和目标行数，只读检查，不修改数据集。\n",
    "label_dir = DATASET_ROOT / \"labels\"\n",
    "label_files = list(label_dir.rglob(\"*.txt\")) if label_dir.exists() else []\n",
    "print({\"label_files\": len(label_files), \"objects\": sum(len(x.read_text(encoding=\"utf-8\").splitlines()) for x in label_files)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 构建数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_root = DATASET_ROOT\n",
    "image_files = [p for p in (dataset_root / \"images\").rglob(\"*\") if p.suffix.lower() in {\".jpg\", \".jpeg\", \".png\", \".tif\", \".tiff\"}] if (dataset_root / \"images\").exists() else []\n",
    "label_files = list((dataset_root / \"labels\").rglob(\"*.txt\")) if (dataset_root / \"labels\").exists() else []\n",
    "image_stems = {p.relative_to(dataset_root / \"images\").with_suffix(\"\") for p in image_files}\n",
    "label_stems = {p.relative_to(dataset_root / \"labels\").with_suffix(\"\") for p in label_files}\n",
    "# 数据按 split 分目录时，仅保留 train/val/test 中与图像同层级的标签。\n",
    "split_names = {\"train\", \"val\", \"test\", \"trainval\"}\n",
    "image_splits = {key.parts[0] for key in image_stems if key.parts and key.parts[0] in split_names}\n",
    "if image_splits:\n",
    "    label_stems = {key for key in label_stems if key.parts and key.parts[0] in image_splits}\n",
    "pair_summary = {\"images\": len(image_files), \"labels\": len(label_files), \"missing_labels\": sorted(image_stems - label_stems)[:10], \"orphan_labels\": sorted(label_stems - image_stems)[:10]}\n",
    "print(\"轻量化实验数据配对:\", pair_summary)\n",
    "if pair_summary[\"missing_labels\"] or pair_summary[\"orphan_labels\"]:\n",
    "    raise ValueError(\"发现图像与标签文件名不匹配，请先修复数据集目录。\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 构建模型变体"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    # YOLO 封装轻量模型的构建、训练、验证、推理和 ONNX 导出接口。\n",
    "    from ultralytics import YOLO\n",
    "    ULTRALYTICS_READY = True\n",
    "except ImportError as exc:\n",
    "    ULTRALYTICS_READY = False; print(\"Ultralytics 未安装:\", exc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 构建训练器与测速函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def benchmark_model(model, image_path, imgsz, repeats=10):\n",
    "    \"\"\"预热后固定图像和后端，返回平均延迟（毫秒）与 FPS。\"\"\"\n",
    "    for _ in range(2): model.predict(source=str(image_path), imgsz=imgsz, verbose=False)\n",
    "    t0 = time.perf_counter()\n",
    "    for _ in range(repeats): model.predict(source=str(image_path), imgsz=imgsz, verbose=False)\n",
    "    latency = (time.perf_counter() - t0) / repeats\n",
    "    return {\"latency_ms\": latency * 1000, \"fps\": 1 / max(latency, 1e-9)}\n",
    "print(\"测速原则：warmup 后固定图像、输入尺寸、batch 和后端。\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7 开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_rows = []; models = {}\n",
    "# 只从完整数据目录选取示例图像，不根据示例图推断总体性能。\n",
    "first_image = find_first_image(DATASET_ROOT / \"images\")\n",
    "if RUN_TRAINING and ULTRALYTICS_READY and first_image:\n",
    "    for variant in MODEL_VARIANTS:\n",
    "        model = YOLO(variant); models[variant] = model\n",
    "        model.train(data=str(DATA_YAML), epochs=int(os.getenv(\"EPOCHS\", \"1\")), imgsz=IMGSZ_LIST[0], batch=int(os.getenv(\"BATCH\", \"2\")), workers=0, seed=SEED, project=str(RUN_ROOT), name=f\"light_{Path(variant).stem}\", exist_ok=True, plots=True)\n",
    "else: print(\"RUN_TRAINING=0、依赖缺失或没有图像：跳过训练。\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8 验证集评估与指标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for variant, model in models.items():\n",
    "    validation = model.val(data=str(DATA_YAML), imgsz=IMGSZ_LIST[0], plots=True)\n",
    "    print(variant, getattr(validation, \"results_dict\", {}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9 测试集推理与 ONNX 导出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for variant, model in models.items():\n",
    "    if first_image: model.predict(source=str(first_image), save=True, project=str(RUN_ROOT), name=f\"predict_{Path(variant).stem}\", exist_ok=True)\n",
    "    exported = model.export(format=\"onnx\", simplify=False); print(variant, \"ONNX:\", exported)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10 延迟、FPS 与精度速度曲线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if models and first_image:\n",
    "    for variant, model in models.items():\n",
    "        for imgsz in IMGSZ_LIST: benchmark_rows.append({\"model\": variant, \"imgsz\": imgsz, **benchmark_model(model, first_image, imgsz)})\n",
    "save_json(\"lightweight_benchmark.json\", benchmark_rows)\n",
    "print(\"基准结果:\", benchmark_rows)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11 结论与部署排错"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"报告需同时比较 mAP、Recall、参数量、模型文件大小、延迟、FPS、输入尺寸和部署后端。\")\n",
    "print_manifest(RUN_ROOT, **ARTIFACTS)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}