{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FDfI95Sh1lW0"
      },
      "source": [
        "# Is it Huggable?\n",
        "*Classify objects as huggable or not.*\n",
        "\n",
        "This notebook has steps to make the model.\n",
        "\n",
        "Just want to play? Use directly on the [website](https://daspartho.github.io/is-it-huggable)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r-GyBdvhzfY2"
      },
      "source": [
        "### Install required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CQdd5Egc-FQV"
      },
      "outputs": [],
      "source": [
        "!pip install -Uqq fastai duckduckgo_search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vgvpU91p0ERn"
      },
      "source": [
        "### Import modules required"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "BD7-yF0l-Y4h"
      },
      "outputs": [],
      "source": [
        "from duckduckgo_search import ddg_images\n",
        "from fastcore.all import *\n",
        "from fastdownload import download_url\n",
        "from fastai.vision.all import *"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WKZC9jY_zOfx"
      },
      "source": [
        "### Use DuckDuckGo to search for images of examples of the two groups"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hqnqTAXWCAn6"
      },
      "outputs": [],
      "source": [
        "def search_images(term, max_images=50):\n",
        "    print(f\"Searching for '{term}'\")\n",
        "    return L(ddg_images(term, max_results=max_images)).itemgot('image')\n",
        "\n",
        "path = Path('huggable_or_not')\n",
        "\n",
        "# examples of both groups\n",
        "categories={\n",
        "   'huggable':['plushie', 'pillow' , 'ballon', 'dog', 'cat', 'bunny', 'snowman', 'bed', 'sofa', 'people', 'baby', 'cloud', 'dolphin', 'horse', 'cow', 'sheep'], \n",
        "   'not huggable':['chainsaw', 'sword', 'cactus', 'barbwire', 'bear', 'snake', 'lion', 'shark', 'fire','knive','fork', 'dinosaur', 'crocodile', 'spider', 'bees', 'porcupine']\n",
        "   }\n",
        "\n",
        "for category in categories:\n",
        "    dest = (path/category)\n",
        "    dest.mkdir(exist_ok=True, parents=True)\n",
        "    for example in categories[category]:\n",
        "        download_images(dest, urls=search_images(f'{example} photo'))\n",
        "        resize_images(path/category, max_size=400, dest=path/category)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bpsp4MTGxBWl"
      },
      "source": [
        "### Remove photos that didn't download correctly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FzuHMc_qD0UO"
      },
      "outputs": [],
      "source": [
        "failed = verify_images(get_image_files(path))\n",
        "failed.map(Path.unlink)\n",
        "len(failed)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eFFr_VE45ihe"
      },
      "source": [
        "### Preparing the data for training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wzSeghRAFYqF"
      },
      "outputs": [],
      "source": [
        "dls = DataBlock(\n",
        "    blocks=(ImageBlock, CategoryBlock), # inputs to our model are images, and the outputs are categories\n",
        "    get_items=get_image_files, \n",
        "    splitter=RandomSplitter(valid_pct=0.2, seed=42), # Split the data into training and validation sets randomly, using 20% of the data for the validation set\n",
        "    get_y=parent_label, # The labels is the name of the parent of each file\n",
        "    item_tfms=RandomResizedCrop(224, min_scale=0.3), # picks a random scaled crop of an image and resize it to 224x224 pixels\n",
        "    batch_tfms=aug_transforms() # applies augmentations to an entire batch\n",
        ").dataloaders(path, bs=32)\n",
        "\n",
        "dls.show_batch()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y-FNWY-3zEF3"
      },
      "source": [
        "### Fine-tune a pretrained neural network to recognise these two groups"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ao0lw2cG2WP"
      },
      "outputs": [],
      "source": [
        "learn = vision_learner(dls, resnet34, metrics=error_rate)\n",
        "learn.fine_tune(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0wZAFpxi7L6z"
      },
      "source": [
        "### Show predictions the model made on images in validation set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aE0vp3jeVtBT"
      },
      "outputs": [],
      "source": [
        "learn.show_results()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gFpqdZr87ZSS"
      },
      "source": [
        "### Download an image from internet for trying the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C0kfX6QUMRoN"
      },
      "outputs": [],
      "source": [
        "term='penguin' # change the search term\n",
        "download_url(search_images(term, max_images=1)[0], 'test.jpg', show_progress=False)\n",
        "Image.open('test.jpg').to_thumb(256,256)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AgOQPzTX7q3o"
      },
      "source": [
        "### Trying the model on the downloaded image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sz1dVCZMHz3N"
      },
      "outputs": [],
      "source": [
        "predict,n,prob = learn.predict(PILImage.create('test.jpg'))\n",
        "print(f\"It's {predict}!\")\n",
        "perc = prob[n]*100\n",
        "print(f\"I'm {perc:.02f}% confident.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lSSjWJq874WE"
      },
      "source": [
        "### Export the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {
        "id": "ae2bc6ac"
      },
      "outputs": [],
      "source": [
        "learn.export('model.pkl')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "model.ipynb",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}