{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Image Classification with DINOv3\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for image classification with state-of-the-art DINOv3 backbone.\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/image_classification.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": "1",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "LightlyTrain can be installed directly via `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "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](https://docs.lightly.ai/train/stable/installation.html) for more details on installation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Train image classification model\n",
    "\n",
    "Training your own classification model is straightforward with LightlyTrain.\n",
    "\n",
    "### Download dataset\n",
    "\n",
    "First download a dataset. The dataset must be organized in class folders, see the [documentation](https://docs.lightly.ai/train/stable/image_classification.html#data) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/alexeygrigorev/clothing-dataset-small"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The dataset looks like this after the download completes:\n",
    "```\n",
    "clothing-dataset-small\n",
    "├── train\n",
    "│   ├── dress\n",
    "│   │   ├── 009b3c31-fb62-45c0-be9a-37a5c238cb88.jpg\n",
    "│   │   ├── 041c6bde-e737-46fd-9586-984c1503941f.jpg\n",
    "│   │   └── ...\n",
    "│   ├── hat\n",
    "│   │   ├── 00d94e21-5891-492e-be0e-792e7338c077.jpg\n",
    "│   │   └── ...\n",
    "│   └── ...\n",
    "├── validation\n",
    "│   ├── dress\n",
    "│   │   ├── 07cddef1-1fc8-47e4-a28a-613e60912590.jpg\n",
    "│   │   └── ...\n",
    "│   └── ...\n",
    "└── test\n",
    "    ├── dress\n",
    "    │   ├── 06a00c0f-5f9a-410d-a7da-3881a9df3a71.jpg\n",
    "    │   └── ...\n",
    "    └── ...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Next, start the training with the `train_image_classification` function. You only have to specify the output directory, model, and input data. LightlyTrain automatically sets the remaining training parameters and applies image augmentations. Of course you can always customize these settings if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "lightly_train.train_image_classification(\n",
    "    out=\"out/my_experiment\",\n",
    "    model=\"dinov3/vitt16\",\n",
    "    steps=1000,  # Small number of steps for demonstration, default is 100_000.\n",
    "    batch_size=16,\n",
    "    data={\n",
    "        \"train\": \"clothing-dataset-small/train\",\n",
    "        \"val\": \"clothing-dataset-small/validation\",\n",
    "        \"classes\": {\n",
    "            0: \"dress\",\n",
    "            1: \"hat\",\n",
    "            2: \"longsleeve\",\n",
    "            3: \"outwear\",\n",
    "            4: \"pants\",\n",
    "            5: \"shirt\",\n",
    "            6: \"shoes\",\n",
    "            7: \"shorts\",\n",
    "            8: \"skirt\",\n",
    "            9: \"t-shirt\",\n",
    "        },\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Once the training is complete, the output directory looks like this:\n",
    "```\n",
    "out/my_experiment\n",
    "├── checkpoints\n",
    "│   ├── best.ckpt\n",
    "│   └── last.ckpt\n",
    "├── events.out.tfevents.1764251158.ef9b159fe4b8.273.0\n",
    "├── exported_models\n",
    "│   ├── exported_best.pt\n",
    "│   └── exported_last.pt\n",
    "└── train.log\n",
    "```\n",
    "\n",
    "The best model checkpoint is saved to `out/my_experiment/exported_models/exported_best.pt`. You can load it for inference like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model for inference\n",
    "model = lightly_train.load_model(\"out/my_experiment/exported_models/exported_best.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision.io import read_image\n",
    "\n",
    "# Run inference\n",
    "image_path = (\n",
    "    \"clothing-dataset-small/validation/dress/07cddef1-1fc8-47e4-a28a-613e60912590.jpg\"\n",
    ")\n",
    "results = model.predict(image_path, topk=1)\n",
    "\n",
    "image = read_image(image_path)\n",
    "plt.imshow(image.permute(1, 2, 0))\n",
    "plt.title(\n",
    "    f\"Predicted class: {model.classes[results['labels'][0].item()]}\\nConfidence: {results['scores'][0]:.1%}\"\n",
    ")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "* [Image Classification Documentation](https://docs.lightly.ai/train/stable/image_classification.html): If you want to learn more about image classification with LightlyTrain.\n",
    "* [Quick Start Guide](https://docs.lightly.ai/train/stable/quick_start_object_detection.html): If you want to learn more about LightlyTrain in general.\n",
    "* [DINOv2 Pretraining](https://docs.lightly.ai/train/stable/pretrain_distill/methods/dinov2.html): If you want to learn how to pretrain foundation models on your data."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}