predict_yolo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import toml\n",
    "import sys\n",
    "import os\n",
    "import time\n",
    "import json\n",
    "import numpy as np\n",
    "from IPython.display import display, Image, clear_output\n",
    "\n",
    "import pynq\n",
    "\n",
    "sys.path.append(os.path.abspath(\"data_management\"))\n",
    "sys.path.append(os.path.abspath(\"onnx_models\"))\n",
    "sys.path.append(os.path.abspath(\"hardware\"))\n",
    "\n",
    "import data_preprocessing as dprep\n",
    "import data_postprocessing as dpostp\n",
    "import onnx_inference as onnx_inf\n",
    "import fpgaconvnet_driver as hw_part\n",
    "\n",
    "%matplotlib inline\n",
    "!sudo usermod -a -G video root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "yolo_cfg = toml.load('yolov5n.toml')\n",
    "inf_exec = yolo_cfg['inf_exec']\n",
    "imgsz = yolo_cfg['input_data']['imgsz']\n",
    "conf_thres = yolo_cfg['predictor']['conf_thres']\n",
    "iou_thres = yolo_cfg['predictor']['iou_thres']\n",
    "classes = yolo_cfg['predictor']['classes']\n",
    "if not classes:\n",
    "    classes = None\n",
    "onnx_model_path = yolo_cfg['onnx_model_path']\n",
    "if inf_exec == \"fpga\":\n",
    "    onnx_model_path = f\"{onnx_model_path.split('.onnx')[0]}_head.onnx\"\n",
    "visualize = yolo_cfg['visualize']\n",
    "out_img_path = yolo_cfg['out_img_path']\n",
    "input_source = yolo_cfg['source']\n",
    "bitstream_path = yolo_cfg['hardware']['bitstream_path']\n",
    "weights_lookup_table = yolo_cfg['hardware']['weights_lookup_table']\n",
    "weights_path = yolo_cfg['hardware']['weights_path']\n",
    "input_bp = yolo_cfg['hardware']['input_binary_point']\n",
    "output_bp = yolo_cfg['hardware']['output_binary_point']\n",
    "fifo_depth = yolo_cfg['hardware']['fifo_depth']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "stride, names, session, output_names = onnx_inf.load_model(onnx_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if inf_exec == \"fpga\":\n",
    "    import fpgaconvnet_driver as hw_part\n",
    "    #########################\n",
    "    ### fpgaConvNet Setup ###\n",
    "    #########################\n",
    "\n",
    "    # initialise partition\n",
    "    partition = hw_part.Partition(bitstream_path, 5) \n",
    "\n",
    "    # add input buffers\n",
    "    partition.add_input_buffer(0, 0, [320, 320, 3], bp=input_bp)\n",
    "\n",
    "    # add output buffers\n",
    "    partition.add_output_buffer(2, 2, [40, 40, 256], bp=output_bp, streams=2)\n",
    "    partition.add_output_buffer(3, 3, [20, 20, 256], bp=output_bp, streams=2)\n",
    "    partition.add_output_buffer(4, 4, [10, 10, 256], bp=output_bp, streams=2)\n",
    "\n",
    "    # create fifos\n",
    "    partition.add_fifo(0, 0, 2, 40*40*64 , burst=fifo_depth[0])\n",
    "    partition.add_fifo(1, 1, 1, 20*20*128, burst=fifo_depth[1])\n",
    "\n",
    "    # # setup hardware\n",
    "    partition.reset_hardware()\n",
    "\n",
    "    # get the lookup table for the weights\n",
    "    with open(weights_lookup_table, \"r\") as f:\n",
    "        lookup = json.load(f)\n",
    "\n",
    "    # iterate over the weights\n",
    "    for layer, idx in lookup.items():\n",
    "\n",
    "        # allocate weights and load them\n",
    "        start_time = time.perf_counter() \n",
    "        partition.reload_weights(idx, f\"{weights_path}/{layer}.dat\")\n",
    "        pred_time = (time.perf_counter() - start_time)*1000\n",
    "        print(f\"[{idx}] {layer} loaded! ({pred_time:.2f} ms)\")\n",
    "\n",
    "    # setup hardware\n",
    "    partition.reset_hardware()\n",
    "    partition.start_hardware() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if input_source == 'webcam':\n",
    "    cap = cv2.VideoCapture(0) \n",
    "else:\n",
    "    cap = cv2.VideoCapture(input_source)\n",
    "\n",
    "# Capture frame-by-frame\n",
    "ret, frame = cap.read()\n",
    "    \n",
    "frame_count = 0\n",
    "try:\n",
    "    while(cap.isOpened()):\n",
    "        # Capture frame-by-frame\n",
    "        ret, frame = cap.read()\n",
    "\n",
    "        # Check that frame is not empty\n",
    "        if ret:\n",
    "            orig_img = frame.copy()\n",
    "\n",
    "            start_time = time.perf_counter()\n",
    "            img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "            img = dprep.img_preprocess(img, imgsz, stride)\n",
    "            preprocess_time = (time.perf_counter() - start_time) * 1000\n",
    "            \n",
    "            if inf_exec == \"cpu\":\n",
    "                predictions, pred_time = onnx_inf.model_inf(img, session, output_names)\n",
    "            elif inf_exec == \"fpga\":\n",
    "                out0, out1, out2, pred_time = hw_part.run_fpgaconvnet(partition, img[0])\n",
    "                predictions = session.run(output_names, {\n",
    "                    \"/model.24/m.0/Conv_output_0\": np.expand_dims(out0[:255,:,:], axis=0),\n",
    "                    \"/model.24/m.1/Conv_output_0\": np.expand_dims(out1[:255,:,:], axis=0),\n",
    "                    \"/model.24/m.2/Conv_output_0\": np.expand_dims(out2[:255,:,:], axis=0),\n",
    "                })[0]\n",
    "\n",
    "            start_time = time.perf_counter()\n",
    "            predictions = dpostp.yolo_nms(predictions, conf_thres, iou_thres, classes) \n",
    "            postprocess_time = (time.perf_counter() - start_time) * 1000\n",
    "\n",
    "            visualize_time = 0\n",
    "            for pred in predictions:\n",
    "                \n",
    "                start_time = time.perf_counter()\n",
    "                pred[:, :4] = dpostp.scale_boxes((imgsz, imgsz), pred[:, :4], orig_img.shape).round()\n",
    "                postprocess_time += (time.perf_counter() - start_time) * 1000\n",
    "                \n",
    "                start_time = time.perf_counter()\n",
    "                for p in pred:\n",
    "\n",
    "                    if visualize:\n",
    "                        cv2.rectangle(orig_img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), (0, 255, 0), 2)\n",
    "                        cv2.putText(orig_img, names[int(p[5])], (int(p[0]), int(p[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)\n",
    "                visualize_time += (time.perf_counter() - start_time) * 1000\n",
    "\n",
    "            start_time = time.perf_counter()\n",
    "            if visualize:\n",
    "\n",
    "                # add the time for inference\n",
    "                cv2.putText(orig_img, f\"latency: {pred_time:.2f} ms\", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)\n",
    "   \n",
    "                _, disp_img = cv2.imencode('.jpg', orig_img)\n",
    "                img_show = Image(data=disp_img)\n",
    "                display(img_show)\n",
    "                clear_output(wait=True)\n",
    "                frame_count += 1\n",
    "        else:\n",
    "            break\n",
    "except KeyboardInterrupt as e:\n",
    "    print(e)\n",
    "    cap.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    cap.release()\n",
    "    cap.destroyAllWindows()\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "    pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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": 2
}