{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8d6b25d5-10b4-4470-b172-2ec670ba7ed7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/botsort/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import argparse\n",
    "import os\n",
    "import os.path as osp\n",
    "import time\n",
    "import cv2\n",
    "import torch\n",
    "\n",
    "from tools.yolox.data.data_augment import preproc\n",
    "from tools.yolox.exp import get_exp\n",
    "from tools.yolox.utils import fuse_model, get_model_info, postprocess\n",
    "from tools.yolox.utils.visualize import plot_tracking\n",
    "from tracker.bot_sort import BoTSORT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fae52375-3770-4fe9-a7a4-b9d854e7c830",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Predictor(object):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        trt_file=None,\n",
    "        decoder=None,\n",
    "        device=torch.device(\"cpu\"),\n",
    "        fp16=False\n",
    "    ):\n",
    "        self.model = model\n",
    "        self.decoder = decoder\n",
    "        self.num_classes = 1\n",
    "        self.confthre = 0.1\n",
    "        self.nmsthre = 0.7\n",
    "        self.test_size = (800, 1440)\n",
    "        self.device = device\n",
    "        self.fp16 = fp16\n",
    "        self.rgb_means = (0.485, 0.456, 0.406)\n",
    "        self.std = (0.229, 0.224, 0.225)\n",
    "\n",
    "    def inference(self, img):\n",
    "        img = cv2.imread(img)\n",
    "        height, width = img.shape[:2]\n",
    "        img, ratio = preproc(img, self.test_size, self.rgb_means, self.std)\n",
    "        img = torch.from_numpy(img).unsqueeze(0).float().to(self.device)\n",
    "        if self.fp16:\n",
    "            img = img.half()  # to FP16\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = self.model(img)\n",
    "            if self.decoder is not None:\n",
    "                outputs = self.decoder(outputs, dtype=outputs.type())\n",
    "            outputs = postprocess(outputs, self.num_classes, self.confthre, self.nmsthre)\n",
    "        return outputs\n",
    "    \n",
    "def get_model(config):\n",
    "    from yolox.models import YOLOPAFPN, YOLOX, YOLOXHead\n",
    "\n",
    "    def init_yolo(M):\n",
    "        for m in M.modules():\n",
    "            if isinstance(m, nn.BatchNorm2d):\n",
    "                m.eps = 1e-3\n",
    "                m.momentum = 0.03\n",
    "\n",
    "    if getattr(config, \"model\", None) is None:\n",
    "        in_channels = [256, 512, 1024]\n",
    "        backbone = YOLOPAFPN(config['depth'], config['width'], in_channels=in_channels)\n",
    "        head = YOLOXHead(config['num_classes'], config['width'], in_channels=in_channels)\n",
    "        model = YOLOX(backbone, head)\n",
    "\n",
    "    model.apply(init_yolo)\n",
    "    model.head.initialize_biases(1e-2)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2ba6a04e-2d5d-496d-8b92-0b4199704490",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'nn' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 8\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# device = torch.device(\"cuda\" if torch.cuda.is_available() is True)\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# if args.tsize is not None:\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;66;03m#     exp.test_size = (args.tsize, args.tsize)\u001b[39;00m\n\u001b[1;32m      6\u001b[0m config \u001b[38;5;241m=\u001b[39m {\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdepth\u001b[39m\u001b[38;5;124m'\u001b[39m : \u001b[38;5;241m1.33\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mwidth\u001b[39m\u001b[38;5;124m'\u001b[39m : \u001b[38;5;241m1.25\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mnum_classes\u001b[39m\u001b[38;5;124m'\u001b[39m : \u001b[38;5;241m1\u001b[39m}\n\u001b[0;32m----> 8\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mget_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcuda\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      9\u001b[0m model\u001b[38;5;241m.\u001b[39meval()\n\u001b[1;32m     11\u001b[0m ckpt_file \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbytetrack_x_mot17.pth.tar\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
      "Cell \u001b[0;32mIn[4], line 51\u001b[0m, in \u001b[0;36mget_model\u001b[0;34m(config)\u001b[0m\n\u001b[1;32m     48\u001b[0m     head \u001b[38;5;241m=\u001b[39m YOLOXHead(config[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mnum_classes\u001b[39m\u001b[38;5;124m'\u001b[39m], config[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mwidth\u001b[39m\u001b[38;5;124m'\u001b[39m], in_channels\u001b[38;5;241m=\u001b[39min_channels)\n\u001b[1;32m     49\u001b[0m     model \u001b[38;5;241m=\u001b[39m YOLOX(backbone, head)\n\u001b[0;32m---> 51\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43minit_yolo\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     52\u001b[0m model\u001b[38;5;241m.\u001b[39mhead\u001b[38;5;241m.\u001b[39minitialize_biases(\u001b[38;5;241m1e-2\u001b[39m)\n\u001b[1;32m     53\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m model\n",
      "File \u001b[0;32m~/anaconda3/envs/botsort/lib/python3.8/site-packages/torch/nn/modules/module.py:667\u001b[0m, in \u001b[0;36mModule.apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    629\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Applies ``fn`` recursively to every submodule (as returned by ``.children()``)\u001b[39;00m\n\u001b[1;32m    630\u001b[0m \u001b[38;5;124;03mas well as self. Typical use includes initializing the parameters of a model\u001b[39;00m\n\u001b[1;32m    631\u001b[0m \u001b[38;5;124;03m(see also :ref:`nn-init-doc`).\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    664\u001b[0m \u001b[38;5;124;03m    )\u001b[39;00m\n\u001b[1;32m    665\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    666\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 667\u001b[0m     \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    668\u001b[0m fn(\u001b[38;5;28mself\u001b[39m)\n\u001b[1;32m    669\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[0;32m~/anaconda3/envs/botsort/lib/python3.8/site-packages/torch/nn/modules/module.py:667\u001b[0m, in \u001b[0;36mModule.apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    629\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Applies ``fn`` recursively to every submodule (as returned by ``.children()``)\u001b[39;00m\n\u001b[1;32m    630\u001b[0m \u001b[38;5;124;03mas well as self. Typical use includes initializing the parameters of a model\u001b[39;00m\n\u001b[1;32m    631\u001b[0m \u001b[38;5;124;03m(see also :ref:`nn-init-doc`).\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    664\u001b[0m \u001b[38;5;124;03m    )\u001b[39;00m\n\u001b[1;32m    665\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    666\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 667\u001b[0m     \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    668\u001b[0m fn(\u001b[38;5;28mself\u001b[39m)\n\u001b[1;32m    669\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "    \u001b[0;31m[... skipping similar frames: Module.apply at line 667 (2 times)]\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/botsort/lib/python3.8/site-packages/torch/nn/modules/module.py:667\u001b[0m, in \u001b[0;36mModule.apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    629\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Applies ``fn`` recursively to every submodule (as returned by ``.children()``)\u001b[39;00m\n\u001b[1;32m    630\u001b[0m \u001b[38;5;124;03mas well as self. Typical use includes initializing the parameters of a model\u001b[39;00m\n\u001b[1;32m    631\u001b[0m \u001b[38;5;124;03m(see also :ref:`nn-init-doc`).\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    664\u001b[0m \u001b[38;5;124;03m    )\u001b[39;00m\n\u001b[1;32m    665\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    666\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 667\u001b[0m     \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    668\u001b[0m fn(\u001b[38;5;28mself\u001b[39m)\n\u001b[1;32m    669\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[0;32m~/anaconda3/envs/botsort/lib/python3.8/site-packages/torch/nn/modules/module.py:668\u001b[0m, in \u001b[0;36mModule.apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    666\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[1;32m    667\u001b[0m     module\u001b[38;5;241m.\u001b[39mapply(fn)\n\u001b[0;32m--> 668\u001b[0m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    669\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "Cell \u001b[0;32mIn[4], line 41\u001b[0m, in \u001b[0;36mget_model.<locals>.init_yolo\u001b[0;34m(M)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minit_yolo\u001b[39m(M):\n\u001b[1;32m     40\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m m \u001b[38;5;129;01min\u001b[39;00m M\u001b[38;5;241m.\u001b[39mmodules():\n\u001b[0;32m---> 41\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(m, \u001b[43mnn\u001b[49m\u001b[38;5;241m.\u001b[39mBatchNorm2d):\n\u001b[1;32m     42\u001b[0m             m\u001b[38;5;241m.\u001b[39meps \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1e-3\u001b[39m\n\u001b[1;32m     43\u001b[0m             m\u001b[38;5;241m.\u001b[39mmomentum \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.03\u001b[39m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'nn' is not defined"
     ]
    }
   ],
   "source": [
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() is True)\n",
    "\n",
    "# if args.tsize is not None:\n",
    "#     exp.test_size = (args.tsize, args.tsize)\n",
    "\n",
    "config = {'depth' : 1.33, 'width' : 1.25, 'num_classes' : 1}\n",
    "\n",
    "model = get_model(config).to(\"cuda\")\n",
    "model.eval()\n",
    "\n",
    "ckpt_file = \"bytetrack_x_mot17.pth.tar\"\n",
    "ckpt = torch.load(ckpt_file, map_location=\"cpu\")\n",
    "# load the model state dict\n",
    "model.load_state_dict(ckpt[\"model\"])\n",
    "model = fuse_model(model)\n",
    "model = model.half()  # to FP16\n",
    "\n",
    "trt_file = None\n",
    "decoder = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9c2ec13-78c9-4969-b494-654bae1d58f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = Predictor(model, trt_file, decoder, device=torch.device(\"cuda\"), fp16=True)\n",
    "current_time = time.localtime()\n",
    "if args.demo == \"image\" or args.demo == \"images\":\n",
    "    image_demo(predictor, vis_folder, current_time, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e85ac8f-9e33-4e7f-a809-ba8c16e0c09e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Detect objects\n",
    "test_size = (800, 1440)\n",
    "utputs, img_info = predictor.inference(img_path)\n",
    "scale = min(test_size[0] / float(img_info['height'], ), test_size[1] / float(img_info['width']))\n",
    "\n",
    "if outputs[0] is not None:\n",
    "    outputs = outputs[0].cpu().numpy()\n",
    "    detections = outputs[:, :7]\n",
    "    detections[:, :4] /= scale"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "botsort",
   "language": "python",
   "name": "botsort"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}