{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "078afd41-6103-4d36-8b39-2bf0cf662599",
   "metadata": {},
   "source": [
    "# Custom YOLOv5 Train and Deploy on Amazon SageMaker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4b2cd73-cadd-4a8b-9c8d-55a36eafce19",
   "metadata": {},
   "source": [
    "In this notebook we will train and deploy custom YOLOv5 object detection CV model with Amazon SageMaker Training Jobs and Endpoints.\n",
    "\n",
    "**Steps:**\n",
    "\n",
    "0. Initial configuration.\n",
    "1. Locate a labeled dataset with YOLOv5 expected format.\n",
    "2. Train the custom YOLOv5 model with SageMaker Training Jobs.\n",
    "3. Deploy the model with SageMaker Endpoints."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33b14bf6-40e5-46b0-a0f6-8f1aa43fa9be",
   "metadata": {},
   "source": [
    "## 0. Initial Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fd63dba-6e16-4bc7-a6a6-a3a80f6ab01a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install -qU sagemaker\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import boto3\n",
    "import sagemaker\n",
    "import uuid\n",
    "import time\n",
    "import cv2\n",
    "import glob\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline \n",
    "from sagemaker.pytorch.estimator import PyTorch\n",
    "from sagemaker.session import TrainingInput\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.utils import name_from_base\n",
    "from sagemaker.pytorch import PyTorchModel\n",
    "from sagemaker.serializers import DataSerializer\n",
    "from sagemaker.deserializers import JSONDeserializer\n",
    "sm_session = sagemaker.Session()\n",
    "role = get_execution_role()\n",
    "s3_resource = boto3.resource('s3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2699467-fa11-4fb6-8c4d-17941ea8f7ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone --quiet https://github.com/ultralytics/yolov5\n",
    "!cp -r helper-code/* yolov5/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef53157e-e980-4f8e-bcf2-879d698de5ab",
   "metadata": {},
   "source": [
    "## 1. Locate a labeled dataset with YOLOv5 expected format."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4810ead6-444b-47bf-a57e-c696e3f3ec44",
   "metadata": {},
   "source": [
    "Before we train a custom YOLOv5 model, we need to have a labeled dataset. In the previous notebook \"0 - Label your dataset with Amazon SageMaker GroundTruth\" you will be able to label your own dataset and transform it into YOLOv5 expected format or use an example custom dataset. Once you have run through one of the two options you will have available the S3 dataset location and labels used.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8120ec27-1e47-4bd1-a925-1ef791f70be4",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_s3_uri = \"\"\n",
    "labels = [\"\",\"\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f2a7284-1cb9-4fba-8333-76280927be70",
   "metadata": {},
   "source": [
    "### Download the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f19a7c76-163a-4894-a0c8-abc0478f6cfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_s3_path(s3_path):\n",
    "    path_parts=s3_path.replace(\"s3://\",\"\").split(\"/\")\n",
    "    bucket=path_parts.pop(0)\n",
    "    key=\"/\".join(path_parts)\n",
    "    return bucket, key\n",
    "\n",
    "def download_dataset(bucket_name, folder):\n",
    "    bucket = s3_resource.Bucket(bucket_name)\n",
    "    for obj in bucket.objects.filter(Prefix = folder):\n",
    "        if not os.path.exists(os.path.dirname(obj.key)):\n",
    "            os.makedirs(os.path.dirname(obj.key))\n",
    "        if os.path.splitext(obj.key)[1]:\n",
    "            bucket.download_file(obj.key, obj.key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fb56366-53e1-43fe-9ecf-77d39da90b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket,dataset_name = split_s3_path(dataset_s3_uri)\n",
    "download_dataset(bucket, dataset_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb704665-7d73-4db5-b54e-e2b971ccfd51",
   "metadata": {},
   "source": [
    "### Lets explore our dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d89de35-c6ba-45ce-b221-cbf4e5ab0ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for filename in glob.iglob(dataset_name + '**', recursive=True):\n",
    "     print(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfe17888-ce81-4c31-81cc-ce95a0212b7c",
   "metadata": {},
   "source": [
    "#### Now let's add these data sources to the data library in the yolov5 folder for our model to train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e1c6d9-dbce-4eb8-b168-fd0266b4786e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"yolov5/data/custom-model.yaml\", 'w') as target:\n",
    "    target.write(\"path: /opt/ml/input/data/training\\n\")\n",
    "    target.write(\"train: images/train\\n\")\n",
    "    target.write(\"val: images/validation\\n\")\n",
    "    target.write(\"names:\\n\")\n",
    "    for i, label in enumerate(labels):\n",
    "        target.write(\"  {}: {}\\n\".format(i, label))\n",
    "        \n",
    "with open('yolov5/data/custom-model.yaml') as file:\n",
    "    lines = file.readlines()\n",
    "    for line in lines:\n",
    "        print(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62a1251b-7e0f-4f96-bc11-d400a0ba04bd",
   "metadata": {},
   "source": [
    "## 3. Train the custom YOLOv5 model with SageMaker Training Jobs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b253a75e-6830-4577-91e0-4b4df8c268f9",
   "metadata": {},
   "source": [
    "#### First let's send our training data to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d788b55-71d1-45ac-9c8b-ccf8208819cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_name = \"yolov5-t\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97834e07-2bb5-4bf4-9eb3-e7630197acab",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = '{}-{}'.format(training_name,str(uuid.uuid4()))\n",
    "print(job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "306f005b-2073-4496-a963-0ab1aadfb27a",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparameters={\n",
    "    \"workers\":\"8\",\n",
    "    \"device\": \"0\",\n",
    "    \"batch-size\": \"8\",\n",
    "    \"epochs\": 50,\n",
    "    \"data\": \"custom-model.yaml\",\n",
    "    \"weights\": \"yolov5s.pt\",\n",
    "    \"project\": \"/opt/ml/model\"\n",
    "}\n",
    "\n",
    "estimator = PyTorch(\n",
    "    framework_version='1.11.0',\n",
    "    py_version='py38',\n",
    "    entry_point='train.py',\n",
    "    source_dir='yolov5',\n",
    "    hyperparameters=hyperparameters,\n",
    "    instance_count=1,\n",
    "    instance_type='ml.g4dn.xlarge',\n",
    "    role=role,\n",
    "    disable_profiler=True, \n",
    "    debugger_hook_config=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "903f4558-1069-48ed-bf19-a208de4eccee",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_input = TrainingInput(dataset_s3_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a238bc87-989b-4e83-9420-1c8a330d64a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit(train_input, job_name=job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d37759f-2c22-4f9a-bd6b-492dad543be1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"Model-\"+job_name\n",
    "model_data = 's3://{}/{}/output/model.tar.gz'.format(sm_session.default_bucket(), job_name)\n",
    "print(model_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7637620a-bc37-495c-b4a1-b8b337f9f49d",
   "metadata": {},
   "source": [
    "## 4. Deploy your model to a SM Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b40a48e-03f4-4e94-a262-f1de14934abb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = PyTorchModel(\n",
    "    entry_point='detect.py',\n",
    "    source_dir='yolov5',\n",
    "    model_data=model_data,\n",
    "    framework_version='1.11.0',\n",
    "    py_version='py38',\n",
    "    role=role,\n",
    "    name=model_name\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b20c386-f663-4263-b8e7-e97759554948",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = model.deploy(initial_instance_count=1, instance_type='ml.c5.large')\n",
    "predictor.deserializer = JSONDeserializer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "549daa11-c68d-4905-ba75-a62f1a6b513b",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.serializer =DataSerializer(content_type=\"image/png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b49fac4-534d-4ed1-a922-fbd03a867585",
   "metadata": {},
   "source": [
    "### Display predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b620af19-e252-4c25-823f-e817e41ad04c",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_files_dir=\"test-images\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36463c30-7767-4414-8b65-3e3d9b2460a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_label (image, box, conf, label):\n",
    "    bbox = np.array(box).astype(np.int32)\n",
    "    cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), [255,0,0], 2, cv2.LINE_AA)\n",
    "    cv2.putText(image, \"{}:{}\".format(label,str(conf)[0:4]), (bbox[0], bbox[1] - 10),  0, 1e-3 * imgHeight, [255,0,0], 2)\n",
    "    \n",
    "def resize_bb(old, new, min_b, max_b):\n",
    "    old = np.array(old)\n",
    "    new = np.array(new)\n",
    "    min_b = np.array(min_b)\n",
    "    max_b = np.array(max_b)\n",
    "    min_xy = min_b/(old/new)\n",
    "    max_xy = max_b/(old/new)\n",
    "    return [int(min_xy[0]),int(min_xy[1]),int(max_xy[0]),int(max_xy[1])]\n",
    "\n",
    "def plot_image(img):\n",
    "    dpi = 80\n",
    "    figsize = imgWidth / float(dpi), imgHeight / float(dpi)    \n",
    "    fig = plt.figure(figsize=figsize)\n",
    "    ax = fig.add_axes([0, 0, 1, 1])\n",
    "    ax.axis('off')\n",
    "    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n",
    "\n",
    "def make_prediction(imgdir,image):\n",
    "    #Get predictions\n",
    "    img_path = \"{}/{}\".format(imgdir,image)\n",
    "    data = open(img_path, 'rb').read()\n",
    "    pr = json.loads(predictor.predict(data))\n",
    "    df = pd.DataFrame(data=pr[\"data\"], index = pr[\"index\"], columns = pr[\"columns\"])\n",
    "    \n",
    "    #Display labels\n",
    "    img = cv2.imread(img_path)\n",
    "    imgHeight,imgWidth,_ = img.shape\n",
    "\n",
    "    for index, row in df.iterrows():\n",
    "        if row['confidence'] > 0.3:\n",
    "            new_boxes = resize_bb([640,640],[imgWidth,imgHeight],[row['xmin'],row['ymin']],[row['xmax'],row['ymax']])\n",
    "            draw_label(img, new_boxes,row[\"confidence\"],row['name'])\n",
    "\n",
    "    plot_image(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ca707e9-58a7-4173-9271-0f3c187cdeca",
   "metadata": {},
   "outputs": [],
   "source": [
    "for image in os.listdir(test_files_dir):\n",
    "    if image.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
    "        make_prediction(test_files_dir,image)"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:eu-west-1:470317259841:image/datascience-1.0"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}