Fix problems in colab && update version

BabitMF · Aug 7, 2024 · 10fde57 · 10fde57
1 parent f02766b
commit 10fde57
Show file tree

Hide file tree

Showing 6 changed files with 381 additions and 362 deletions.
diff --git a/bmf/demo/face_detect/facedetect_demo_colab.ipynb b/bmf/demo/face_detect/facedetect_demo_colab.ipynb
@@ -54,7 +54,6 @@
         "%cd ..\n",
         "!rm -rf python uff graphsurgeon onnx_graphsurgeon\n",
         "%cd /content\n",
-        "!pip install tensorrt --force-reinstall\n",
         "%env LIBRARY_PATH=/usr/local/cuda/lib64/stubs:/content/trt/TensorRT-8.6.1.6/lib\n",
         "%env LD_LIBRARY_PATH=/content/trt/TensorRT-8.6.1.6/lib:/usr/lib64-nvidia:/content/trt/TensorRT-8.6.1.6/lib\n",
         "%env PATH=/opt/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/tools/node/bin:/tools/google-cloud-sdk/bin:/content/trt/TensorRT-8.6.1.6/bin"
@@ -179,6 +178,8 @@
       },
       "outputs": [],
       "source": [
+        "%%writefile /content/trt_face_detect.py\n",
+        "\n",
         "import tensorrt as trt\n",
         "import torch\n",
         "import torch.nn.functional as F\n",
@@ -529,6 +530,8 @@
       },
       "outputs": [],
       "source": [
+        "%%writefile /content/trt.py\n",
+        "\n",
         "import sys\n",
         "import torch\n",
         "import numpy as np\n",
@@ -552,7 +555,7 @@
         "            \"input_shapes\": {\n",
         "                \"input\": [1, 3, 480, 640]\n",
         "            }\n",
-        "        }, entry=\"__main__.trt_face_detect\").encode(\n",
+        "        }).encode(\n",
         "        None, {\n",
         "            \"output_path\": \"./trt_out.mp4\",\n",
         "            \"video_params\": {\n",
@@ -566,6 +569,17 @@
         "    main()\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PkEp8bKrAXQX"
+      },
+      "outputs": [],
+      "source": [
+        "!python3 trt.py"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

diff --git a/bmf/demo/video_enhance/bmf-enhance-demo.ipynb b/bmf/demo/video_enhance/bmf-enhance-demo.ipynb