Tool_Wear_Detection.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyM5d0bxwhFFKSiHkdDkQiEr",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/khazzan/toolwearID/blob/main/Tool_Wear_Detection.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RufJptsgrQpe"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "!pip install -q tflite-model-maker\n",
        "!sudo apt -y install libportaudio2\n",
        "!pip install -q --use-deprecated=legacy-resolver tflite-model-maker\n",
        "!pip install -q pycocotools\n",
        "\n",
        "!pip install -q opencv-python-headless==4.1.2.30\n",
        "!pip uninstall -y tensorflow && pip install -q tensorflow==2.8.0\n",
        "!git clone https://github.com/khazzan/toolwearID.git\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import os\n",
        "\n",
        "from tflite_model_maker.config import ExportFormat\n",
        "from tflite_model_maker import model_spec\n",
        "from tflite_model_maker import object_detector\n",
        "\n",
        "import tensorflow as tf\n",
        "assert tf.__version__.startswith('2')\n",
        "\n",
        "tf.get_logger().setLevel('ERROR')\n",
        "from absl import logging\n",
        "logging.set_verbosity(logging.ERROR)"
      ],
      "metadata": {
        "id": "Rf8z1ckBrdk7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset_is_split = True\n",
        "use_custom_dataset = True\n",
        "if not use_custom_dataset:\n",
        "  train_data, validation_data, test_data = object_detector.DataLoader.from_csv('wear_ml_use.csv')\n",
        "label_map = {1: 'flank_wear'} \n",
        "\n",
        "train_images_dir = 'toolwearID/split-dataset/train/images'\n",
        "train_annotations_dir = 'toolwearID/split-dataset/train/annotations'\n",
        "val_images_dir = 'toolwearID/split-dataset/validation/images'\n",
        "val_annotations_dir = 'toolwearID/split-dataset/validation/annotations'\n",
        "test_images_dir = 'toolwearID/split-dataset/test/images'\n",
        "test_annotations_dir = 'toolwearID/split-dataset/test/annotations'\n",
        "\n",
        "train_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        train_images_dir, train_annotations_dir, label_map=label_map)\n",
        "validation_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        val_images_dir, val_annotations_dir, label_map=label_map)\n",
        "test_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        test_images_dir, test_annotations_dir, label_map=label_map)"
      ],
      "metadata": {
        "id": "l0xT0adNrheU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "images_path = test_images_dir if dataset_is_split else os.path.join(test_dir, \"images\")\n",
        "filenames = os.listdir(os.path.join(images_path))\n",
        "random_index = random.randint(0, len(filenames) -1)\n",
        "!python3 -m pip install --extra-index-url https://google-coral.github.io/py-repo/ pycoral"
      ],
      "metadata": {
        "id": "Tn2ChXpPrnlc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "from PIL import ImageFont\n",
        "\n",
        "!python3 -m pip install tflite-runtime\n",
        "!pip install tflite-support\n",
        "import tflite_runtime.interpreter as tflite \n",
        "\n",
        "\n",
        "from pycoral.adapters import common\n",
        "from pycoral.adapters import detect\n",
        "from pycoral.utils.dataset import read_label_file\n",
        "\n",
        "def draw_objects(draw, objs, scale_factor, labels):\n",
        "  \"\"\"Draws the bounding box and label for each object.\"\"\"\n",
        "  COLORS = np.random.randint(0, 255, size = (len(labels), 3), dtype = np.uint8)\n",
        "  for obj in objs:\n",
        "    bbox = obj.bbox\n",
        "    color = tuple(int(c) for c in COLORS[obj.id])\n",
        "    draw.rectangle([(bbox.xmin * scale_factor, bbox.ymin * scale_factor),\n",
        "                    (bbox.xmax * scale_factor, bbox.ymax * scale_factor)],\n",
        "                   outline = color, width = 3)\n",
        "    font = ImageFont.truetype(\"LiberationSans-Regular.ttf\", size = 15)\n",
        "    draw.text((bbox.xmin * scale_factor + 4, bbox.ymin * scale_factor + 4),\n",
        "              '%s\\n%.2f' % (labels.get(obj.id, obj.id), obj.score),\n",
        "              fill = color, font = font)\n",
        "\n",
        "# Load the TF Lite model\n",
        "\n",
        "LABELS_FILENAME = 'toolwearID/tf-lite_model/wear-labels.txt'\n",
        "TFLITE_FILENAME = 'toolwearID/tf-lite_model/efficientdet-lite-wear.tflite'\n",
        "\n",
        "labels = read_label_file(LABELS_FILENAME)\n",
        "interpreter = tflite.Interpreter(TFLITE_FILENAME)\n",
        "interpreter.allocate_tensors()"
      ],
      "metadata": {
        "id": "-AmNfnMZrtTH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "INPUT_IMAGE = os.path.join(images_path, filenames[10])\n",
        "image = Image.open(INPUT_IMAGE)\n",
        "_, scale = common.set_resized_input(\n",
        "    interpreter, image.size, lambda size: image.resize(size, Image.ANTIALIAS))\n",
        "\n",
        "# Run inference\n",
        "interpreter.invoke()\n",
        "objs = detect.get_objects(interpreter, score_threshold = 0.4, image_scale = scale)\n",
        "\n",
        "# Resize again to a reasonable size for display\n",
        "display_width = 500\n",
        "scale_factor = display_width / image.width\n",
        "height_ratio = image.height / image.width\n",
        "image = image.resize((display_width, int(display_width * height_ratio)))\n",
        "draw_objects(ImageDraw.Draw(image), objs, scale_factor, labels)"
      ],
      "metadata": {
        "id": "x7qWUDr2r4BO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image"
      ],
      "metadata": {
        "id": "ma8O8xyDr6kN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(objs)\n",
        "ratio = 0.00771 # px to mm\n",
        "wear_edge = \"{:.2f}\".format(((objs[0][2][2]) - (objs[0][2][0])) * ratio)\n",
        "vb_max = \"{:.2f}\".format(((objs[0][2][3]) - (objs[0][2][1])) * ratio)\n",
        "print(str(wear_edge) + \" mm\")\n",
        "print(str(vb_max) + \" mm\")"
      ],
      "metadata": {
        "id": "F0EAYzwhr9iK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KZizM3i7uONi"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}