added yolov8s vision model training notebook

.gitignore

@@ -1,3 +1,5 @@
 **/*.env
 **/__pycache__/
-*.pyc
+*.pyc
+*.jpg
+*.png

ModelTrain/main.ipynb

@@ -0,0 +1,665 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1370c868",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "YOLO MODEL TRAINING PIPELINE FOR URBAN ISSUES DETECTION\n",
+    "================================================================================\n",
+    "Autonomous City Issue Resolution Agent\n",
+    "\n",
+    "This file is divided into sections. Copy each section into a separate \n",
+    "Jupyter Notebook cell and run sequentially.\n",
+    "\n",
+    "Classes:\n",
+    "    0: Damaged Road Issues\n",
+    "    1: Pothole Issues  \n",
+    "    2: Illegal Parking Issues\n",
+    "    3: Broken Road Sign Issues\n",
+    "    4: Fallen Trees\n",
+    "    5: Littering/Garbage on Public Places\n",
+    "    6: Vandalism Issues\n",
+    "    7: Dead Animal Pollution\n",
+    "    8: Damaged Concrete Structures\n",
+    "    9: Damaged Electric Wires and Poles\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "\n",
+    "\n",
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 1: IMPORTS AND SETUP\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "import os\n",
+    "import shutil\n",
+    "import yaml\n",
+    "import random\n",
+    "from pathlib import Path\n",
+    "from tqdm import tqdm\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import cv2\n",
+    "from PIL import Image\n",
+    "import pandas as pd\n",
+    "import seaborn as sns\n",
+    "from ultralytics import YOLO\n",
+    "import torch\n",
+    "\n",
+    "BASE_DIR = Path(r\"D:\\B.Tech\\ProjeX\\24.HackTheThrone\")\n",
+    "DATASET_DIR = BASE_DIR / \"Dataset\"\n",
+    "MODEL_DIR = BASE_DIR / \"Model\"\n",
+    "MERGED_DIR = BASE_DIR / \"Dataset_Merged\"\n",
+    "IMG_SIZE = 640\n",
+    "\n",
+    "print(f\"PyTorch version: {torch.__version__}\")\n",
+    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
+    "print(f\"CUDA version: {torch.version.cuda}\")\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
+    "    print(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")\n",
+    "    DEVICE = 0\n",
+    "else:\n",
+    "    print(\"\\n[!] GPU NOT DETECTED - Training will use CPU (SLOW)\")\n",
+    "    print(\"    To install PyTorch 2.9.0 with CUDA 13.0 (Windows/Linux):\")\n",
+    "    print(\"    pip uninstall torch torchvision torchaudio -y\")\n",
+    "    print(\"    pip cache purge\")\n",
+    "    print(\"    pip install torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0 --index-url https://download.pytorch.org/whl/cu130\")\n",
+    "    print(\"    (For CPU-only or other CUDA builds like cu126/cu128, see https://pytorch.org/get-started/locally/)\")\n",
+    "    DEVICE = \"cpu\"\n",
+    "\n",
+    "print(f\"\\nBase directory: {BASE_DIR}\")\n",
+    "print(f\"Device: {DEVICE}\")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d32f748e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 2: DATASET CONFIGURATION\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "CLASS_NAMES = {\n",
+    "    0: \"Damaged Road Issues\",\n",
+    "    1: \"Pothole Issues\",\n",
+    "    2: \"Illegal Parking Issues\",\n",
+    "    3: \"Broken Road Sign Issues\",\n",
+    "    4: \"Fallen Trees\",\n",
+    "    5: \"Littering/Garbage on Public Places\",\n",
+    "    6: \"Vandalism Issues\",\n",
+    "    7: \"Dead Animal Pollution\",\n",
+    "    8: \"Damaged Concrete Structures\",\n",
+    "    9: \"Damaged Electric Wires and Poles\"\n",
+    "}\n",
+    "\n",
+    "DATASET_MAPPING = {\n",
+    "    \"Potholes and RoadCracks/Potholes and RoadCracks\": {\n",
+    "        \"class_map\": {1: 1},\n",
+    "        \"name\": \"Potholes\"\n",
+    "    },\n",
+    "    \"Garbage/Garbage\": {\n",
+    "        \"class_map\": {5: 5},\n",
+    "        \"name\": \"Garbage\"\n",
+    "    },\n",
+    "    \"FallenTrees/FallenTrees\": {\n",
+    "        \"class_map\": {4: 4},\n",
+    "        \"name\": \"FallenTrees\"\n",
+    "    },\n",
+    "    \"DamagedElectricalPoles/DamagedElectricalPoles\": {\n",
+    "        \"class_map\": {9: 9},\n",
+    "        \"name\": \"DamagedElectricalPoles\"\n",
+    "    },\n",
+    "    \"Damaged concrete structures/Damaged concrete structures\": {\n",
+    "        \"class_map\": {8: 8},\n",
+    "        \"name\": \"DamagedConcrete\"\n",
+    "    },\n",
+    "    \"DamagedRoadSigns/DamagedRoadSigns\": {\n",
+    "        \"class_map\": {0: 3, 1: 3},\n",
+    "        \"name\": \"DamagedRoadSigns\"\n",
+    "    },\n",
+    "    \"DeadAnimalsPollution/DeadAnimalsPollution\": {\n",
+    "        \"class_map\": {7: 7},\n",
+    "        \"name\": \"DeadAnimals\"\n",
+    "    },\n",
+    "    \"Graffitti/Graffitti\": {\n",
+    "        \"class_map\": {6: 6},\n",
+    "        \"name\": \"Graffiti\"\n",
+    "    },\n",
+    "    \"IllegalParking/IllegalParking\": {\n",
+    "        \"class_map\": {0: 2, 1: 2, 2: 2},\n",
+    "        \"name\": \"IllegalParking\"\n",
+    "    }\n",
+    "}\n",
+    "\n",
+    "print(f\"Total classes: {len(CLASS_NAMES)}\")\n",
+    "for idx, name in CLASS_NAMES.items():\n",
+    "    print(f\"  {idx}: {name}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "756043d6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 3: CREATE MERGED DATASET DIRECTORY STRUCTURE\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "def create_merged_dataset_structure():\n",
+    "    for split in [\"train\", \"valid\", \"test\"]:\n",
+    "        (MERGED_DIR / \"images\" / split).mkdir(parents=True, exist_ok=True)\n",
+    "        (MERGED_DIR / \"labels\" / split).mkdir(parents=True, exist_ok=True)\n",
+    "    print(f\"Created merged dataset structure at: {MERGED_DIR}\")\n",
+    "\n",
+    "create_merged_dataset_structure()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f33064a9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 4: MERGE AND RELABEL DATASETS\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "def relabel_annotation(label_path, class_map, is_segmentation=False):\n",
+    "    if not label_path.exists():\n",
+    "        return None\n",
+    "    \n",
+    "    with open(label_path, 'r') as f:\n",
+    "        lines = f.readlines()\n",
+    "    \n",
+    "    new_lines = []\n",
+    "    for line in lines:\n",
+    "        line = line.strip()\n",
+    "        if not line:\n",
+    "            continue\n",
+    "        \n",
+    "        parts = line.split()\n",
+    "        if len(parts) < 5:\n",
+    "            continue\n",
+    "            \n",
+    "        old_class = int(parts[0])\n",
+    "        if old_class not in class_map:\n",
+    "            continue\n",
+    "            \n",
+    "        new_class = class_map[old_class]\n",
+    "        \n",
+    "        if len(parts) == 5:\n",
+    "            cx, cy, bw, bh = map(float, parts[1:5])\n",
+    "        else:\n",
+    "            coords = list(map(float, parts[1:]))\n",
+    "            xs = coords[0::2]\n",
+    "            ys = coords[1::2]\n",
+    "            x_min, x_max = min(xs), max(xs)\n",
+    "            y_min, y_max = min(ys), max(ys)\n",
+    "            cx = (x_min + x_max) / 2.0\n",
+    "            cy = (y_min + y_max) / 2.0\n",
+    "            bw = x_max - x_min\n",
+    "            bh = y_max - y_min\n",
+    "        \n",
+    "        new_line = f\"{new_class} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}\"\n",
+    "        new_lines.append(new_line)\n",
+    "    \n",
+    "    return new_lines if new_lines else None\n",
+    "\n",
+    "def merge_datasets():\n",
+    "    stats = {split: {cls: 0 for cls in CLASS_NAMES.keys()} for split in [\"train\", \"valid\", \"test\"]}\n",
+    "    \n",
+    "    for dataset_path, config in tqdm(DATASET_MAPPING.items(), desc=\"Merging datasets\"):\n",
+    "        dataset_full_path = DATASET_DIR / dataset_path\n",
+    "        class_map = config[\"class_map\"]\n",
+    "        dataset_name = config[\"name\"]\n",
+    "        \n",
+    "        for split in [\"train\", \"valid\", \"test\"]:\n",
+    "            images_dir = dataset_full_path / split / \"images\"\n",
+    "            labels_dir = dataset_full_path / split / \"labels\"\n",
+    "            \n",
+    "            if not images_dir.exists():\n",
+    "                print(f\"Skipping {dataset_name}/{split} - images not found\")\n",
+    "                continue\n",
+    "            \n",
+    "            image_files = list(images_dir.glob(\"*.[jJ][pP][gG]\")) + \\\n",
+    "                         list(images_dir.glob(\"*.[jJ][pP][eE][gG]\")) + \\\n",
+    "                         list(images_dir.glob(\"*.[pP][nN][gG]\"))\n",
+    "            \n",
+    "            for img_path in image_files:\n",
+    "                label_name = img_path.stem + \".txt\"\n",
+    "                label_path = labels_dir / label_name\n",
+    "                \n",
+    "                new_lines = relabel_annotation(label_path, class_map)\n",
+    "                if new_lines is None:\n",
+    "                    continue\n",
+    "                \n",
+    "                new_img_name = f\"{dataset_name}_{img_path.name}\"\n",
+    "                new_label_name = f\"{dataset_name}_{img_path.stem}.txt\"\n",
+    "                \n",
+    "                dst_img = MERGED_DIR / \"images\" / split / new_img_name\n",
+    "                dst_label = MERGED_DIR / \"labels\" / split / new_label_name\n",
+    "                \n",
+    "                shutil.copy2(img_path, dst_img)\n",
+    "                \n",
+    "                with open(dst_label, 'w') as f:\n",
+    "                    f.write('\\n'.join(new_lines))\n",
+    "                \n",
+    "                for line in new_lines:\n",
+    "                    cls = int(line.split()[0])\n",
+    "                    stats[split][cls] += 1\n",
+    "    \n",
+    "    return stats\n",
+    "\n",
+    "def merged_dataset_exists():\n",
+    "    for split in [\"train\", \"valid\", \"test\"]:\n",
+    "        images_dir = MERGED_DIR / \"images\" / split\n",
+    "        labels_dir = MERGED_DIR / \"labels\" / split\n",
+    "        if not images_dir.exists() or not labels_dir.exists():\n",
+    "            return False\n",
+    "        if not list(images_dir.glob(\"*\")) or not list(labels_dir.glob(\"*.txt\")):\n",
+    "            return False\n",
+    "    return True\n",
+    "\n",
+    "if merged_dataset_exists():\n",
+    "    print(f\"Merged dataset already exists at: {MERGED_DIR}\")\n",
+    "    print(\"Skipping dataset merge.\")\n",
+    "    stats = None\n",
+    "else:\n",
+    "    print(\"Merging datasets...\")\n",
+    "    stats = merge_datasets()\n",
+    "    \n",
+    "    print(\"\\n Dataset Statistics:\")\n",
+    "    for split, class_stats in stats.items():\n",
+    "        total = sum(class_stats.values())\n",
+    "        print(f\"\\n{split.upper()}: {total} annotations\")\n",
+    "        for cls, count in class_stats.items():\n",
+    "            if count > 0:\n",
+    "                print(f\"  {cls}: {CLASS_NAMES[cls]} - {count}\")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "71e9fa04",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 5: GENERATE DATA.YAML CONFIG FILE\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "data_yaml = {\n",
+    "    \"path\": str(MERGED_DIR),\n",
+    "    \"train\": \"images/train\",\n",
+    "    \"val\": \"images/valid\",\n",
+    "    \"test\": \"images/test\",\n",
+    "    \"nc\": len(CLASS_NAMES),\n",
+    "    \"names\": list(CLASS_NAMES.values())\n",
+    "}\n",
+    "\n",
+    "yaml_path = MERGED_DIR / \"data.yaml\"\n",
+    "with open(yaml_path, 'w') as f:\n",
+    "    yaml.dump(data_yaml, f, default_flow_style=False)\n",
+    "\n",
+    "print(f\"Created data.yaml at: {yaml_path}\")\n",
+    "print(\"\\nContents:\")\n",
+    "with open(yaml_path, 'r') as f:\n",
+    "    print(f.read())\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "680070a4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 6: VALIDATE DATASET INTEGRITY\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "def validate_dataset():\n",
+    "    issues = []\n",
+    "    \n",
+    "    for split in [\"train\", \"valid\", \"test\"]:\n",
+    "        images_dir = MERGED_DIR / \"images\" / split\n",
+    "        labels_dir = MERGED_DIR / \"labels\" / split\n",
+    "        \n",
+    "        image_files = list(images_dir.glob(\"*\"))\n",
+    "        label_files = list(labels_dir.glob(\"*.txt\"))\n",
+    "        \n",
+    "        image_stems = {f.stem for f in image_files}\n",
+    "        label_stems = {f.stem for f in label_files}\n",
+    "        \n",
+    "        missing_labels = image_stems - label_stems\n",
+    "        missing_images = label_stems - image_stems\n",
+    "        \n",
+    "        if missing_labels:\n",
+    "            issues.append(f\"{split}: {len(missing_labels)} images missing labels\")\n",
+    "        if missing_images:\n",
+    "            issues.append(f\"{split}: {len(missing_images)} labels missing images\")\n",
+    "        \n",
+    "        print(f\"{split}: {len(image_files)} images, {len(label_files)} labels\")\n",
+    "    \n",
+    "    if issues:\n",
+    "        print(\"\\n Issues found:\")\n",
+    "        for issue in issues:\n",
+    "            print(f\"  - {issue}\")\n",
+    "    else:\n",
+    "        print(\"\\n All files validated successfully!\")\n",
+    "    \n",
+    "    return len(issues) == 0\n",
+    "\n",
+    "validate_dataset()\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3e0466f7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 7: VISUALIZE SAMPLE IMAGES WITH ANNOTATIONS\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "def visualize_samples(n_samples=6):\n",
+    "    train_images = list((MERGED_DIR / \"images\" / \"train\").glob(\"*\"))\n",
+    "    sample_images = random.sample(train_images, min(n_samples, len(train_images)))\n",
+    "    \n",
+    "    fig, axes = plt.subplots(2, 3, figsize=(15, 10))\n",
+    "    axes = axes.flatten()\n",
+    "    \n",
+    "    colors = plt.cm.tab10(np.linspace(0, 1, 10))\n",
+    "    \n",
+    "    for idx, img_path in enumerate(sample_images):\n",
+    "        img = cv2.imread(str(img_path))\n",
+    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
+    "        h, w = img.shape[:2]\n",
+    "        \n",
+    "        label_path = MERGED_DIR / \"labels\" / \"train\" / (img_path.stem + \".txt\")\n",
+    "        \n",
+    "        if label_path.exists():\n",
+    "            with open(label_path, 'r') as f:\n",
+    "                for line in f:\n",
+    "                    parts = line.strip().split()\n",
+    "                    if len(parts) >= 5:\n",
+    "                        cls = int(parts[0])\n",
+    "                        \n",
+    "                        if len(parts) == 5:\n",
+    "                            cx, cy, bw, bh = map(float, parts[1:5])\n",
+    "                            x1 = int((cx - bw/2) * w)\n",
+    "                            y1 = int((cy - bh/2) * h)\n",
+    "                            x2 = int((cx + bw/2) * w)\n",
+    "                            y2 = int((cy + bh/2) * h)\n",
+    "                        else:\n",
+    "                            coords = list(map(float, parts[1:]))\n",
+    "                            xs = [coords[i] * w for i in range(0, len(coords), 2)]\n",
+    "                            ys = [coords[i] * h for i in range(1, len(coords), 2)]\n",
+    "                            x1, x2 = int(min(xs)), int(max(xs))\n",
+    "                            y1, y2 = int(min(ys)), int(max(ys))\n",
+    "                        \n",
+    "                        color = tuple(int(c * 255) for c in colors[cls][:3])\n",
+    "                        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)\n",
+    "                        cv2.putText(img, CLASS_NAMES[cls][:15], (x1, y1-5), \n",
+    "                                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
+    "        \n",
+    "        axes[idx].imshow(img)\n",
+    "        axes[idx].set_title(img_path.name[:30])\n",
+    "        axes[idx].axis('off')\n",
+    "    \n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(MODEL_DIR / \"sample_annotations.png\", dpi=150)\n",
+    "    plt.show()\n",
+    "    print(f\"Saved sample visualization to: {MODEL_DIR / 'sample_annotations.png'}\")\n",
+    "\n",
+    "visualize_samples()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8ebfcace",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 9: EVALUATE MODEL ON VALIDATION SET\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "best_model_path = MODEL_DIR / \"urban_issues_yolov8\" / \"weights\" / \"best.pt\"\n",
+    "model = YOLO(str(best_model_path))\n",
+    "\n",
+    "val_results = model.val(\n",
+    "    data=str(MERGED_DIR / \"data.yaml\"),\n",
+    "    split=\"val\",\n",
+    "    imgsz=IMG_SIZE,\n",
+    "    batch=32,\n",
+    "    device=DEVICE,\n",
+    "    workers=4,\n",
+    "    save_json=True,\n",
+    "    plots=True\n",
+    ")\n",
+    "\n",
+    "print(\"\\nValidation Results:\")\n",
+    "print(f\"mAP50: {val_results.box.map50:.4f}\")\n",
+    "print(f\"mAP50-95: {val_results.box.map:.4f}\")\n",
+    "print(f\"Precision: {val_results.box.mp:.4f}\")\n",
+    "print(f\"Recall: {val_results.box.mr:.4f}\")\n",
+    "\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "86cb025d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 10: RUN INFERENCE ON TEST SET\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "test_images_dir = MERGED_DIR / \"images\" / \"test\"\n",
+    "test_images = list(test_images_dir.glob(\"*\"))[:20]\n",
+    "\n",
+    "results = model.predict(\n",
+    "    source=test_images,\n",
+    "    save=True,\n",
+    "    save_txt=True,\n",
+    "    project=str(MODEL_DIR),\n",
+    "    name=\"test_predictions\",\n",
+    "    exist_ok=True,\n",
+    "    conf=0.25,\n",
+    "    iou=0.45\n",
+    ")\n",
+    "\n",
+    "print(f\"\\nTest predictions saved to: {MODEL_DIR / 'test_predictions'}\")\n",
+    "\n",
+    "for r in results[:3]:\n",
+    "    print(f\"\\nImage: {Path(r.path).name}\")\n",
+    "    if r.boxes is not None:\n",
+    "        for box in r.boxes:\n",
+    "            cls = int(box.cls[0])\n",
+    "            conf = float(box.conf[0])\n",
+    "            print(f\"  - {CLASS_NAMES[cls]}: {conf:.2f}\")\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3e333a9d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 11: GENERATE CONFUSION MATRIX AND ANALYSIS\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "confusion_matrix_path = MODEL_DIR / \"urban_issues_yolov8\" / \"confusion_matrix.png\"\n",
+    "if confusion_matrix_path.exists():\n",
+    "    img = Image.open(confusion_matrix_path)\n",
+    "    plt.figure(figsize=(12, 10))\n",
+    "    plt.imshow(img)\n",
+    "    plt.axis('off')\n",
+    "    plt.title(\"Confusion Matrix\")\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()\n",
+    "else:\n",
+    "    print(\"Confusion matrix not found. Will be generated after training.\")\n",
+    "\n",
+    "results_csv = MODEL_DIR / \"urban_issues_yolov8\" / \"results.csv\"\n",
+    "if results_csv.exists():\n",
+    "    df = pd.read_csv(results_csv)\n",
+    "    df.columns = df.columns.str.strip()\n",
+    "    \n",
+    "    fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
+    "    \n",
+    "    if 'train/box_loss' in df.columns:\n",
+    "        axes[0, 0].plot(df['epoch'], df['train/box_loss'], label='Box Loss')\n",
+    "        axes[0, 0].plot(df['epoch'], df['train/cls_loss'], label='Class Loss')\n",
+    "        axes[0, 0].set_xlabel('Epoch')\n",
+    "        axes[0, 0].set_ylabel('Loss')\n",
+    "        axes[0, 0].set_title('Training Losses')\n",
+    "        axes[0, 0].legend()\n",
+    "        axes[0, 0].grid(True)\n",
+    "    \n",
+    "    if 'metrics/mAP50(B)' in df.columns:\n",
+    "        axes[0, 1].plot(df['epoch'], df['metrics/mAP50(B)'], label='mAP50')\n",
+    "        axes[0, 1].plot(df['epoch'], df['metrics/mAP50-95(B)'], label='mAP50-95')\n",
+    "        axes[0, 1].set_xlabel('Epoch')\n",
+    "        axes[0, 1].set_ylabel('mAP')\n",
+    "        axes[0, 1].set_title('Validation mAP')\n",
+    "        axes[0, 1].legend()\n",
+    "        axes[0, 1].grid(True)\n",
+    "    \n",
+    "    if 'metrics/precision(B)' in df.columns:\n",
+    "        axes[1, 0].plot(df['epoch'], df['metrics/precision(B)'], label='Precision')\n",
+    "        axes[1, 0].plot(df['epoch'], df['metrics/recall(B)'], label='Recall')\n",
+    "        axes[1, 0].set_xlabel('Epoch')\n",
+    "        axes[1, 0].set_ylabel('Score')\n",
+    "        axes[1, 0].set_title('Precision & Recall')\n",
+    "        axes[1, 0].legend()\n",
+    "        axes[1, 0].grid(True)\n",
+    "    \n",
+    "    if 'val/box_loss' in df.columns:\n",
+    "        axes[1, 1].plot(df['epoch'], df['val/box_loss'], label='Val Box Loss')\n",
+    "        axes[1, 1].plot(df['epoch'], df['val/cls_loss'], label='Val Class Loss')\n",
+    "        axes[1, 1].set_xlabel('Epoch')\n",
+    "        axes[1, 1].set_ylabel('Loss')\n",
+    "        axes[1, 1].set_title('Validation Losses')\n",
+    "        axes[1, 1].legend()\n",
+    "        axes[1, 1].grid(True)\n",
+    "    \n",
+    "    plt.tight_layout()\n",
+    "    plt.savefig(MODEL_DIR / \"training_metrics.png\", dpi=150)\n",
+    "    plt.show()\n",
+    "    print(f\"Saved training metrics to: {MODEL_DIR / 'training_metrics.png'}\")\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad0f65a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "================================================================================\n",
+    "SECTION 12: EXPORT OPTIMIZED MODEL\n",
+    "================================================================================\n",
+    "\"\"\"\n",
+    "best_model = YOLO(str(best_model_path))\n",
+    "\n",
+    "onnx_path = best_model.export(format=\"onnx\", simplify=True, dynamic=False)\n",
+    "print(f\"Exported ONNX model to: {onnx_path}\")\n",
+    "\n",
+    "torchscript_path = best_model.export(format=\"torchscript\")\n",
+    "print(f\"Exported TorchScript model to: {torchscript_path}\")\n",
+    "\n",
+    "model_info = {\n",
+    "    \"model_name\": \"Urban Issues YOLOv8 Detector\",\n",
+    "    \"version\": \"1.0\",\n",
+    "    \"classes\": CLASS_NAMES,\n",
+    "    \"input_size\": IMG_SIZE,\n",
+    "    \"best_weights\": str(best_model_path),\n",
+    "    \"onnx_path\": str(onnx_path),\n",
+    "    \"torchscript_path\": str(torchscript_path),\n",
+    "    \"val_mAP50\": float(val_results.box.map50),\n",
+    "    \"val_mAP50_95\": float(val_results.box.map)\n",
+    "}\n",
+    "\n",
+    "with open(MODEL_DIR / \"model_info.yaml\", 'w') as f:\n",
+    "    yaml.dump(model_info, f, default_flow_style=False)\n",
+    "\n",
+    "print(f\"\\nModel info saved to: {MODEL_DIR / 'model_info.yaml'}\")\n",
+    "print(\"\\n\" + \"=\"*60)\n",
+    "print(\"TRAINING PIPELINE COMPLETE!\")\n",
+    "print(\"=\"*60)\n",
+    "print(f\"Best model: {best_model_path}\")\n",
+    "print(f\"ONNX export: {onnx_path}\")\n",
+    "print(f\"Validation mAP50: {val_results.box.map50:.4f}\")\n",
+    "\"\"\"\n",
+    "================================================================================\n",
+    "END OF PIPELINE\n",
+    "================================================================================\n",
+    "\"\"\"\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}