1_SSD.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SSD.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "SE_Si1tBwicy"
      ],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyPU4jGYq4iJZZtFFtqedC/S",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/gustavsma/MD/blob/main/SSD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initial setup"
      ],
      "metadata": {
        "id": "V88LWs_eut53"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from google.colab import drive\n",
        "\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "7-AthHvju1o5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "CUSTOM_MODEL_NAME = 'deep_weeds' \n",
        "PRETRAINED_MODEL_NAME = 'ssd_resnet50_v1_fpn_1024x1024_coco17_tpu-8'\n",
        "PRETRAINED_MODEL_URL = 'http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_resnet50_v1_fpn_1024x1024_coco17_tpu-8.tar.gz'\n",
        "TF_RECORD_SCRIPT_NAME = 'generate_tfrecord.py'\n",
        "LABEL_MAP_NAME = 'label_map.pbtxt'"
      ],
      "metadata": {
        "id": "Q_97mUuCvW-n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "paths = {\n",
        "    'WORKSPACE_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace'),\n",
        "    'SCRIPTS_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow','scripts'),\n",
        "    'APIMODEL_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow','models'),\n",
        "    'ANNOTATION_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','annotations'),\n",
        "    'IMAGE_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','images_old'),\n",
        "    'MODEL_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','models'),\n",
        "    'PRETRAINED_MODEL_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','pre-trained-models'),\n",
        "    'CHECKPOINT_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME), \n",
        "    'OUTPUT_PATH': os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'export'), \n",
        "    'TFJS_PATH':os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'tfjsexport'), \n",
        "    'TFLITE_PATH':os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'tfliteexport'), \n",
        "    'PROTOC_PATH':os.path.join('gdrive', 'MyDrive', 'Tensorflow','protoc')\n",
        " }\n",
        "\n",
        "files = {\n",
        "    'PIPELINE_CONFIG':os.path.join('gdrive', 'MyDrive', 'Tensorflow', 'workspace','models', CUSTOM_MODEL_NAME, 'pipeline.config'),\n",
        "    'TF_RECORD_SCRIPT': os.path.join(paths['SCRIPTS_PATH'], TF_RECORD_SCRIPT_NAME), \n",
        "    'LABELMAP': os.path.join(paths['ANNOTATION_PATH'], LABEL_MAP_NAME)\n",
        "}\n",
        "\n",
        "for path in paths.values():\n",
        "    if not os.path.exists(path):\n",
        "      !mkdir -p {path}"
      ],
      "metadata": {
        "id": "tH0ZO_O9vS7w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%bash\n",
        "cd gdrive/MyDrive/Tensorflow/models/research\n",
        "cp object_detection/packages/tf2/setup.py .\n",
        "python -m pip install --use-feature=2020-resolver ."
      ],
      "metadata": {
        "id": "ZYqvTgeJoD0n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import object_detection"
      ],
      "metadata": {
        "id": "mKfbbfOcm2E7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# !wget {PRETRAINED_MODEL_URL}\n",
        "# !mv {PRETRAINED_MODEL_NAME+'.tar.gz'} {paths['PRETRAINED_MODEL_PATH']}\n",
        "# !cd {paths['PRETRAINED_MODEL_PATH']} && tar -zxvf {PRETRAINED_MODEL_NAME+'.tar.gz'}"
      ],
      "metadata": {
        "id": "hL9BE11YwF5p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip uninstall opencv-python-headless==4.5.5.62 \n",
        "!pip install opencv-python-headless==4.1.2.30"
      ],
      "metadata": {
        "id": "UmfmdZgI02wd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "VERIFICATION_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'builders', 'model_builder_tf2_test.py')\n",
        "!python {VERIFICATION_SCRIPT}"
      ],
      "metadata": {
        "id": "jkl23486wL5s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prepare annotations"
      ],
      "metadata": {
        "id": "vfioqnLNv6mi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "labels = [{'name':'weed', 'id':1}]\n",
        "\n",
        "with open(files['LABELMAP'], 'w') as f:\n",
        "    for label in labels:\n",
        "        f.write('item { \\n')\n",
        "        f.write('\\tname:\\'{}\\'\\n'.format(label['name']))\n",
        "        f.write('\\tid:{}\\n'.format(label['id']))\n",
        "        f.write('}\\n')"
      ],
      "metadata": {
        "id": "5HB6JB_gwTUh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if not os.path.exists(files['TF_RECORD_SCRIPT']):\n",
        "    !git clone https://github.com/nicknochnack/GenerateTFRecord {paths['SCRIPTS_PATH']}\n",
        "\n",
        "!python {files['TF_RECORD_SCRIPT']} -x {os.path.join(paths['IMAGE_PATH'], 'train')} -l {files['LABELMAP']} -o {os.path.join(paths['ANNOTATION_PATH'], 'train.record')} \n",
        "!python {files['TF_RECORD_SCRIPT']} -x {os.path.join(paths['IMAGE_PATH'], 'val')} -l {files['LABELMAP']} -o {os.path.join(paths['ANNOTATION_PATH'], 'test.record')} "
      ],
      "metadata": {
        "id": "4_Q7qloMwVMA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prepare config"
      ],
      "metadata": {
        "id": "SE_Si1tBwicy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "if os.name =='posix':\n",
        "    !cp {os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'pipeline.config')} {os.path.join(paths['CHECKPOINT_PATH'])}\n",
        "if os.name == 'nt':\n",
        "    !copy {os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'pipeline.config')} {os.path.join(paths['CHECKPOINT_PATH'])}"
      ],
      "metadata": {
        "id": "Ya9BSlFdwuPu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.protos import pipeline_pb2\n",
        "from google.protobuf import text_format"
      ],
      "metadata": {
        "id": "Eq_mz_kgwxU8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = config_util.get_configs_from_pipeline_file(files['PIPELINE_CONFIG'])\n",
        "\n",
        "pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()\n",
        "with tf.io.gfile.GFile(files['PIPELINE_CONFIG'], \"r\") as f:                                                                                                                                                                                                                     \n",
        "    proto_str = f.read()                                                                                                                                                                                                                                          \n",
        "    text_format.Merge(proto_str, pipeline_config)  \n",
        "\n",
        "pipeline_config.model.ssd.num_classes = len(labels)\n",
        "pipeline_config.train_config.batch_size = 2\n",
        "pipeline_config.train_config.fine_tune_checkpoint = os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'checkpoint', 'ckpt-0')\n",
        "pipeline_config.train_config.fine_tune_checkpoint_type = \"detection\"\n",
        "pipeline_config.train_input_reader.label_map_path= files['LABELMAP']\n",
        "pipeline_config.train_input_reader.tf_record_input_reader.input_path[:] = [os.path.join(paths['ANNOTATION_PATH'], 'train.record')]\n",
        "pipeline_config.eval_input_reader[0].label_map_path = files['LABELMAP']\n",
        "pipeline_config.eval_input_reader[0].tf_record_input_reader.input_path[:] = [os.path.join(paths['ANNOTATION_PATH'], 'test.record')]\n",
        "\n",
        "config_text = text_format.MessageToString(pipeline_config)                                                                                                                                                                                                        \n",
        "with tf.io.gfile.GFile(files['PIPELINE_CONFIG'], \"wb\") as f:                                                                                                                                                                                                                     \n",
        "    f.write(config_text)"
      ],
      "metadata": {
        "id": "TAWvl1M2w3IY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = config_util.get_configs_from_pipeline_file(files['PIPELINE_CONFIG'])\n",
        "print(config)"
      ],
      "metadata": {
        "id": "zzwfMtTDxEAv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Train"
      ],
      "metadata": {
        "id": "uKaq6QDOxHyh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "TRAINING_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'model_main_tf2.py')\n",
        "command = \"python {} --model_dir={} --pipeline_config_path={} --checkpoint_every_n=1000 --num_train_steps=100000\".format(TRAINING_SCRIPT, paths['CHECKPOINT_PATH'],files['PIPELINE_CONFIG'])\n",
        "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
        "tf.config.experimental.set_memory_growth(gpus[0], True)"
      ],
      "metadata": {
        "id": "ExFilMWWxMn2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir={paths['CHECKPOINT_PATH']}"
      ],
      "metadata": {
        "id": "18RbMikyxbxG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!{command}"
      ],
      "metadata": {
        "id": "ToTuu65PxS49"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Eval"
      ],
      "metadata": {
        "id": "Do7jRXdoxdFv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "command = \"python {} --model_dir={} --pipeline_config_path={} --checkpoint_dir={}\".format(TRAINING_SCRIPT, paths['CHECKPOINT_PATH'],files['PIPELINE_CONFIG'], paths['CHECKPOINT_PATH'])"
      ],
      "metadata": {
        "id": "s_T4Jf98xeDF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!{command}"
      ],
      "metadata": {
        "id": "GqpjUuuSxhH1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# DeepWeeds model"
      ],
      "metadata": {
        "id": "BfeVeVkrQgP1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# dw_model = tf.keras.models.load_model(os.path.join(paths['PRETRAINED_MODEL_PATH'], 'resnet.hdf5'))\n",
        "# tf.saved_model.save(dw_model, os.path.join(paths['PRETRAINED_MODEL_PATH'], 'deep_weeds'))"
      ],
      "metadata": {
        "id": "7Wm64BcBQflH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load from checkpoint"
      ],
      "metadata": {
        "id": "_0a-FhH13N4T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.builders import model_builder\n",
        "from object_detection.utils import config_util\n",
        "\n",
        "# Load pipeline config and build a detection model\n",
        "configs = config_util.get_configs_from_pipeline_file(files['PIPELINE_CONFIG'])\n",
        "detection_model = model_builder.build(model_config=configs['model'], is_training=False)\n",
        "\n",
        "# Restore checkpoint\n",
        "ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)\n",
        "ckpt.restore(os.path.join(paths['CHECKPOINT_PATH'], 'ckpt-84')).expect_partial()\n",
        "\n",
        "@tf.function\n",
        "def detect_fn(image):\n",
        "    image, shapes = detection_model.preprocess(image)\n",
        "    prediction_dict = detection_model.predict(image, shapes)\n",
        "    detections = detection_model.postprocess(prediction_dict, shapes)\n",
        "    return detections"
      ],
      "metadata": {
        "id": "eAPnsoAh3ORM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Inference"
      ],
      "metadata": {
        "id": "jUgsunfi3UhV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2 \n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "%matplotlib inline\n",
        "\n",
        "category_index = label_map_util.create_category_index_from_labelmap(files['LABELMAP'])"
      ],
      "metadata": {
        "id": "cqyA-0sL3U4q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "IMAGE_PATH = os.path.join(paths['IMAGE_PATH'], 'test', 'pm0_1_0.jpg')\n",
        "\n",
        "tf.config.run_functions_eagerly(True)\n",
        "\n",
        "img = cv2.imread(IMAGE_PATH)\n",
        "image_np = np.array(img)\n",
        "\n",
        "input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)\n",
        "detections = detect_fn(input_tensor)\n",
        "\n",
        "num_detections = int(detections.pop('num_detections'))\n",
        "detections = {key: value[0, :num_detections].numpy()\n",
        "              for key, value in detections.items()}\n",
        "detections['num_detections'] = num_detections\n",
        "\n",
        "# detection_classes should be ints.\n",
        "detections['detection_classes'] = detections['detection_classes'].astype(np.int64)\n",
        "\n",
        "label_id_offset = 1\n",
        "image_np_with_detections = image_np.copy()\n",
        "\n",
        "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "            image_np_with_detections,\n",
        "            detections['detection_boxes'],\n",
        "            detections['detection_classes']+label_id_offset,\n",
        "            detections['detection_scores'],\n",
        "            category_index,\n",
        "            use_normalized_coordinates=True,\n",
        "            max_boxes_to_draw=20,\n",
        "            min_score_thresh=.2,\n",
        "            agnostic_mode=False)\n",
        "\n",
        "plt.imshow(cv2.cvtColor(image_np_with_detections, cv2.COLOR_BGR2RGB))\n",
        "plt.rcParams[\"figure.figsize\"] = (20,20)\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "uXbL_xCi3bB5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}