{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# LightlyTrain with RF-DETR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "In this notebook we will demonstrate how you can use [LightlyTrain](https://docs.lightly.ai/train/stable/index.html) to pretrain an [RF-DETR model](https://github.com/roboflow/rf-detr) from Roboflow. For now, `rfdetr` only supports training with datasets in COCO JSON format. To this end, for pretraining we use the raw images (**no labels**) from [the COCO-minitrain dataset](https://github.com/giddyyupp/coco-minitrain), a subset of the COCO dataset with 25k images, and for fine-tuning we use the Roboflow's [Coconut Custom Dataset](https://universe.roboflow.com/ravi-mgvlz/coconut-custom-dataset).\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/rfdetr.ipynb)\n",
    "\n",
    "> **Important**: When running on Google Colab make sure to select a GPU runtime for faster processing. You can do this by going to `Runtime` > `Change runtime type` and selecting a GPU hardware accelerator."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "You can install `lightly_train` directly from PyPI using pip with support of `rfdetr` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"lightly-train[rfdetr]\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "> **Important**: LightlyTrain is officially supported on\n",
    "> - Linux: CPU or CUDA\n",
    "> - MacOS: CPU only\n",
    "> - Windows (experimental): CPU or CUDA\n",
    ">\n",
    "> We are planning to support MPS for MacOS.\n",
    ">\n",
    "> Check the installation instructions for more details on installation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Pretrain on COCO-minitrain-25k Dataset\n",
    "\n",
    "We use [the COCO-minitrain dataset](https://github.com/giddyyupp/coco-minitrain), a subset of the COCO dataset with 25k images for pretraining the RF-DETR model.\n",
    "\n",
    "### Download the Dataset\n",
    "\n",
    "We can download the COCO-minitrain dataset (25k images) directly from HuggingFace\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://huggingface.co/datasets/bryanbocao/coco_minitrain/resolve/main/coco_minitrain_25k.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "... unzip it..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip coco_minitrain_25k.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "... and since LightlyTrain does not require any labels, we can also confidently delete all the labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf coco_minitrain_25k/labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Pretrain an RF-DETR Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Pretraining an RF-DETR model with LightlyTrain is straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    lightly_train.pretrain(\n",
    "        out=\"out/my_experiment\",  # Output directory.\n",
    "        data=\"coco_minitrain_25k/images\",  # Directory with images.\n",
    "        model=\"rfdetr/rf-detr-base\",  # Pass the RF-DETR model.\n",
    "        epochs=5,  # Number of epochs to train\n",
    "        batch_size=16,  # Batch size\n",
    "        overwrite=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Fine-tune on Coconuts Custom Dataset\n",
    "\n",
    "We use Roboflow's [Coconut Custom Dataset](https://universe.roboflow.com/ravi-mgvlz/coconut-custom-dataset) for fine-tuning.\n",
    "\n",
    "### Download the Dataset\n",
    "\n",
    "The dataset can be directly downloaded via Roboflow API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "from roboflow import Roboflow\n",
    "\n",
    "rf = Roboflow(api_key=\"your_roboflow_api_key\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "project = rf.workspace(\"ravi-mgvlz\").project(\"coconut-custom-dataset\")\n",
    "version = project.version(3)\n",
    "finetune_dataset = version.download(\"coco\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "### Fine-tune an RF-DETR Model\n",
    "\n",
    "You can directly use the `rfdetr` package for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rfdetr import RFDETRBase\n",
    "\n",
    "model = RFDETRBase(\n",
    "    pretrain_weights=\"out/my_experiment/exported_models/exported_last.pt\"\n",
    ")\n",
    "model.train(dataset_dir=finetune_dataset.location, epochs=10, batch_size=4, lr=1e-4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}