{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Khmer & Math Detection - Test YOLOv8 OBB Model\n",
    "\n",
    "Upload a document image to detect **Khmer text** (red boxes) and **LaTeX math expressions** (blue boxes).\n",
    "\n",
    "Model: `krotreaksmey/Yolo_Detect_khmer_math` | mAP50: 0.995"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ultralytics huggingface_hub Pillow -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO\n",
    "from huggingface_hub import hf_hub_download\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "from IPython.display import display\n",
    "import os\n",
    "\n",
    "# Download model\n",
    "path = hf_hub_download(\n",
    "    repo_id=\"krotreaksmey/Yolo_Detect_khmer_math\",\n",
    "    filename=\"best_yolo_khmer_math.pt\"\n",
    ")\n",
    "model = YOLO(path)\n",
    "print(f\"Model loaded: {path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CLASSES = {0: \"khmer\", 1: \"expression\"}\n",
    "COLORS = {0: (255, 80, 80), 1: (80, 180, 255)}\n",
    "\n",
    "def detect_and_show(image_path, conf=0.3):\n",
    "    img = Image.open(image_path).convert(\"RGB\")\n",
    "    results = model.predict(img, conf=conf, imgsz=1280)\n",
    "    \n",
    "    draw = ImageDraw.Draw(img)\n",
    "    try:\n",
    "        font = ImageFont.truetype(\"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf\", 16)\n",
    "    except:\n",
    "        font = ImageFont.load_default()\n",
    "    \n",
    "    count = 0\n",
    "    for result in results:\n",
    "        if result.obb is not None:\n",
    "            for i in range(len(result.obb.cls)):\n",
    "                cls_id = int(result.obb.cls[i])\n",
    "                confidence = float(result.obb.conf[i])\n",
    "                name = CLASSES.get(cls_id, f\"class_{cls_id}\")\n",
    "                color = COLORS.get(cls_id, (255, 255, 0))\n",
    "                points = result.obb.xyxyxyxy[i][0].tolist()\n",
    "                pts = [(p[0], p[1]) for p in points]\n",
    "                draw.polygon(pts, outline=color, width=3)\n",
    "                label = f\"{name} {confidence:.0%}\"\n",
    "                tx, ty = pts[0]\n",
    "                draw.text((tx, ty - 20), label, fill=color, font=font)\n",
    "                count += 1\n",
    "    \n",
    "    display(img)\n",
    "    print(f\"\\n{count} detections found.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload your own image\n",
    "from google.colab import files\n",
    "uploaded = files.upload()\n",
    "\n",
    "for filename in uploaded:\n",
    "    print(f\"\\nProcessing: {filename}\")\n",
    "    detect_and_show(filename, conf=0.3)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}