YOLOv5_tomato_roboflow_Classification_Tutorial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/danielsoy/ALOCC-CVPR2018/blob/master/YOLOv5_tomato_roboflow_Classification_Tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# YOLOv5 Classification Tutorial\n",
        "\n",
        "YOLOv5 supports classification tasks too. This is the official YOLOv5 classification notebook tutorial. YOLOv5 is maintained by [Ultralytics](https://github.com/ultralytics/yolov5).\n",
        "\n",
        "This notebook covers:\n",
        "\n",
        "*   Inference with out-of-the-box YOLOv5 classification on ImageNet\n",
        "*  [Training YOLOv5 classification](https://blog.roboflow.com//train-YOLOv5-classification-custom-data) on custom data\n",
        "\n",
        "*Looking for custom data? Explore over 66M community datasets on [Roboflow Universe](https://universe.roboflow.com).*\n",
        "\n",
        "This notebook was created with Google Colab. [Click here](https://colab.research.google.com/drive/1FiSNz9f_nT8aFtDEU3iDAQKlPT8SCVni?usp=sharing) to run it."
      ],
      "metadata": {
        "id": "5GYQX3of4QiW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n",
        "\n",
        "Pull in respective libraries to prepare the notebook environment."
      ],
      "metadata": {
        "id": "-PJ8vlYXbWtN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pIM7fOwm8A7l"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/ultralytics/yolov5  # clone\n",
        "%cd yolov5\n",
        "%pip install -qr requirements.txt  # install\n",
        "\n",
        "import torch\n",
        "import utils\n",
        "display = utils.notebook_init()  # checks"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. Infer on ImageNet\n",
        "\n",
        "To demonstrate YOLOv5 classification, we'll leverage an already trained model. In this case, we'll download the ImageNet trained models pretrained on ImageNet using YOLOv5 Utils."
      ],
      "metadata": {
        "id": "i_DrUi2nmF40"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from utils.downloads import attempt_download\n",
        "\n",
        "p5 = ['n', 's', 'm', 'l', 'x']  # P5 models\n",
        "cls = [f'{x}-cls' for x in p5]  # classification models\n",
        "\n",
        "for x in cls:\n",
        "    attempt_download(f'weights/yolov5{x}.pt')"
      ],
      "metadata": {
        "id": "o2scLEh6EYnL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, we can infer on an example image from the ImageNet dataset."
      ],
      "metadata": {
        "id": "Fn2_a38DmZ2H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Download example image\n",
        "import requests\n",
        "image_url = \"https://i.imgur.com/OczPfaz.jpg\"\n",
        "img_data = requests.get(image_url).content\n",
        "with open('bananas.jpg', 'wb') as handler:\n",
        "    handler.write(img_data)"
      ],
      "metadata": {
        "id": "L9objhVHnS-h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#Infer using classify/predict.py\n",
        "!python classify/predict.py --weights ./weigths/yolov5s-cls.pt --source bananas.jpg"
      ],
      "metadata": {
        "id": "qqxF5pHCrLd3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "From the output, we can see the ImageNet trained model correctly predicts the class `banana` with `0.95` confidence."
      ],
      "metadata": {
        "id": "yQmj7IXqo3kk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. (Optional) Validate\n",
        "\n",
        "Use the `classify/val.py` script to run validation for the model. This will show us the model's performance on each class.\n",
        "\n",
        "First, we need to download ImageNet."
      ],
      "metadata": {
        "id": "5EosQzyDCk3W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # WARNING: takes ~20 minutes\n",
        "# !bash data/scripts/get_imagenet.sh --val"
      ],
      "metadata": {
        "id": "HwAYptjCq-C_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # run the validation script\n",
        "# !python classify/val.py --weights ./weigths/yolov5s-cls.pt --data ../datasets/imagenet"
      ],
      "metadata": {
        "id": "CoHdKXWc8hrD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The output shows accuracy metrics for the ImageNet validation dataset including per class accuracy."
      ],
      "metadata": {
        "id": "r2coOcIjuzCO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3. Train On Custom Data\n",
        "\n",
        "To train on custom data, we need to prepare a dataset with custom labels.\n",
        "\n",
        "To prepare custom data, we'll use [Roboflow](https://roboflow.com). Roboflow enables easy dataset prep with your team, including labeling, formatting into the right export format, deploying, and active learning with a `pip` package. \n",
        "\n",
        "If you need custom data, there are over 66M open source images from the community on [Roboflow Universe](https://universe.roboflow.com).\n",
        "\n",
        "(For more guidance, here's a detailed blog on [training YOLOv5 classification on custom data](https://blog.roboflow.com/train-YOLOv5-classification-custom-data).)\n",
        "\n",
        "\n",
        "Create a free Roboflow account, upload your data, and label. \n",
        "\n",
        "![](https://s4.gifyu.com/images/fruit-labeling.gif)"
      ],
      "metadata": {
        "id": "9bXHHYeVDCXg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load Custom Dataset\n",
        "\n",
        "Next, we'll export our dataset into the right directory structure for training YOLOv5 classification to load into this notebook. Select the `Export` button at the top of the version page, `Folder Structure` type, and `show download code`.\n",
        "\n",
        "The ensures all our directories are in the right format:\n",
        "\n",
        "```\n",
        "dataset\n",
        "├── train\n",
        "│   ├── class-one\n",
        "│   │   ├── IMG_123.jpg\n",
        "│   └── class-two\n",
        "│       ├── IMG_456.jpg\n",
        "├── valid\n",
        "│   ├── class-one\n",
        "│   │   ├── IMG_789.jpg\n",
        "│   └── class-two\n",
        "│       ├── IMG_101.jpg\n",
        "├── test\n",
        "│   ├── class-one\n",
        "│   │   ├── IMG_121.jpg\n",
        "│   └── class-two\n",
        "│       ├── IMG_341.jpg\n",
        "```\n",
        "\n",
        "![](https://i.imgur.com/BF9BNR8.gif)\n",
        "\n",
        "\n",
        "Copy and paste that snippet into the cell below."
      ],
      "metadata": {
        "id": "Cu6-lrukD6Hc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Ensure we're in the right directory to download our custom dataset\n",
        "import os\n",
        "os.makedirs(\"../datasets/\", exist_ok=True)\n",
        "%cd ../datasets/"
      ],
      "metadata": {
        "id": "6IIgJbP7G6Th"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# REPLACE the below with your exported code snippet from above\n",
        "!pip install roboflow\n",
        "from roboflow import Roboflow\n",
        "\n",
        "#rf = Roboflow(api_key=\"YOUR API KEY\")\n",
        "#project = rf.workspace(\"yolov5-classification\").project(\"banana-ripeness-classification\")\n",
        "#dataset = project.version(1).download(\"folder\")"
      ],
      "metadata": {
        "id": "He6JwHIlG-W_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rf = Roboflow(api_key=\"8PQLqPZJ3S8IN0OVaY8V\")\n",
        "project = rf.workspace().project(\"tomato-classification-jc5kz\")\n",
        "model = project.version(2).model"
      ],
      "metadata": {
        "id": "ZU_xnHocuuj1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#Save the dataset name to the environment so we can use it in a system call later\n",
        "dataset_name = dataset.location.split(os.sep)[-1]\n",
        "os.environ[\"DATASET_NAME\"] = dataset_name"
      ],
      "metadata": {
        "id": "wLQbThFICpn4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train On Custom Data 🎉\n",
        "Here, we use the DATASET_NAME environment variable to pass our dataset to the `--data` parameter.\n",
        "\n",
        "Note: we're training for 100 epochs here. We're also starting training from the pretrained weights. Larger datasets will likely benefit from longer training. "
      ],
      "metadata": {
        "id": "-5z7Yv42FGrK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd ../yolov5\n",
        "!python classify/train.py --model yolov5s-cls.pt --data $DATASET_NAME --epochs 100 --img 128 --pretrained weights/yolov5s-cls.pt"
      ],
      "metadata": {
        "id": "MXWTTN2BEaqe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Validate Your Custom Model\n",
        "\n",
        "Repeat step 2 from above to test and validate your custom model."
      ],
      "metadata": {
        "id": "HHUFGeLbGd98"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python classify/val.py --weights runs/train-cls/exp/weights/best.pt --data ../datasets/$DATASET_NAME"
      ],
      "metadata": {
        "id": "DIV7ydyKGZFL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Infer With Your Custom Model"
      ],
      "metadata": {
        "id": "uH5tJNpEsi6g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Get the path of an image from the test or validation set\n",
        "if os.path.exists(os.path.join(dataset.location, \"test\")):\n",
        "  split_path = os.path.join(dataset.location, \"test\")\n",
        "else:\n",
        "  os.path.join(dataset.location, \"valid\")\n",
        "example_class = os.listdir(split_path)[0]\n",
        "example_image_name = os.listdir(os.path.join(split_path, example_class))[0]\n",
        "example_image_path = os.path.join(split_path, example_class, example_image_name)\n",
        "os.environ[\"TEST_IMAGE_PATH\"] = example_image_path\n",
        "\n",
        "print(f\"Inferring on an example of the class '{example_class}'\")\n",
        "\n",
        "#Infer\n",
        "!python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source $TEST_IMAGE_PATH"
      ],
      "metadata": {
        "id": "81lK1hU_sk54"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can see the inference results show ~3ms inference and the respective classes predicted probabilities."
      ],
      "metadata": {
        "id": "DdGuG-1kNjWT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## (OPTIONAL) Improve Our Model with Active Learning\n",
        "\n",
        "Now that we've trained our model once, we will want to continue to improve its performance. Improvement is largely dependent on improving our dataset.\n",
        "\n",
        "We can programmatically upload example failure images back to our custom dataset based on conditions (like seeing an underrpresented class or a low confidence score) using the same `pip` package."
      ],
      "metadata": {
        "id": "I38IM6NXKNN9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # Upload example image\n",
        "# project.upload(image_path)\n"
      ],
      "metadata": {
        "id": "HycgSEnYKo0J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # Example upload code \n",
        "# min_conf = float(\"inf\")\n",
        "# for pred in results:\n",
        "#     if pred[\"score\"] < min_conf:\n",
        "#         min_conf = pred[\"score\"]\n",
        "# if min_conf < 0.4:\n",
        "#     project.upload(image_path)"
      ],
      "metadata": {
        "id": "VwXDoz_vLK3V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# (BONUS) YOLOv5 classify/predict.py Accepts Several Input Methods\n",
        "- Webcam: `python classify/predict.py --weights yolov5s-cls.pt --source 0`\n",
        "- Image `python classify/predict.py --weights yolov5s-cls.pt --source img.jpg`\n",
        "- Video: `python classify/predict.py --weights yolov5s-cls.pt --source vid.mp4`\n",
        "- Directory: `python classify/predict.py --weights yolov5s-cls.pt --source path/`\n",
        "- Glob: `python classify/predict.py --weights yolov5s-cls.pt --source 'path/*.jpg'`\n",
        "- YouTube: `python classify/predict.py --weights yolov5s-cls.pt --source 'https://youtu.be/Zgi9g1ksQHc'`\n",
        "- RTSP, RTMP, HTTP stream: `python classify/predict.py --weights yolov5s-cls.pt --source 'rtsp://example.com/media.mp4'`"
      ],
      "metadata": {
        "id": "aYlfaHDusN-j"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Directory Example"
      ],
      "metadata": {
        "id": "iKSP-SNTvcLJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Directory infer\n",
        "os.environ[\"TEST_CLASS_PATH\"] = test_class_path = os.path.join(*os.environ[\"TEST_IMAGE_PATH\"].split(os.sep)[:-1])\n",
        "print(f\"Infering on all images from the directory {os.environ['TEST_CLASS_PATH']}\")\n",
        "!python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source /$TEST_CLASS_PATH/"
      ],
      "metadata": {
        "id": "lwSoHcHcvjeD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "###YouTube Example"
      ],
      "metadata": {
        "id": "kCCao9t8se8i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#YouTube infer\n",
        "!python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source 'https://www.youtube.com/watch?v=7AlYA4ItA74'"
      ],
      "metadata": {
        "id": "heebjpJBsakV"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}