python/demos/run_custom_object_detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8a98475a-d000-4fa2-ba17-b51e33e72a93",
   "metadata": {},
   "source": [
    "# 転移学習モデルの推論実行"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0062cb29-2cac-414f-a7ae-269512de5c0d",
   "metadata": {},
   "source": [
    "https://github.com/pf-robotics/kachaka-transfer-learning に従い作成したonnxファイルを\"predictor.onnx\"として本ノートブックと同じディレクトリに配置して下さい。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed27f4e-7637-4317-bdc4-16ff825309f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "\n",
    "import kachaka_api\n",
    "import numpy as np\n",
    "import onnxruntime\n",
    "from IPython.display import clear_output, display\n",
    "from kachaka_api.util.vision import get_bbox_drawn_image_with_user_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b94f64-6deb-4323-b090-94da6f746aca",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = \"predictor.onnx\"\n",
    "session = onnxruntime.InferenceSession(model)\n",
    "input_name_list = [out.name for out in session.get_inputs()]\n",
    "output_name_list = [out.name for out in session.get_outputs()]\n",
    "\n",
    "class_num = 1\n",
    "label_names = [f\"{i}_class\" for i in range(class_num)]\n",
    "# up to 3 classes\n",
    "colors = [\"red\", \"blue\", \"green\"]\n",
    "scale_h = 720.0 / 224\n",
    "scale_w = 1280.0 / 398\n",
    "_SCORE_THRESHOLD = 0.5\n",
    "\n",
    "client = kachaka_api.aio.KachakaApiClient()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92191b5b-727f-462c-98d7-e3d2ff62a788",
   "metadata": {},
   "outputs": [],
   "source": [
    "stream_i = client.front_camera_ros_compressed_image.stream()\n",
    "stream_d = client.object_detection_features.stream()\n",
    "\n",
    "while True:\n",
    "    image, (features_header, features) = await asyncio.gather(\n",
    "        anext(stream_i), anext(stream_d)\n",
    "    )\n",
    "    inputs = {\n",
    "        feature.name.replace(\"_out\", \"\"): np.array(\n",
    "            feature.data, dtype=np.float32\n",
    "        ).reshape(feature.shape)\n",
    "        for feature in features\n",
    "    }\n",
    "    (bboxes, labels, scores) = session.run(output_name_list, inputs)\n",
    "    img = get_bbox_drawn_image_with_user_labels(\n",
    "        image,\n",
    "        label_names,\n",
    "        bboxes,\n",
    "        labels,\n",
    "        colors,\n",
    "        scores,\n",
    "        scale_h,\n",
    "        scale_w,\n",
    "        _SCORE_THRESHOLD,\n",
    "    )\n",
    "    clear_output(wait=True)\n",
    "    display(img)"
   ]
  }
 ],
 "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": 5
}