bytetrack-inference-yolox.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ByteTrack Inference with Amazon SageMaker\n",
    "\n",
    "This notebook will demonstrate how to create an endpoint for real time inference with the trained FairMOT model. We will first deploy the trained model in Sagemaker using [BYOS](https://sagemaker-examples.readthedocs.io/en/latest/sagemaker-script-mode/sagemaker-script-mode.html) mode by using custom inference scripts. And then apply inference on each frame of the video by invoking the endpoint. The inference result will be saved to a local directory.\n",
    "SageMaker provided prebuilt containers for various frameworks like Scikit-learn, PyTorch, and XGBoost. For this example, we will use PyTorch prebuilt containers by defining a PyTorchModel instance. \n",
    "\n",
    "## 1. SageMaker Initialization \n",
    "First we upgrade SageMaker to the latest version. If your notebook is already using latest Sagemaker 2.x API, you may skip the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pip install --upgrade pip\n",
    "! python3 -m pip install --upgrade sagemaker\n",
    "! pip install cython\n",
    "! pip install loguru\n",
    "! pip install thop\n",
    "! pip install lap\n",
    "! pip install numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Get the numpy version\n",
    "numpy_version=$(python -c \"import numpy; print(numpy.__version__)\")\n",
    "\n",
    "# Check the version and perform different operations\n",
    "if [[ $numpy_version > \"1.19.5\" ]]; then\n",
    "    git clone https://github.com/samson-wang/cython_bbox.git\n",
    "    sed -i 's/DTYPE\\ =\\ np.float\\b/DTYPE\\ =\\ np.float64/g' cython_bbox/src/cython_bbox.pyx\n",
    "    cd cython_bbox && pip install .\n",
    "else\n",
    "    pip install cython_bbox\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import libraries and get execution role "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import json\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "\n",
    "role = (\n",
    "    get_execution_role()\n",
    ")  # provide a pre-existing role ARN as an alternative to creating a new role\n",
    "print(f\"SageMaker Execution Role:{role}\")\n",
    "\n",
    "client = boto3.client('sts')\n",
    "account = client.get_caller_identity()['Account']\n",
    "print(f'AWS account:{account}')\n",
    "\n",
    "session = boto3.session.Session()\n",
    "aws_region = session.region_name\n",
    "print(f\"AWS region:{aws_region}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Deploy YOLOX model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You need to complete training job on [bytetrack-training.ipynb](bytetrack-training.ipynb) before running the following steps. Script Mode in SageMaker allows you to take control of the training and inference process without having to go through the trouble of creating and maintaining your own docker containers. Here, since we want to use a custom pytorch algorithm, we just use the AWS-provided Pytorch container and pass our own inference code. On your behalf, the SageMaker Python SDK will package this entry point script (sagemaker-serving/code/inferece.py), upload it to S3, and set two environment variables that are read at runtime and load the custom inference functions from the entry point script. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the s3 path for the model trained in [bytetrack-training.ipynb](bytetrack-training.ipynb) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%store -r s3_model_uri"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare tracking scripts\n",
    "we reuse the tracking processing scripts in the original [ByteTrack repo](https://github.com/ifzhang/ByteTrack/tree/main)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bucket = \"sagemaker-us-east-1-822507008821\"\n",
    "prefix = \"sm-bytetrack\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use pretrained model in our inference endpoint, we need to convert pretrained model format. We can use the defined YOLOX model in [ByteTrack](https://github.com/ifzhang/ByteTrack.git). Once endpoint deployment is done, we also use it to test endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "git clone --filter=blob:none --no-checkout --depth 1 --sparse https://github.com/ifzhang/ByteTrack.git && \\\n",
    "cd ByteTrack && \\\n",
    "git sparse-checkout set yolox && \\\n",
    "git checkout && \\\n",
    "cd ..\n",
    "cp -r ByteTrack/yolox yolox\n",
    "cp container-batch-inference/byte_tracker.py yolox/tracker/\n",
    "\n",
    "numpy_version=$(python -c \"import numpy; print(numpy.__version__)\")\n",
    "\n",
    "# Check the version and perform different operations\n",
    "if [[ $numpy_version > \"1.19.5\" ]]; then\n",
    "    sed -i 's/np.float\\b/np.float64/g' yolox/tracker/byte_tracker.py\n",
    "    sed -i 's/np.float\\b/np.float64/g' yolox/tracker/matching.py\n",
    "fi\n",
    "\n",
    "sudo rm -r ByteTrack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy an endpoint with the pretrained model (Optional)\n",
    "Note that if you just go through this solution and didn't train the model enough in training phase, model may not perform well with the trained model. You can use the pre-trained model to deploy an endpoint and demostrate the tracking solution. We can download pretrained yolox model and create a `model.tar.gz` file, and upload it to S3 Bucket. <span style=\"color:red\">Skip this section if you use the model trained on your own data.</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the below pretrained models.\n",
    "- [bytetrack_x_mot17](https://drive.google.com/u/0/uc?id=1P4mY0Yyd3PPTybgZkjMYhFri88nTmJX5)\n",
    "- [bytetrack_m_mot17](https://drive.google.com/u/0/uc?id=11Zb0NN_Uu7JwUd9e6Nk8o2_EUfxWqsun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install gdown\n",
    "# bytetrack_m_mot17\n",
    "!gdown https://drive.google.com/u/0/uc?id=11Zb0NN_Uu7JwUd9e6Nk8o2_EUfxWqsun\n",
    "\n",
    "# bytetrack_x_mot17\n",
    "#!gdown https://drive.google.com/u/0/uc?id=1P4mY0Yyd3PPTybgZkjMYhFri88nTmJX5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from yolox.models import YOLOPAFPN, YOLOX, YOLOXHead\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "from yolox.utils import load_ckpt\n",
    "\n",
    "model_name = \"bytetrack_m_mot17.pth.tar\"\n",
    "#model_name = \"bytetrack_x_mot17.pth.tar\"\n",
    "\n",
    "if model_name == \"bytetrack_x_mot17.pth.tar\":\n",
    "    # X\n",
    "    depth = 1.33\n",
    "    width = 1.25\n",
    "elif model_name == \"bytetrack_m_mot17.pth.tar\":\n",
    "    # M\n",
    "    depth = 0.67\n",
    "    width = 0.75\n",
    "num_classes = 1\n",
    "\n",
    "in_channels = [256, 512, 1024]\n",
    "backbone = YOLOPAFPN(depth, width, in_channels=in_channels)\n",
    "head = YOLOXHead(num_classes, width, in_channels=in_channels)\n",
    "model = YOLOX(backbone, head)\n",
    "\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "ckpt = torch.load(model_name, map_location=device)\n",
    "model.load_state_dict(ckpt[\"model\"])\n",
    "model = model.to(device)\n",
    "\n",
    "input_shape = [1, 3, 800, 1440]\n",
    "trace = torch.jit.trace(model.float().eval(), torch.zeros(input_shape).to(device).float())\n",
    "trace.save(\"model.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Package model into `model.tar.gz` and upload it to s3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3_model_uri = f\"s3://{bucket}/{prefix}/pretrained-model/model.tar.gz\"\n",
    "\n",
    "!mkdir pretrained_model\n",
    "!mv model.pth pretrained_model\n",
    "!cd pretrained_model && tar -czvf  model.tar.gz .\n",
    "!aws s3 cp pretrained_model/model.tar.gz $s3_model_uri\n",
    "!rm -r pretrained_model\n",
    "!rm $model_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an instance of PyTorchModel\n",
    "\n",
    "Inside inference.py, we defined 4 functions: `model_fn`, `input_fn`, `predict_fn`, and `output_fn`. These function handlers are automatically loaded and executed at runtime. The argument variables for these function handlers are predefined by SageMaker prebuilt containers.\n",
    "- `model_fn` handler loads the model according to s3 path\n",
    "- `input_fn` handler defines steps to pre-process the image passed by the requests\n",
    "- `predict_fn` handler defines the model forward computing steps\n",
    "- `output_fn` handler defines the post-processing steps after getting the inference results. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat sagemaker-serving/code/inference.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pytorch_model = PyTorchModel(\n",
    "    model_data=s3_model_uri,\n",
    "    role=role,\n",
    "    source_dir=\"sagemaker-serving/code\",\n",
    "    entry_point=\"inference.py\",\n",
    "    framework_version=\"1.7.1\",\n",
    "    py_version=\"py3\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In this cell, you would need to define the endpoint name.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "endpoint_name = \"yolox-detection-20230608-9\"\n",
    "pytorch_model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=\"ml.p3.2xlarge\",\n",
    "    endpoint_name=endpoint_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Run Multi-Object Tracking with YOLOX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell defines a function which captures each frame of a video, passes each frame to the inference endpoint defined in the previous step and saves the resulting frames to a local directory (save_folder). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from yolox.tracker.byte_tracker import BYTETracker\n",
    "import cv2\n",
    "import time\n",
    "from yolox.tracking_utils.timer import Timer\n",
    "import os.path as osp\n",
    "import os\n",
    "import torch\n",
    "from yolox.utils.visualize import plot_tracking\n",
    "\n",
    "sm_runtime = boto3.Session().client(\"sagemaker-runtime\")\n",
    "\n",
    "def imageflow_demo(endpoint_name, video_path=\"\", save_folder=\"\"):\n",
    "    cap = cv2.VideoCapture(video_path)\n",
    "    \n",
    "    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)  # float\n",
    "    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float\n",
    "    \n",
    "    print(f\"width: {width}, height: {height}\")\n",
    "    \n",
    "    fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "    \n",
    "    os.makedirs(save_folder, exist_ok=True)\n",
    "    save_path = osp.join(save_folder, video_path.split(\"/\")[-1])\n",
    "    \n",
    "    print(f\"video save_path is {save_path}\")\n",
    "    vid_writer = cv2.VideoWriter(\n",
    "        save_path, cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )\n",
    "    \n",
    "    aspect_ratio_thresh = 1.6\n",
    "    min_box_area = 10\n",
    "    tracker = BYTETracker(\n",
    "        frame_rate=30,\n",
    "        track_thresh=0.5,\n",
    "        track_buffer=30,\n",
    "        mot20=False,\n",
    "        match_thresh=0.8\n",
    "    )\n",
    "    timer = Timer()\n",
    "    frame_id = 0\n",
    "    results = []\n",
    "    while True:\n",
    "        ret_val, frame = cap.read()\n",
    "        if ret_val:\n",
    "            cv2.imwrite(f'datasets/frame_{frame_id}.png', frame)\n",
    "            with open(f\"datasets/frame_{frame_id}.png\", \"rb\") as f:\n",
    "                payload = f.read()\n",
    "            \n",
    "            timer.tic()\n",
    "            response = sm_runtime.invoke_endpoint(\n",
    "                EndpointName=endpoint_name, ContentType=\"application/x-image\", Body=payload\n",
    "            )\n",
    "            outputs = json.loads(response[\"Body\"].read().decode())\n",
    "            \n",
    "            if outputs[0] is not None:\n",
    "                online_targets = tracker.update(torch.as_tensor(outputs[0]), [height, width], (800, 1440))\n",
    "                online_tlwhs = []\n",
    "                online_ids = []\n",
    "                online_scores = []\n",
    "                for t in online_targets:\n",
    "                    tlwh = t.tlwh\n",
    "                    tid = t.track_id\n",
    "                    vertical = tlwh[2] / tlwh[3] > aspect_ratio_thresh\n",
    "                    if tlwh[2] * tlwh[3] > min_box_area and not vertical:\n",
    "                        online_tlwhs.append(tlwh)\n",
    "                        online_ids.append(tid)\n",
    "                        online_scores.append(t.score)\n",
    "                        results.append(\n",
    "                            f\"{frame_id},{tid},{tlwh[0]:.2f},{tlwh[1]:.2f},{tlwh[2]:.2f},{tlwh[3]:.2f},{t.score:.2f},-1,-1,-1\\n\"\n",
    "                        )\n",
    "                timer.toc()\n",
    "                online_im = plot_tracking(\n",
    "                    frame, online_tlwhs, online_ids, frame_id=frame_id + 1, fps=1. / timer.average_time\n",
    "                )\n",
    "            else:\n",
    "                timer.toc()\n",
    "                online_im = frame\n",
    "            if frame_id % 20 == 0:\n",
    "                print('Processing frame {} ({:.2f} fps)'.format(frame_id, 1./ timer.average_time))\n",
    "            \n",
    "            vid_writer.write(online_im)\n",
    "        else:\n",
    "            break\n",
    "        frame_id += 1\n",
    "\n",
    "    res_file = osp.join(save_folder, f\"log.txt\")\n",
    "    with open(res_file, 'w') as f:\n",
    "        f.writelines(results)\n",
    "    print(f\"save results to {res_file}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download a video from a public source. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir datasets\n",
    "!wget https://raw.githubusercontent.com/ifzhang/FairMOT/master/videos/MOT16-03.mp4 -O datasets/MOT16-03.mp4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apply the demo function on the downloaded video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "video_path=\"datasets/MOT16-03.mp4\"\n",
    "save_folder=\"track_res\"\n",
    "endpoint_name = \"yolox-detection-20230608-9\"\n",
    "imageflow_demo(endpoint_name, video_path, save_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}