scripts/rebuild_notebook.py

"""Rebuild notebooks/02_training.ipynb from scratch with all fixes."""
import json
import os

nb = {
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
        "language_info": {"name": "python", "version": "3.10.0"},
        "colab": {"provenance": [], "gpuType": "T4"},
        "accelerator": "GPU",
    },
    "cells": [],
}


def md(cell_id, source):
    return {"cell_type": "markdown", "id": cell_id, "metadata": {}, "source": source}


def code(cell_id, source):
    return {
        "cell_type": "code",
        "id": cell_id,
        "metadata": {},
        "source": source,
        "outputs": [],
        "execution_count": None,
    }


cells = nb["cells"]

# ── Intro ────────────────────────────────────────────────────────────────────
cells.append(
    md(
        "markdown-intro",
        "# Phase 2 \u2014 YOLOv8s Fine-Tuning & ONNX Export\n"
        "\n"
        "**Run on Google Colab with T4 GPU** (Runtime \u2192 Change runtime type \u2192 T4 GPU)\n"
        "\n"
        "**Prerequisite:** `01_data_exploration.ipynb` must have been run \u2014 dataset saved to "
        "Drive at `TrafficVision/datasets/traffic_yolo/`\n"
        "\n"
        "Steps:\n"
        "1. Mount Drive & configure paths\n"
        "2. Install dependencies\n"
        "3. Verify GPU\n"
        "4. Copy dataset to Colab SSD (fast I/O via parallel zip)\n"
        "5. Fine-tune YOLOv8s (pretrained COCO \u2192 our 4-class traffic dataset)\n"
        "6. Plot training curves\n"
        "7. Evaluate \u2014 mAP, per-class metrics, confusion matrix, sample predictions\n"
        "8. Export to ONNX\n"
        "9. Verify ONNX inference speed\n"
        "10. Save weights to Drive\n"
        "\n"
        "### SSD copy strategy\n"
        "- **First run ever**: Reads dataset from Drive using 32 parallel threads (much faster "
        "than sequential FUSE reads). Creates `traffic_yolo.zip` on Drive. Takes ~5-20 min once.\n"
        "- **Every future session**: Copies the single zip file from Drive to SSD (~1-3 min) "
        "then extracts locally. No per-file FUSE latency.\n"
        "- **Training**: `cache=False` \u2014 SSD reads are ~500 MB/s so no RAM cache needed.",
    )
)

# ── 1. Mount Drive ──────────────────────────────────────────────────────────
cells.append(md("markdown-drive", "## 1. Mount Google Drive"))
cells.append(
    code(
        "cell-mount",
        'import os, sys\n'
        '\n'
        "IN_COLAB = 'google.colab' in sys.modules or 'COLAB_BACKEND_VERSION' in os.environ\n"
        '\n'
        'if IN_COLAB:\n'
        '    from google.colab import drive\n'
        "    drive.mount('/content/drive')\n"
        "    DRIVE_BASE = '/content/drive/MyDrive/TrafficVision'\n"
        'else:\n'
        "    _repo = os.path.abspath(os.path.join(os.getcwd(), '..'))\n"
        "    DRIVE_BASE = os.path.join(_repo, 'data', 'TrafficVision')\n"
        "    print('[LOCAL MODE] Using local data/ directory')\n"
        '\n'
        "DATASET_YAML = os.path.join(DRIVE_BASE, 'datasets', 'traffic_yolo', 'dataset.yaml')\n"
        "WEIGHTS_DIR  = os.path.join(DRIVE_BASE, 'weights')\n"
        "RUNS_DIR     = '/content/runs' if IN_COLAB else os.path.join("
        "_repo if not IN_COLAB else '/content', 'runs')\n"
        '\n'
        'os.makedirs(WEIGHTS_DIR, exist_ok=True)\n'
        'os.makedirs(RUNS_DIR, exist_ok=True)\n'
        '\n'
        'assert os.path.exists(DATASET_YAML), (\n'
        "    f'dataset.yaml not found at {DATASET_YAML}\\n'\n"
        "    'Run 01_data_exploration.ipynb first to prepare the dataset.'\n"
        ')\n'
        '\n'
        "print(f'Drive base   : {DRIVE_BASE}')\n"
        "print(f'Dataset yaml : {DATASET_YAML}')\n"
        "print(f'Weights dir  : {WEIGHTS_DIR}')\n"
        "print(f'Runs dir     : {RUNS_DIR}')",
    )
)

# ── 2. Install ───────────────────────────────────────────────────────────────
cells.append(md("markdown-install", "## 2. Install Dependencies"))
cells.append(code("cell-install", "!pip install -q ultralytics onnx onnxruntime-gpu"))

# ── 3. Verify GPU ───────────────────────────────────────────────────────────
cells.append(md("markdown-gpu", "## 3. Verify GPU"))
cells.append(
    code(
        "cell-gpu",
        'import torch\n'
        '\n'
        "print(f'PyTorch version : {torch.__version__}')\n"
        "print(f'CUDA available  : {torch.cuda.is_available()}')\n"
        'if torch.cuda.is_available():\n'
        '    props = torch.cuda.get_device_properties(0)\n'
        "    print(f'GPU             : {props.name}')\n"
        "    print(f'VRAM            : {props.total_memory / 1e9:.1f} GB')\n"
        'else:\n'
        "    print('WARNING: No GPU - enable GPU in Runtime > Change runtime type')\n"
        "    print('Training will be extremely slow on CPU.')",
    )
)

# ── 4. Copy dataset to SSD ──────────────────────────────────────────────────
cells.append(
    md(
        "markdown-copy-dataset",
        "## 4. Copy Dataset to Colab SSD\n"
        "\n"
        "**Why not train directly from Drive?**  \n"
        "Drive FUSE reads each small image file individually at ~0.2 MB/s with ~50ms per-file "
        "latency. 5800 images \u00d7 1.3 s/image = ~2 hours just for the RAM cache build every session.\n"
        "\n"
        "**The fix**: Copy the whole dataset as a single zip file to Colab's NVMe SSD, then "
        "train directly from SSD at 500 MB/s with `cache=False`.\n"
        "\n"
        "| Scenario | Time |\n"
        "|----------|------|\n"
        "| First run \u2014 parallel zip creation | ~5\u201320 min (once ever) |\n"
        "| Future sessions \u2014 copy zip + extract | ~1\u20133 min |\n"
        "| Training per epoch (SSD, cache=False) | ~3\u20135 min |",
    )
)

cells.append(
    code(
        "cell-copy-dataset",
        'import os, time, shutil, zipfile, threading, subprocess\n'
        'import yaml as _yaml\n'
        'from pathlib import Path\n'
        'from concurrent.futures import ThreadPoolExecutor, as_completed\n'
        '\n'
        "LOCAL_EXTRACT  = '/content/traffic_yolo_local'  # extraction target\n"
        "DRIVE_DATASETS = f'{DRIVE_BASE}/datasets'\n"
        "DRIVE_ZIP      = f'{DRIVE_DATASETS}/traffic_yolo.zip'\n"
        "LOCAL_ZIP      = '/content/traffic_yolo.zip'\n"
        "SRC_DIR        = f'{DRIVE_DATASETS}/traffic_yolo'\n"
        '\n'
        '\n'
        'def _find_dataset_root(base):\n'
        '    """Find the directory under base that contains images/."""\n'
        '    base = Path(base)\n'
        '    if (base / "images").exists():\n'
        '        return base\n'
        '    for sub in sorted(base.iterdir()):\n'
        '        if sub.is_dir() and (sub / "images").exists():\n'
        '            return sub\n'
        '    raise RuntimeError(\n'
        '        f"Cannot find images/ under {base}.\\n"\n'
        '        f"Contents: {[p.name for p in base.iterdir()]}"\n'
        '    )\n'
        '\n'
        '\n'
        'if IN_COLAB:\n'
        '    # Force Drive FUSE to populate directory listings before any exists() check.\n'
        '    # Without this, os.path.exists() can return False on files that are there.\n'
        '    subprocess.run(["ls", DRIVE_BASE],     capture_output=True)\n'
        '    subprocess.run(["ls", DRIVE_DATASETS], capture_output=True)\n'
        '    print(f"Drive datasets dir: {os.listdir(DRIVE_DATASETS) if os.path.isdir(DRIVE_DATASETS) else \"NOT FOUND\"}")\n'
        '\n'
        '    # ── Step 1: get dataset onto local SSD ───────────────────────────\n'
        '    if os.path.exists(LOCAL_EXTRACT):\n'
        "        print('Dataset already on SSD — skipping copy.')\n"
        '\n'
        '    elif os.path.exists(DRIVE_ZIP):\n'
        '        # Validate local zip — delete if missing, empty, or corrupt\n'
        '        _zip_ok = False\n'
        '        if os.path.exists(LOCAL_ZIP) and os.path.getsize(LOCAL_ZIP) > 0:\n'
        '            try:\n'
        '                import zipfile as _zfmod\n'
        '                with _zfmod.ZipFile(LOCAL_ZIP) as _ztest:\n'
        '                    _zip_ok = len(_ztest.namelist()) > 0\n'
        '            except Exception:\n'
        '                _zip_ok = False\n'
        '        if not _zip_ok:\n'
        '            if os.path.exists(LOCAL_ZIP):\n'
        '                print(f"Local zip invalid/corrupt — re-copying from Drive...")\n'
        '                os.remove(LOCAL_ZIP)\n'
        '            sz = os.path.getsize(DRIVE_ZIP) / 1e6\n'
        "            print(f'Copying zip from Drive ({sz:.0f} MB) — ~1-3 min...')\n"
        '            t0 = time.time()\n'
        '            shutil.copy2(DRIVE_ZIP, LOCAL_ZIP)\n'
        "            print(f'Copied in {time.time()-t0:.0f}s')\n"
        '        else:\n'
        "            print(f'Local zip valid — skipping Drive copy.')\n"
        "        print('Extracting to SSD...')\n"
        '        t0 = time.time()\n'
        "        subprocess.run(['unzip', '-q', LOCAL_ZIP, '-d', LOCAL_EXTRACT], check=True)\n"
        "        print(f'Extracted in {time.time()-t0:.0f}s')\n"
        '\n'
        '    else:\n'
        '        # One-time: parallel zip creation (32 threads overlap Drive FUSE latency)\n'
        "        print('One-time setup: archiving dataset with 32 parallel threads...')\n"
        "        print('(Happens ONCE — future sessions copy the zip in ~1-3 min)')\n"
        '        src   = Path(SRC_DIR)\n'
        '        files = sorted(f for f in src.rglob("*") if f.is_file())\n'
        "        print(f'Found {len(files)} files to archive...')\n"
        '        t0      = time.time()\n'
        '        lock    = threading.Lock()\n'
        '        counter = [0]\n'
        '\n'
        '        def _archive_file(f):\n'
        '            data    = f.read_bytes()\n'
        '            arcname = str(f.relative_to(src))\n'
        '            with lock:\n'
        '                _zf.writestr(arcname, data)\n'
        '                counter[0] += 1\n'
        '                if counter[0] % 1000 == 0:\n'
        "                    print(f'  {counter[0]}/{len(files)} files archived...')\n"
        '\n'
        '        with zipfile.ZipFile(LOCAL_ZIP, "w", zipfile.ZIP_STORED) as _zf:\n'
        '            with ThreadPoolExecutor(max_workers=32) as pool:\n'
        '                futures = [pool.submit(_archive_file, f) for f in files]\n'
        '                for fut in as_completed(futures):\n'
        '                    fut.result()\n'
        '\n'
        '        elapsed = time.time() - t0\n'
        '        size_mb = os.path.getsize(LOCAL_ZIP) / 1e6\n'
        "        print(f'Zip: {size_mb:.0f} MB in {elapsed:.0f}s ({elapsed/60:.1f} min)')\n"
        "        print('Saving zip to Drive for future sessions...')\n"
        '        shutil.copy2(LOCAL_ZIP, DRIVE_ZIP)\n'
        "        print('Extracting to SSD...')\n"
        "        subprocess.run(['unzip', '-q', LOCAL_ZIP, '-d', LOCAL_EXTRACT], check=True)\n"
        "        print('Done.')\n"
        '\n'
        '    # ── Step 2: find actual dataset root (handles zip subdirectory wrappers) ──\n'
        '    _root = _find_dataset_root(LOCAL_EXTRACT)\n'
        "    print(f'Dataset root: {_root}')\n"
        '\n'
        '    # ── Step 3: write dataset.yaml pointing at the correct local root ──\n'
        "    LOCAL_YAML = '/content/dataset.yaml'\n"
        '    _yaml_src = _root / "dataset.yaml"\n'
        '    if not _yaml_src.exists():\n'
        '        _yaml_src = Path(SRC_DIR) / "dataset.yaml"\n'
        '    with open(_yaml_src) as _f:\n'
        '        _cfg = _yaml.safe_load(_f)\n'
        '\n'
        '    _cfg["path"] = str(_root)  # absolute path to the correct root\n'
        '\n'
        '    # Verify/fix split paths against what actually exists on disk\n'
        '    _SPLIT_ALT = {\n'
        "        'train': ['images/train', 'train'],\n"
        "        'val':   ['images/val', 'images/valid', 'images/validation',\n"
        "                  'val', 'valid', 'validation', 'images/test', 'test'],\n"
        "        'test':  ['images/test', 'test', 'images/val', 'images/valid'],\n"
        '    }\n'
        "    for _split, _alts in _SPLIT_ALT.items():\n"
        '        _declared = _cfg.get(_split)\n'
        '        if not _declared:\n'
        '            continue\n'
        '        if (_root / _declared).exists():\n'
        '            print(f"  {_split}: {_declared!r} OK")\n'
        '        else:\n'
        '            for _alt in _alts:\n'
        '                if (_root / _alt).exists():\n'
        '                    print(f"  {_split}: {_declared!r} -> {_alt!r}")\n'
        '                    _cfg[_split] = _alt\n'
        '                    break\n'
        '            else:\n'
        '                print(f"  WARNING: no directory found for {_split!r} under {_root}")\n'
        '\n'
        '    # Final assertion — catch bad paths before training starts\n'
        '    for _split in ("train", "val"):\n'
        '        _p = _root / _cfg[_split]\n'
        '        assert _p.exists(), (\n'
        '            f"Split {_split!r} path does not exist: {_p}\\n"\n'
        '            f"Contents of {_root}: {[x.name for x in _root.iterdir()]}"\n'
        '        )\n'
        '\n'
        "    with open(LOCAL_YAML, 'w') as _f:\n"
        '        _yaml.dump(_cfg, _f, default_flow_style=False)\n'
        '    DATASET_YAML = LOCAL_YAML\n'
        "    print(f'dataset.yaml OK — path={_root}, train={_cfg[\"train\"]}, val={_cfg[\"val\"]}')\n"
        '\n'
        'else:\n'
        "    print('[LOCAL] Skipping dataset copy')",
    )
)

# ── 5. Fine-tune ────────────────────────────────────────────────────────────
cells.append(
    md(
        "markdown-train",
        "## 5. Fine-tune YOLOv8s\n"
        "\n"
        "Starting from `yolov8s.pt` (pretrained on COCO 80 classes) and fine-tuning on our "
        "4-class traffic dataset. The pretrained backbone gives us strong feature extraction "
        "out of the box; fine-tuning adapts the head to our specific classes.\n"
        "\n"
        "**`cache=False`**: Data is on local SSD (500 MB/s). No RAM cache needed \u2014 "
        "disk reads are ~instant compared to a GPU training step.\n"
        "\n"
        "**Augmentation:** `degrees=0` (fixed cameras), `flipud=0` (no upside-down vehicles), "
        "`mosaic=1.0` (small/occluded vehicles), `scale=0.5` (distance variation)\n"
        "\n"
        "**Checkpoints:** After every epoch, `last.pt` and `best.pt` are copied to Drive. "
        "If interrupted, the next run auto-resumes from Drive checkpoint.",
    )
)

cells.append(
    code(
        "cell-train",
        'import os, shutil, subprocess\n'
        'from ultralytics import YOLO\n'
        '\n'
        "MODEL_NAME = 'yolov8s_traffic'\n"
        "BEST_PT    = f'{RUNS_DIR}/{MODEL_NAME}/weights/best.pt'\n"
        "LAST_PT    = f'{RUNS_DIR}/{MODEL_NAME}/weights/last.pt'\n"
        "DRIVE_LAST = f'{WEIGHTS_DIR}/last.pt'\n"
        "DRIVE_BEST = f'{WEIGHTS_DIR}/best.pt'\n"
        '\n'
        "print(f'Training data : {DATASET_YAML}')\n"
        "assert '/content/' in DATASET_YAML or not IN_COLAB, (\n"
        "    'DATASET_YAML still points to Drive! Run cell-copy-dataset first.'\n"
        ')\n'
        '\n'
        '# Force Drive FUSE to sync before checking for checkpoints\n'
        'if IN_COLAB:\n'
        '    subprocess.run(["ls", WEIGHTS_DIR], capture_output=True)\n'
        "    print(f'Weights dir contents: {os.listdir(WEIGHTS_DIR) if os.path.isdir(WEIGHTS_DIR) else \"empty\"}')\n"
        '\n'
        'def save_to_drive(trainer):\n'
        '    if not IN_COLAB:\n'
        '        return\n'
        '    if os.path.exists(LAST_PT):\n'
        '        shutil.copy2(LAST_PT, DRIVE_LAST)\n'
        '    if os.path.exists(BEST_PT):\n'
        '        shutil.copy2(BEST_PT, DRIVE_BEST)\n'
        '    if trainer.epoch % 5 == 0:\n'
        "        print(f'[Drive] Checkpoints saved at epoch {trainer.epoch + 1}')\n"
        '\n'
        'if os.path.exists(DRIVE_LAST):\n'
        "    print(f'Resuming from {DRIVE_LAST}')\n"
        '    model = YOLO(DRIVE_LAST)\n'
        '    resume = True\n'
        'else:\n'
        "    print('Starting fresh from yolov8s.pt')\n"
        "    model = YOLO('yolov8s.pt')\n"
        '    resume = False\n'
        '\n'
        "model.add_callback('on_train_epoch_end', save_to_drive)\n"
        '\n'
        'results = model.train(\n'
        '    data=DATASET_YAML,\n'
        '    epochs=50,\n'
        '    imgsz=640,\n'
        '    batch=16,\n'
        '    cache=False,     # data is on local SSD — reads are instant, no RAM cache needed\n'
        '    device=0,\n'
        '    workers=4,\n'
        '    patience=15,\n'
        '    save=True,\n'
        '    save_period=10,\n'
        '    project=RUNS_DIR,\n'
        '    name=MODEL_NAME,\n'
        '    exist_ok=True,\n'
        '    resume=resume,\n'
        '    hsv_h=0.015,\n'
        '    hsv_s=0.7,\n'
        '    hsv_v=0.4,\n'
        '    degrees=0.0,\n'
        '    translate=0.1,\n'
        '    scale=0.5,\n'
        '    flipud=0.0,\n'
        '    fliplr=0.5,\n'
        '    mosaic=1.0,\n'
        '    mixup=0.1,\n'
        '    copy_paste=0.1,\n'
        '    plots=True,\n'
        '    verbose=True,\n'
        ')',
    )
)

# ── 6. Training Curves ──────────────────────────────────────────────────────
cells.append(md("markdown-curves", "## 6. Training Curves"))
cells.append(
    code(
        "cell-curves",
        'import pandas as pd\n'
        'import matplotlib.pyplot as plt\n'
        'import os\n'
        '\n'
        "results_csv = os.path.join(RUNS_DIR, MODEL_NAME, 'results.csv')\n"
        "assert os.path.exists(results_csv), f'results.csv not found at {results_csv}'\n"
        '\n'
        'df = pd.read_csv(results_csv)\n'
        'df.columns = df.columns.str.strip()\n'
        '\n'
        'fig, axes = plt.subplots(2, 3, figsize=(18, 10))\n'
        'axes = axes.flatten()\n'
        '\n'
        'metrics = [\n'
        "    ('train/box_loss',      'Train Box Loss',      '#e84c4c'),\n"
        "    ('train/cls_loss',      'Train Class Loss',     '#e8a44c'),\n"
        "    ('train/dfl_loss',      'Train DFL Loss',       '#8be84c'),\n"
        "    ('metrics/mAP50(B)',    'Val mAP@0.5',          '#4c8be8'),\n"
        "    ('metrics/mAP50-95(B)', 'Val mAP@0.5:0.95',    '#be84c8'),\n"
        "    ('val/box_loss',        'Val Box Loss',          '#4ce8e8'),\n"
        ']\n'
        '\n'
        'for ax, (col, title, color) in zip(axes, metrics):\n'
        '    if col in df.columns:\n'
        "        ax.plot(df['epoch'], df[col], color=color, linewidth=2)\n"
        '        ax.set_title(title, fontsize=12)\n'
        "        ax.set_xlabel('Epoch')\n"
        '        ax.grid(True, alpha=0.3)\n'
        "        best_val = df[col].min() if 'loss' in col else df[col].max()\n"
        "        best_ep  = df.loc[df[col].idxmax() if 'mAP' in col else df[col].idxmin(), 'epoch']\n"
        "        ax.axvline(best_ep, color='red', linestyle='--', alpha=0.5,\n"
        "                   label=f'Best: {best_val:.4f} @ ep{int(best_ep)}')\n"
        '        ax.legend(fontsize=9)\n'
        '    else:\n'
        "        ax.set_title(f'{title} (not found)', fontsize=10)\n"
        "        ax.axis('off')\n"
        '\n'
        "plt.suptitle('Training Curves - YOLOv8s Traffic', fontsize=15, fontweight='bold')\n"
        'plt.tight_layout()\n'
        "save_path = os.path.join(DRIVE_BASE, 'datasets', '05_training_curves.png')\n"
        "plt.savefig(save_path, dpi=120, bbox_inches='tight')\n"
        'plt.show()\n'
        "print('Saved:', save_path)",
    )
)

# ── 7. Evaluate ─────────────────────────────────────────────────────────────
cells.append(md("markdown-eval", "## 7. Evaluate on Validation Set"))
cells.append(
    code(
        "cell-eval",
        'import os\n'
        'from ultralytics import YOLO\n'
        '\n'
        "assert os.path.exists(BEST_PT), f'Best weights not found: {BEST_PT}'\n"
        '\n'
        'model   = YOLO(BEST_PT)\n'
        'metrics = model.val(\n'
        '    data=DATASET_YAML,\n'
        '    imgsz=640,\n'
        '    device=0,\n'
        '    plots=True,\n'
        '    save_json=False,\n'
        ')\n'
        '\n'
        "CLASS_NAMES = ['car', 'bus', 'motorcycle', 'truck']\n"
        '\n'
        'print()\n'
        "print('=' * 45)\n"
        "print(f'  Overall mAP@0.5     : {metrics.box.map50:.4f}')\n"
        "print(f'  Overall mAP@0.5:0.95: {metrics.box.map:.4f}')\n"
        "print(f'  Precision           : {metrics.box.mp:.4f}')\n"
        "print(f'  Recall              : {metrics.box.mr:.4f}')\n"
        "print('=' * 45)\n"
        "print('  Per-class AP@0.5:')\n"
        'for i, (name, ap) in enumerate(zip(CLASS_NAMES, metrics.box.ap50)):\n'
        "    print(f'    {name:<12}: {ap:.4f}')\n"
        "print('=' * 45)",
    )
)

# ── Per-class AP ─────────────────────────────────────────────────────────────
cells.append(md("markdown-perclass", "### Per-class AP"))
cells.append(
    code(
        "cell-perclass",
        'import matplotlib.pyplot as plt\n'
        '\n'
        "CLASS_NAMES  = ['car', 'bus', 'motorcycle', 'truck']\n"
        "CLASS_COLORS = ['#4c8be8', '#e8a44c', '#e84c4c', '#8be84c']\n"
        'ap50_vals = list(metrics.box.ap50)\n'
        '\n'
        'fig, ax = plt.subplots(figsize=(8, 5))\n'
        "bars = ax.bar(CLASS_NAMES, ap50_vals, color=CLASS_COLORS, edgecolor='none', width=0.5)\n"
        'ax.set_ylim(0, 1.05)\n'
        "ax.set_ylabel('AP @ IoU=0.50', fontsize=12)\n"
        "ax.set_title(f'Per-Class AP@0.5  |  mAP@0.5 = {metrics.box.map50:.4f}', fontsize=13)\n"
        "ax.axhline(metrics.box.map50, color='red', linestyle='--', linewidth=1.5,\n"
        "           label=f'mAP@0.5 = {metrics.box.map50:.4f}')\n"
        'ax.legend()\n'
        'for bar, val in zip(bars, ap50_vals):\n'
        '    ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,\n'
        "            f'{val:.3f}', ha='center', fontsize=11, fontweight='bold')\n"
        'plt.tight_layout()\n'
        "save_path = os.path.join(DRIVE_BASE, 'datasets', '06_per_class_ap.png')\n"
        "plt.savefig(save_path, dpi=120, bbox_inches='tight')\n"
        'plt.show()\n'
        "print('Saved:', save_path)",
    )
)

# ── Sample Predictions ───────────────────────────────────────────────────────
cells.append(md("markdown-sample-preds", "### Sample Predictions"))
cells.append(
    code(
        "cell-sample-preds",
        'import cv2, random, yaml\n'
        'import matplotlib.pyplot as plt\n'
        'import matplotlib.patches as patches\n'
        'from pathlib import Path\n'
        '\n'
        'with open(DATASET_YAML) as f:\n'
        '    ds_cfg = yaml.safe_load(f)\n'
        '\n'
        "val_img_dir = Path(ds_cfg['path']) / ds_cfg['val']\n"
        "val_imgs    = sorted(val_img_dir.glob('*.jpg'))\n"
        'sample_imgs = random.sample(val_imgs, min(6, len(val_imgs)))\n'
        '\n'
        "CLASS_NAMES = ['car', 'bus', 'motorcycle', 'truck']\n"
        '\n'
        'fig, axes = plt.subplots(2, 3, figsize=(18, 10))\n'
        'axes = axes.flatten()\n'
        '\n'
        'for ax, img_path in zip(axes, sample_imgs):\n'
        '    result  = model(img_path, verbose=False, conf=0.3)[0]\n'
        '    img_rgb = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)\n'
        '    ax.imshow(img_rgb)\n'
        '    for box in result.boxes:\n'
        '        cls_id = int(box.cls)\n'
        '        conf   = float(box.conf)\n'
        '        x1, y1, x2, y2 = box.xyxy[0].tolist()\n'
        "        color = ['#6464ff', '#ffb464', '#64ff64', '#ffc864'][cls_id % 4]\n"
        '        rect  = patches.Rectangle((x1, y1), x2 - x1, y2 - y1,\n'
        "                                   linewidth=1.5, edgecolor=color, facecolor='none')\n"
        '        ax.add_patch(rect)\n'
        "        ax.text(x1, y1 - 3, f'{CLASS_NAMES[cls_id]} {conf:.2f}',\n"
        "                color=color, fontsize=7, fontweight='bold')\n"
        '    ax.set_title(img_path.name, fontsize=8)\n'
        "    ax.axis('off')\n"
        '\n'
        "plt.suptitle('Sample Predictions on Validation Set (conf > 0.3)', fontsize=14, fontweight='bold')\n"
        'plt.tight_layout()\n'
        "save_path = os.path.join(DRIVE_BASE, 'datasets', '07_sample_predictions.png')\n"
        "plt.savefig(save_path, dpi=120, bbox_inches='tight')\n"
        'plt.show()\n'
        "print('Saved:', save_path)",
    )
)

# ── 8. ONNX Export ───────────────────────────────────────────────────────────
cells.append(
    md(
        "markdown-export",
        "## 8. Export to ONNX\n"
        "\n"
        "- `opset=17` \u2014 compatible with ONNX Runtime 1.16+\n"
        "- `simplify=True` \u2014 removes redundant ops\n"
        "- `dynamic=False` \u2014 fixed batch=1, better GPU optimization\n"
        "- `half=False` \u2014 FP32 for maximum compatibility",
    )
)

cells.append(
    code(
        "cell-export",
        'from ultralytics import YOLO\n'
        'import os\n'
        '\n'
        'model = YOLO(BEST_PT)\n'
        '\n'
        'export_path = model.export(\n'
        "    format='onnx',\n"
        '    imgsz=640,\n'
        '    opset=17,\n'
        '    simplify=True,\n'
        '    dynamic=False,\n'
        '    half=False,\n'
        ')\n'
        '\n'
        'ONNX_PATH = str(export_path)\n'
        "print(f'ONNX exported to: {ONNX_PATH}')\n"
        '\n'
        'size_mb = os.path.getsize(ONNX_PATH) / 1e6\n'
        "print(f'File size: {size_mb:.1f} MB')\n"
        "assert size_mb > 5, 'ONNX file suspiciously small - export may have failed'",
    )
)

# ── 9. Verify ONNX ──────────────────────────────────────────────────────────
cells.append(md("markdown-verify", "## 9. Verify ONNX Inference"))
cells.append(
    code(
        "cell-verify",
        'import onnxruntime as ort\n'
        'import numpy as np\n'
        'import time\n'
        '\n'
        "providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']\n"
        'sess      = ort.InferenceSession(ONNX_PATH, providers=providers)\n'
        '\n'
        "print(f'Providers active : {sess.get_providers()}')\n"
        "print(f'Input  name      : {sess.get_inputs()[0].name}')\n"
        "print(f'Input  shape     : {sess.get_inputs()[0].shape}')\n"
        "print(f'Output shape     : {sess.get_outputs()[0].shape}')\n"
        '\n'
        'dummy      = np.random.rand(1, 3, 640, 640).astype(np.float32)\n'
        "input_name = sess.get_inputs()[0].name\n"
        '\n'
        'for _ in range(3):\n'
        '    sess.run(None, {input_name: dummy})\n'
        '\n'
        'N = 100\n'
        't0 = time.perf_counter()\n'
        'for _ in range(N):\n'
        '    sess.run(None, {input_name: dummy})\n'
        'elapsed = time.perf_counter() - t0\n'
        '\n'
        "print(f'\\nBenchmark ({N} frames):')\n"
        "print(f'  Latency : {elapsed / N * 1000:.1f} ms/frame')\n"
        "print(f'  FPS     : {N / elapsed:.1f}')\n"
        "print(f'\\nONNX inference verified.')",
    )
)

# ── 10. Save to Drive ───────────────────────────────────────────────────────
cells.append(md("markdown-save", "## 10. Save Weights to Google Drive"))
cells.append(
    code(
        "cell-save",
        'import shutil\n'
        'from pathlib import Path\n'
        '\n'
        'src_pt   = Path(BEST_PT)\n'
        'src_onnx = Path(ONNX_PATH)\n'
        '\n'
        "dst_pt   = Path(WEIGHTS_DIR) / 'yolov8s_traffic_best.pt'\n"
        "dst_onnx = Path(WEIGHTS_DIR) / 'yolov8s_traffic.onnx'\n"
        '\n'
        'shutil.copy2(src_pt,   dst_pt)\n'
        'shutil.copy2(src_onnx, dst_onnx)\n'
        '\n'
        "print('Saved to Google Drive:')\n"
        "print(f'  PT   : {dst_pt}   ({dst_pt.stat().st_size / 1e6:.1f} MB)')\n"
        "print(f'  ONNX : {dst_onnx} ({dst_onnx.stat().st_size / 1e6:.1f} MB)')\n"
        '\n'
        'val_plots_dir = Path(RUNS_DIR) / MODEL_NAME\n'
        "for png in val_plots_dir.glob('*.png'):\n"
        "    dst = Path(DRIVE_BASE) / 'datasets' / png.name\n"
        '    shutil.copy2(png, dst)\n'
        "    print(f'  Plot : {dst.name}')\n"
        '\n'
        "print('\\nAll weights and plots saved to Drive.')",
    )
)

# ── Done ─────────────────────────────────────────────────────────────────────
cells.append(
    md(
        "markdown-done",
        "## Done \u2014 Phase 2 Complete\n"
        "\n"
        "### What was produced\n"
        "\n"
        "| File | Location on Drive |\n"
        "|------|-------------------|\n"
        "| Best PyTorch weights | `TrafficVision/weights/yolov8s_traffic_best.pt` |\n"
        "| ONNX model | `TrafficVision/weights/yolov8s_traffic.onnx` |\n"
        "| Training curves | `TrafficVision/datasets/05_training_curves.png` |\n"
        "| Per-class AP | `TrafficVision/datasets/06_per_class_ap.png` |\n"
        "| Sample predictions | `TrafficVision/datasets/07_sample_predictions.png` |\n"
        "| Confusion matrix | `TrafficVision/datasets/confusion_matrix.png` |\n"
        "\n"
        "### Next steps\n"
        "\n"
        "1. **Download the ONNX model** from Google Drive to your local machine:\n"
        "   ```\n"
        "   models/weights/yolov8s_traffic.onnx\n"
        "   ```\n"
        "2. **Start Phase 3** \u2014 build the core inference engine:\n"
        "   - `core/detector.py` \u2014 ONNX Runtime wrapper\n"
        "   - `core/tracker.py` \u2014 ByteTrack via supervision\n"
        "   - `core/analytics.py` \u2014 counting, speed, anomalies\n"
        "   - `core/pipeline.py` \u2014 orchestrator",
    )
)

# Write
out_path = os.path.join(os.path.dirname(__file__), '..', 'notebooks', '02_training.ipynb')
out_path = os.path.normpath(out_path)
with open(out_path, 'w', encoding='utf-8') as f:
    json.dump(nb, f, indent=1, ensure_ascii=True)

print(f"Written {len(cells)} cells to {out_path}")