{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": { "provenance": [], "gpuType": "T4" },
  "kernelspec": { "name": "python3", "display_name": "Python 3" },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# XRayVision AI — YOLOv8 Fracture Detection Training\n",
    "**FYP: Web-Based AI Chatbot for Automated X-Ray Fracture Detection**  \n",
    "Minhaj University Lahore — BSSE 8th Semester\n",
    "\n",
    "---\n",
    "\n",
    "## Before running: pick a dataset source\n",
    "\n",
    "### Option A — Roboflow (Recommended)\n",
    "1. Sign up free at [roboflow.com](https://roboflow.com)\n",
    "2. Go to [universe.roboflow.com](https://universe.roboflow.com)\n",
    "3. Search **bone fracture detection** → open any dataset\n",
    "4. Click **Download** → choose **YOLOv8 format** → download zip\n",
    "5. Upload the zip to Google Drive\n",
    "6. **Run cells: 1 → 2 → 3 → 4A → 5 → 6 → 7 → 8**\n",
    "\n",
    "### Option B — FracAtlas (No signup needed)\n",
    "1. Go to [huggingface.co/datasets/FracAtlas/FracAtlas](https://huggingface.co/datasets/FracAtlas/FracAtlas)\n",
    "2. Click **Files and versions** → download the zip\n",
    "3. Upload zip to Google Drive\n",
    "4. **Run cells: 1 → 2 → 3 → 4B → 5 → 6 → 7 → 8**\n",
    "\n",
    "---\n",
    "**First: Runtime → Change runtime type → T4 GPU → Save**"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 1: Verify GPU ────────────────────────────────────────────\n",
    "!nvidia-smi\n",
    "import torch\n",
    "print(f'CUDA available : {torch.cuda.is_available()}')\n",
    "print(f'GPU            : {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"NONE — go to Runtime > Change runtime type > T4\"}')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 2: Install dependencies ──────────────────────────────────\n",
    "!pip install ultralytics -q\n",
    "from ultralytics import YOLO\n",
    "print('ultralytics OK')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 3: Mount Google Drive ────────────────────────────────────\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "print('Drive mounted')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 4A: Extract Roboflow dataset (Option A) ──────────────────\n",
    "# !! Skip this cell if using FracAtlas (Option B) !!\n",
    "import os, zipfile\n",
    "\n",
    "# Change this to your zip file path in Google Drive\n",
    "ZIP_PATH    = '/content/drive/MyDrive/bone-fracture-detection.zip'\n",
    "EXTRACT_DIR = '/content/dataset'\n",
    "\n",
    "os.makedirs(EXTRACT_DIR, exist_ok=True)\n",
    "print(f'Extracting {ZIP_PATH} ...')\n",
    "with zipfile.ZipFile(ZIP_PATH, 'r') as zf:\n",
    "    zf.extractall(EXTRACT_DIR)\n",
    "print('Done!')\n",
    "\n",
    "# Show top-level structure\n",
    "for entry in sorted(os.listdir(EXTRACT_DIR)):\n",
    "    print(f'  {entry}')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 4B: Extract & convert FracAtlas (Option B) ───────────────\n",
    "# !! Skip this cell if using Roboflow (Option A) !!\n",
    "import os, json, shutil, zipfile\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "\n",
    "ZIP_PATH    = '/content/drive/MyDrive/FracAtlas.zip'\n",
    "EXTRACT_DIR = '/content/fracatlas_raw'\n",
    "YOLO_DIR    = '/content/dataset'\n",
    "\n",
    "# 1. Extract\n",
    "print('Extracting...')\n",
    "with zipfile.ZipFile(ZIP_PATH, 'r') as zf:\n",
    "    zf.extractall(EXTRACT_DIR)\n",
    "\n",
    "# 2. Locate images and COCO annotation files\n",
    "def find_file(base, name):\n",
    "    for p in Path(base).rglob(name):\n",
    "        return str(p)\n",
    "    return None\n",
    "\n",
    "images_root  = find_file(EXTRACT_DIR, 'images') or EXTRACT_DIR\n",
    "train_json   = find_file(EXTRACT_DIR, 'train.json')\n",
    "test_json    = find_file(EXTRACT_DIR, 'test.json')  or find_file(EXTRACT_DIR, 'val.json')\n",
    "\n",
    "print(f'Images root : {images_root}')\n",
    "print(f'Train JSON  : {train_json}')\n",
    "print(f'Val JSON    : {test_json}')\n",
    "\n",
    "# 3. COCO → YOLO converter\n",
    "def coco_to_yolo(coco_json_path, src_images_dir, out_images_dir, out_labels_dir):\n",
    "    os.makedirs(out_images_dir, exist_ok=True)\n",
    "    os.makedirs(out_labels_dir, exist_ok=True)\n",
    "\n",
    "    with open(coco_json_path) as f:\n",
    "        coco = json.load(f)\n",
    "\n",
    "    id2info = {img['id']: img for img in coco['images']}\n",
    "    # group annotations by image_id\n",
    "    img_anns = {}\n",
    "    for ann in coco.get('annotations', []):\n",
    "        img_anns.setdefault(ann['image_id'], []).append(ann)\n",
    "\n",
    "    for img_id, img_info in id2info.items():\n",
    "        fname   = img_info['file_name']\n",
    "        W, H    = img_info['width'], img_info['height']\n",
    "        src_img = os.path.join(src_images_dir, os.path.basename(fname))\n",
    "\n",
    "        if not os.path.exists(src_img):\n",
    "            # search recursively\n",
    "            found = find_file(EXTRACT_DIR, os.path.basename(fname))\n",
    "            if found:\n",
    "                src_img = found\n",
    "            else:\n",
    "                continue\n",
    "\n",
    "        stem = Path(fname).stem\n",
    "        shutil.copy(src_img, os.path.join(out_images_dir, os.path.basename(fname)))\n",
    "\n",
    "        label_lines = []\n",
    "        for ann in img_anns.get(img_id, []):\n",
    "            cat_id = ann['category_id'] - 1   # 0-indexed\n",
    "            x, y, w, h = ann['bbox']          # COCO: top-left x,y + w,h in pixels\n",
    "            xc = (x + w / 2) / W\n",
    "            yc = (y + h / 2) / H\n",
    "            wn = w / W\n",
    "            hn = h / H\n",
    "            label_lines.append(f'{cat_id} {xc:.6f} {yc:.6f} {wn:.6f} {hn:.6f}')\n",
    "\n",
    "        label_path = os.path.join(out_labels_dir, stem + '.txt')\n",
    "        with open(label_path, 'w') as lf:\n",
    "            lf.write('\\n'.join(label_lines))\n",
    "\n",
    "    print(f'  Converted {len(id2info)} images → {out_images_dir}')\n",
    "\n",
    "print('\\nConverting train split...')\n",
    "coco_to_yolo(train_json,\n",
    "             images_root,\n",
    "             f'{YOLO_DIR}/images/train',\n",
    "             f'{YOLO_DIR}/labels/train')\n",
    "\n",
    "print('Converting val split...')\n",
    "coco_to_yolo(test_json,\n",
    "             images_root,\n",
    "             f'{YOLO_DIR}/images/valid',\n",
    "             f'{YOLO_DIR}/labels/valid')\n",
    "\n",
    "print('\\nFracAtlas → YOLO conversion complete!')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 5: Build dataset.yaml ────────────────────────────────────\n",
    "import os, yaml\n",
    "\n",
    "DATASET_DIR = '/content/dataset'\n",
    "\n",
    "# Find images root\n",
    "def find_images_root(base):\n",
    "    for root, dirs, _ in os.walk(base):\n",
    "        if 'images' in dirs:\n",
    "            return root\n",
    "    return base\n",
    "\n",
    "dataset_root = find_images_root(DATASET_DIR)\n",
    "images_dir   = os.path.join(dataset_root, 'images')\n",
    "val_name     = 'valid' if os.path.exists(os.path.join(images_dir, 'valid')) else 'val'\n",
    "has_train    = os.path.exists(os.path.join(images_dir, 'train'))\n",
    "has_val      = os.path.exists(os.path.join(images_dir, val_name))\n",
    "\n",
    "# Count images\n",
    "def count_imgs(path):\n",
    "    if not os.path.exists(path): return 0\n",
    "    return sum(1 for f in os.listdir(path) if f.lower().endswith(('.jpg','.jpeg','.png')))\n",
    "\n",
    "print(f'Dataset root : {dataset_root}')\n",
    "print(f'Train images : {count_imgs(os.path.join(images_dir, \"train\"))}')\n",
    "print(f'Val images   : {count_imgs(os.path.join(images_dir, val_name))}')\n",
    "\n",
    "# Auto-detect class names from data.yaml if present (Roboflow includes one)\n",
    "roboflow_yaml = os.path.join(dataset_root, 'data.yaml')\n",
    "if os.path.exists(roboflow_yaml):\n",
    "    with open(roboflow_yaml) as f:\n",
    "        rf = yaml.safe_load(f)\n",
    "    CLASSES = rf.get('names', ['fracture'])\n",
    "    print(f'Classes from Roboflow yaml: {CLASSES}')\n",
    "else:\n",
    "    # FracAtlas classes\n",
    "    CLASSES = ['Elbow Positive', 'Fingers Positive', 'Forearm Fracture',\n",
    "               'Humerus Fracture', 'Shoulder Fracture', 'Wrist Positive']\n",
    "    print(f'Using FracAtlas classes: {CLASSES}')\n",
    "\n",
    "cfg = {\n",
    "    'path'  : dataset_root,\n",
    "    'train' : 'images/train' if has_train else 'images',\n",
    "    'val'   : f'images/{val_name}' if has_val else 'images/train',\n",
    "    'nc'    : len(CLASSES),\n",
    "    'names' : CLASSES,\n",
    "}\n",
    "\n",
    "YAML_PATH = '/content/fracture_dataset.yaml'\n",
    "with open(YAML_PATH, 'w') as f:\n",
    "    yaml.dump(cfg, f, default_flow_style=False)\n",
    "\n",
    "print(f'\\ndataset.yaml saved to {YAML_PATH}')\n",
    "!cat /content/fracture_dataset.yaml"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 6: TRAIN ─────────────────────────────────────────────────\n",
    "# T4 GPU: ~1 hour for 50 epochs, ~2 hours for 100 epochs\n",
    "from ultralytics import YOLO\n",
    "\n",
    "model = YOLO('yolov8n.pt')   # swap to yolov8s.pt for +2% mAP, +30min\n",
    "\n",
    "results = model.train(\n",
    "    data      = '/content/fracture_dataset.yaml',\n",
    "    epochs    = 50,\n",
    "    batch     = 32,\n",
    "    imgsz     = 640,\n",
    "    project   = '/content/drive/MyDrive/xrayvision_training',\n",
    "    name      = 'fracture_v1',\n",
    "    device    = 0,\n",
    "    # Medical X-ray augmentations\n",
    "    hsv_h     = 0.0,    # X-rays are grayscale\n",
    "    hsv_s     = 0.0,\n",
    "    hsv_v     = 0.4,    # brightness variation (exposure)\n",
    "    degrees   = 10.0,   # slight rotation\n",
    "    fliplr    = 0.5,\n",
    "    mosaic    = 0.5,\n",
    "    patience  = 15,     # early stopping\n",
    "    save_period = 10,\n",
    "    plots     = True,\n",
    "    verbose   = True,\n",
    "    exist_ok  = True,\n",
    ")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 7: Validate ──────────────────────────────────────────────\n",
    "from ultralytics import YOLO\n",
    "\n",
    "BEST = '/content/drive/MyDrive/xrayvision_training/fracture_v1/weights/best.pt'\n",
    "model   = YOLO(BEST)\n",
    "metrics = model.val(data='/content/fracture_dataset.yaml')\n",
    "\n",
    "print('\\n========= FINAL RESULTS =========')\n",
    "print(f'mAP50     : {metrics.box.map50:.4f}')\n",
    "print(f'mAP50-95  : {metrics.box.map:.4f}')\n",
    "print(f'Precision : {metrics.box.mp:.4f}')\n",
    "print(f'Recall    : {metrics.box.mr:.4f}')\n",
    "print('\\nPer-class:')\n",
    "for i, name in enumerate(metrics.names.values()):\n",
    "    ap = metrics.box.ap50[i] if i < len(metrics.box.ap50) else 0\n",
    "    print(f'  {name:<30} AP50={ap:.4f}')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ── Cell 8: Download best.pt ──────────────────────────────────────\n",
    "BEST = '/content/drive/MyDrive/xrayvision_training/fracture_v1/weights/best.pt'\n",
    "\n",
    "from google.colab import files\n",
    "files.download(BEST)\n",
    "\n",
    "print()\n",
    "print('After downloading:')\n",
    "print('  1. Copy best.pt  →  backend/models/fracture_yolov8.pt')\n",
    "print('  2. Edit backend/.env:')\n",
    "print('       YOLO_WEIGHTS_PATH=models/fracture_yolov8.pt')\n",
    "print('       ALLOW_GENERIC_YOLO_WEIGHTS=false')\n",
    "print('  3. Restart the FastAPI server')"
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}