♻️ [Refactor] Code of examples, use PostProccess

examples/notebook_inference.ipynb

@@ -6,15 +6,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import sys\n",
+    "from pathlib import Path\n",
+    "\n",
     "import torch\n",
     "from hydra import compose, initialize\n",
     "from PIL import Image \n",
     "\n",
-    "# Ensure that the necessary repository is cloned and installed. You may need to run: \n",
-    "# git clone git@github.com:WongKinYiu/YOLO.git\n",
-    "# cd YOLO \n",
-    "# pip install .\n",
-    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes, Vec2Box"
+    "project_root = Path().resolve().parent\n",
+    "sys.path.append(str(project_root))\n",
+    "\n",
+    "from yolo import AugmentationComposer, Config, create_model, custom_logger, draw_bboxes, Vec2Box, PostProccess\n",
+    "from yolo.utils.bounding_box_utils import Anc2Box"
    ]
   },
   {
@@ -25,7 +28,7 @@
    "source": [
     "CONFIG_PATH = \"../yolo/config\"\n",
     "CONFIG_NAME = \"config\"\n",
-    "MODEL = \"v9-c\"\n",
+    "MODEL = \"v7-base\"\n",
     "\n",
     "DEVICE = 'cuda:0'\n",
     "CLASS_NUM = 80\n",
@@ -45,7 +48,9 @@
     "    cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\", f\"model={MODEL}\"])\n",
     "    model = create_model(cfg.model, class_num=CLASS_NUM).to(device)\n",
     "    transform = AugmentationComposer([], cfg.image_size)\n",
-    "    vec2box = Vec2Box(model, cfg.image_size, device)"
+    "    converter = Anc2Box(model, cfg.model.anchor, cfg.image_size, device)\n",
+    "    # converter = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)\n",
+    "    post_proccess = PostProccess(converter, cfg.task.nms)"
    ]
   },
   {
@@ -57,7 +62,7 @@
     "pil_image = Image.open(IMAGE_PATH)\n",
     "image, bbox, rev_tensor = transform(pil_image)\n",
     "image = image.to(device)[None]\n",
-    "rev_tensor = rev_tensor.to(device)"
+    "rev_tensor = rev_tensor.to(device)[None]"
    ]
   },
   {
@@ -68,10 +73,8 @@
    "source": [
     "with torch.no_grad():\n",
     "    predict = model(image)\n",
-    "    pred_class, _, pred_bbox = vec2box(predict[\"Main\"])\n",
+    "    pred_bbox = post_proccess(predict, rev_tensor)\n",
     "\n",
-    "pred_bbox = (pred_bbox / rev_tensor[0] - rev_tensor[None, None, 1:]) \n",
-    "pred_bbox = bbox_nms(pred_class, pred_bbox, cfg.task.nms)\n",
     "draw_bboxes(pil_image, pred_bbox, idx2label=cfg.class_list)"
    ]
   },
@@ -83,6 +86,23 @@
     "\n",
     "![image](../demo/images/output/visualize.png)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

examples/notebook_smallobject.ipynb

@@ -22,7 +22,6 @@
     "import torch\n",
     "from hydra import compose, initialize\n",
     "from PIL import Image \n",
-    "from einops import rearrange\n",
     "\n",
     "# Ensure that the necessary repository is cloned and installed. You may need to run: \n",
     "# git clone [email protected]:WongKinYiu/YOLO.git\n",
@@ -30,8 +29,8 @@
     "# pip install .\n",
     "project_root = Path().resolve().parent\n",
     "sys.path.append(str(project_root))\n",
-    "from yolo.config.config import NMSConfig\n",
-    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes, Vec2Box"
+    "\n",
+    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes, Vec2Box, NMSConfig, PostProccess"
    ]
   },
   {
@@ -63,7 +62,9 @@
     "    cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\", f\"model={MODEL}\"])\n",
     "    model = create_model(cfg.model, class_num=CLASS_NUM).to(device)\n",
     "    transform = AugmentationComposer([], cfg.image_size)\n",
-    "    vec2box = Vec2Box(model, cfg.image_size, device)"
+    "    vec2box = Vec2Box(model, cfg.image_size, device)\n",
+    "    post_proccess = PostProccess(vec2box, NMSConfig(0.5, 0.9))\n",
+    "    "
    ]
   },
   {
@@ -75,7 +76,7 @@
     "pil_image = Image.open(IMAGE_PATH)\n",
     "image, bbox, rev_tensor = transform(pil_image)\n",
     "image = image.to(device)[None]\n",
-    "rev_tensor = rev_tensor.to(device)"
+    "rev_tensor = rev_tensor.to(device)[None]"
    ]
   },
   {
@@ -114,7 +115,9 @@
     "    pred_class, _, pred_bbox = vec2box(predict[\"Main\"])\n",
     "pred_bbox[1:] = (pred_bbox[1: ] + total_shift[:, None]) / SLIDE\n",
     "pred_bbox = pred_bbox.view(1, -1, 4)\n",
-    "pred_class = pred_class.view(1, -1, 80)"
+    "pred_class = pred_class.view(1, -1, 80)\n",
+    "pred_bbox = (pred_bbox - rev_tensor[:, None, 1:]) / rev_tensor[:, 0:1, None]\n",
+    "predict_box = bbox_nms(pred_class, pred_bbox, NMSConfig(0.3, 0.5))\n"
    ]
   },
   {
@@ -123,7 +126,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "pred_bbox = (pred_bbox / rev_tensor[0] - rev_tensor[None, None, 1:]) "
+    "draw_bboxes(pil_image, predict_box, idx2label=cfg.class_list)"
    ]
   },
   {
@@ -131,10 +134,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "predict_box = bbox_nms(pred_class, pred_bbox, NMSConfig(0.5, 0.5))\n",
-    "draw_bboxes(pil_image, predict_box, idx2label=cfg.class_list)"
-   ]
+   "source": []
   }
  ],
  "metadata": {