{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "992db2dd",
   "metadata": {},
   "source": [
    "# 实验1：Faster R-CNN 水平框检测\n",
    "\n",
    "本实验按完整检测工程组织：**数据准备 → 数据检查 → 构建数据集 → 构建模型 → 构建训练器 → 开始训练 → 评估验证集 → 推理与可视化 → 性能曲线**。\n",
    "\n",
    "`RUN_TRAINING=0` 只执行数据检查和模型构建；`RUN_TRAINING=1` 才启动训练。教学子集仅用于控制课堂运行规模，不替代完整数据集的正式实验。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f1de7c2",
   "metadata": {},
   "source": [
    "## 1 实验参数与环境准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21025c60",
   "metadata": {},
   "outputs": [],
   "source": [
    "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 set_seed, build_coco_teaching_subset, coco_summary, print_manifest, find_first_image\n",
    "\n",
    "SEED = int(os.getenv(\"SEED\", \"42\"))\n",
    "set_seed(SEED)\n",
    "RUN_TRAINING = bool(int(os.getenv(\"RUN_TRAINING\", \"0\")))\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",
    "    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",
    "    \"\"\"Save loss/learning-rate curves when a real run produced history.\"\"\"\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": "code",
   "execution_count": null,
   "id": "7ee4397c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"建议环境：Python 3.9、PyTorch 2.0、torchvision 0.15；GPU 训练优先使用 CUDA。\")\n",
    "print(\"本次运行目录:\", RUN_ROOT.resolve())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74f31110",
   "metadata": {},
   "source": [
    "## 2 DIOR 数据准备：DIOR → COCO 格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b5812fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SOURCE_COCO_JSON 是 DIOR 转换得到的完整 COCO 标注；函数只读取源数据，\n",
    "# 再把图像—标注配对复制到 RUN_ROOT 下的教学子集目录。\n",
    "SOURCE_COCO_JSON = Path(os.getenv(\"COCO_JSON\", str(DATA_ROOT / \"annotations.json\")))\n",
    "SOURCE_IMAGE_DIR = Path(os.getenv(\"IMAGE_DIR\", str(DATA_ROOT / \"images\")))\n",
    "if not SOURCE_COCO_JSON.exists():\n",
    "    candidates = [DATA_ROOT / \"annotations\" / \"train.json\", DATA_ROOT / \"train.json\"]\n",
    "    SOURCE_COCO_JSON = next((path for path in candidates if path.exists()), SOURCE_COCO_JSON)\n",
    "if not SOURCE_COCO_JSON.exists():\n",
    "    raise FileNotFoundError(f\"DIOR 的 COCO 标注不存在: {SOURCE_COCO_JSON.resolve()}\")\n",
    "if not SOURCE_IMAGE_DIR.exists():\n",
    "    raise FileNotFoundError(f\"DIOR 图像目录不存在: {SOURCE_IMAGE_DIR.resolve()}\")\n",
    "SUBSET_ROOT = RUN_ROOT / \"faster_rcnn_subset\"\n",
    "subset_info = build_coco_teaching_subset(\n",
    "    SOURCE_COCO_JSON, SUBSET_ROOT / \"annotations.json\",\n",
    "    max_images=SUBSET_MAX_IMAGES, max_classes=SUBSET_MAX_CLASSES, seed=SEED,\n",
    "    source_image_dir=SOURCE_IMAGE_DIR,\n",
    ")\n",
    "COCO_JSON = Path(subset_info[\"annotation\"])\n",
    "IMAGE_DIR = SUBSET_ROOT\n",
    "ARTIFACTS[\"subset_manifest\"] = str(subset_info[\"manifest\"])\n",
    "summary = coco_summary(COCO_JSON)\n",
    "ARTIFACTS[\"coco_summary\"] = summary\n",
    "print(\"DIOR → COCO 教学子集准备完成:\", subset_info, summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8a42f8a",
   "metadata": {},
   "source": [
    "## 3 数据检查与标注可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1360fa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "with COCO_JSON.open(encoding=\"utf-8\") as stream:\n",
    "    coco = json.load(stream)\n",
    "print(\"图像数:\", len(coco[\"images\"]), \"标注数:\", len(coco[\"annotations\"]), \"类别:\", coco[\"categories\"])\n",
    "ann_counts = {}\n",
    "for ann in coco[\"annotations\"]:\n",
    "    ann_counts[ann[\"category_id\"]] = ann_counts.get(ann[\"category_id\"], 0) + 1\n",
    "print(\"类别实例统计:\", ann_counts)\n",
    "first = find_first_image(IMAGE_DIR)\n",
    "if first:\n",
    "    try:\n",
    "        from PIL import Image\n",
    "        import matplotlib.pyplot as plt\n",
    "        image = Image.open(first).convert(\"RGB\")\n",
    "        plt.figure(figsize=(5, 5)); plt.imshow(image); plt.title(first.name); plt.axis(\"off\"); plt.show()\n",
    "    except ImportError:\n",
    "        print(\"Pillow 或 matplotlib 未安装，跳过图像显示。\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4705c23d",
   "metadata": {},
   "source": [
    "## 4 构建数据集与 DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8363141e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "try:\n",
    "    import torch\n",
    "    from torch.utils.data import Dataset, DataLoader\n",
    "    TORCH_READY = True\n",
    "except ImportError as exc:\n",
    "    torch = None; Dataset = object; DataLoader = None; TORCH_READY = False\n",
    "    print(\"PyTorch 未安装:\", exc)\n",
    "\n",
    "cat_ids = sorted(item[\"id\"] for item in coco[\"categories\"])\n",
    "cat_to_label = {cat_id: i + 1 for i, cat_id in enumerate(cat_ids)}\n",
    "images = {item[\"id\"]: item for item in coco[\"images\"]}\n",
    "ann_by_image = {image_id: [] for image_id in images}\n",
    "for ann in coco[\"annotations\"]: ann_by_image.setdefault(ann[\"image_id\"], []).append(ann)\n",
    "\n",
    "class CocoDetection(Dataset):\n",
    "    \"\"\"把 COCO xywh 转为 torchvision xyxy；不在这里做随机增强，便于复核。\"\"\"\n",
    "    def __init__(self, image_dir, records, annotations): self.image_dir, self.records, self.annotations = Path(image_dir), records, annotations\n",
    "    def __len__(self): return len(self.records)\n",
    "    def __getitem__(self, index):\n",
    "        record = self.records[index]\n",
    "        image = torch.as_tensor(np.asarray(Image.open(self.image_dir / record[\"file_name\"]).convert(\"RGB\")).transpose(2, 0, 1), dtype=torch.float32) / 255.0\n",
    "        boxes, labels, areas = [], [], []\n",
    "        for ann in self.annotations.get(record[\"id\"], []):\n",
    "            x, y, w, h = ann[\"bbox\"]; boxes.append([x, y, x + w, y + h]); labels.append(cat_to_label[ann[\"category_id\"]]); areas.append(w * h)\n",
    "        target = {\"boxes\": torch.tensor(boxes, dtype=torch.float32), \"labels\": torch.tensor(labels, dtype=torch.int64), \"image_id\": torch.tensor([record[\"id\"]]), \"area\": torch.tensor(areas, dtype=torch.float32), \"iscrowd\": torch.zeros(len(boxes), dtype=torch.int64)}\n",
    "        return image, target\n",
    "\n",
    "def collate_fn(batch): return tuple(zip(*batch))\n",
    "records = list(images.values()); split = max(1, int(len(records) * .8))\n",
    "if TORCH_READY:\n",
    "    train_ds = CocoDetection(IMAGE_DIR, records[:split], ann_by_image); val_ds = CocoDetection(IMAGE_DIR, records[split:] or records[:1], ann_by_image)\n",
    "    train_loader = DataLoader(train_ds, batch_size=int(os.getenv(\"BATCH\", \"2\")), shuffle=True, collate_fn=collate_fn)\n",
    "    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, collate_fn=collate_fn)\n",
    "    print(\"训练/验证图像数:\", len(train_ds), len(val_ds))\n",
    "else:\n",
    "    train_ds = val_ds = train_loader = val_loader = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70585f2e",
   "metadata": {},
   "source": [
    "## 5 构建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41895e84",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = None; device = None\n",
    "if TORCH_READY:\n",
    "    try:\n",
    "        from torchvision.models.detection import fasterrcnn_resnet50_fpn\n",
    "        from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights\n",
    "        from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "        weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT if os.getenv(\"USE_PRETRAINED\", \"0\") == \"1\" else None\n",
    "        model = fasterrcnn_resnet50_fpn(weights=weights, min_size=320, max_size=640)\n",
    "        model.roi_heads.box_predictor = FastRCNNPredictor(model.roi_heads.box_predictor.cls_score.in_features, len(cat_ids) + 1)\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"); model.to(device)\n",
    "        print(\"模型构建完成:\", type(model).__name__, \"device=\", device)\n",
    "    except ImportError as exc:\n",
    "        print(\"torchvision detection 未就绪:\", exc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3b2e747",
   "metadata": {},
   "source": [
    "## 6 构建训练器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b7e25e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = scheduler = None\n",
    "if model is not None:\n",
    "    optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=float(os.getenv(\"LR\", \"0.005\")), momentum=.9, weight_decay=.0005)\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=.1)\n",
    "    print(\"优化器:\", type(optimizer).__name__, \"初始学习率:\", optimizer.param_groups[0][\"lr\"])\n",
    "\n",
    "def train_one_epoch(model, loader, optimizer, device):\n",
    "    model.train(); losses = []\n",
    "    for image_batch, target_batch in loader:\n",
    "        image_batch = [x.to(device) for x in image_batch]; target_batch = [{k: v.to(device) for k, v in t.items()} for t in target_batch]\n",
    "        loss_dict = model(image_batch, target_batch); loss = sum(loss_dict.values())\n",
    "        optimizer.zero_grad(); loss.backward(); optimizer.step(); losses.append(float(loss.detach().cpu()))\n",
    "    return sum(losses) / max(1, len(losses))\n",
    "\n",
    "def evaluate_predictions(model, loader, device):\n",
    "    model.eval(); count = 0; boxes = 0\n",
    "    with torch.no_grad():\n",
    "        for image_batch, _ in loader:\n",
    "            output = model([x.to(device) for x in image_batch]); count += len(output); boxes += sum(len(item[\"boxes\"]) for item in output)\n",
    "    return {\"images\": count, \"predictions\": boxes}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e1daff1",
   "metadata": {},
   "source": [
    "## 7 开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f32016c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = []\n",
    "if model is None:\n",
    "    print(\"模型不可用，完成数据链路检查后跳过训练。\")\n",
    "elif RUN_TRAINING:\n",
    "    for epoch in range(int(os.getenv(\"EPOCHS\", \"1\"))):\n",
    "        loss = train_one_epoch(model, train_loader, optimizer, device); scheduler.step()\n",
    "        row = {\"epoch\": epoch + 1, \"loss\": loss, \"lr\": optimizer.param_groups[0][\"lr\"]}; history.append(row); print(row)\n",
    "    checkpoint = RUN_ROOT / \"faster_rcnn_checkpoint.pt\"\n",
    "    torch.save({\"model\": model.state_dict(), \"categories\": summary[\"category_names\"], \"seed\": SEED}, checkpoint)\n",
    "    ARTIFACTS[\"checkpoint\"] = str(checkpoint)\n",
    "else:\n",
    "    print(\"RUN_TRAINING=0：跳过长时间训练。\")\n",
    "save_json(\"faster_rcnn_history.json\", history)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dfb5976",
   "metadata": {},
   "source": [
    "## 8 评估验证集与计算指标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8662b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation = {}\n",
    "if model is not None:\n",
    "    evaluation = evaluate_predictions(model, val_loader, device)\n",
    "    print(\"教学子集验证集预测统计:\", evaluation)\n",
    "else:\n",
    "    print(\"无模型输出；正式运行时应接入 COCOeval，报告 mAP50、mAP50:95、Precision、Recall。\")\n",
    "save_json(\"faster_rcnn_evaluation.json\", evaluation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba09f0bf",
   "metadata": {},
   "source": [
    "## 9 测试集推理与结果可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d9ab54",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_record = {\"status\": \"skipped\", \"reason\": \"模型未训练或依赖未安装\"}\n",
    "if model is not None and RUN_TRAINING and first:\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        image = torch.as_tensor(np.asarray(Image.open(first).convert(\"RGB\")).transpose(2, 0, 1), dtype=torch.float32).to(device) / 255.0\n",
    "        output = model([image])[0]\n",
    "    inference_record = {\"image\": str(first), \"boxes\": int(len(output[\"boxes\"])), \"score_threshold\": .25}\n",
    "    print(\"推理结果:\", inference_record)\n",
    "save_json(\"faster_rcnn_inference.json\", inference_record)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b69035",
   "metadata": {},
   "source": [
    "## 10 绘制损失、学习率与结果曲线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd06f73b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_history(history, \"faster_rcnn_curves.png\")\n",
    "print(\"最终产物:\", ARTIFACTS)\n",
    "print_manifest(RUN_ROOT, **ARTIFACTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d3ba1fd",
   "metadata": {},
   "source": [
    "### 结果解释\n",
    "\n",
    "报告需区分数据检查、教学子集训练和完整数据正式测试三类结果。正式结果至少包含逐类别 AP、mAP50、mAP50:95、Precision、Recall、FPS、显存、错误样本和运行命令。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}