🚀 [Update] the code with Vec2Box in deploy part

examples/notebook_inference.ipynb

@@ -10,7 +10,11 @@
     "from hydra import compose, initialize\n",
     "from PIL import Image \n",
     "\n",
-    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes"
+    "# Ensure that the necessary repository is cloned and installed. You may need to run: \n",
+    "# git clone [email protected]:WongKinYiu/YOLO.git\n",
+    "# cd YOLO \n",
+    "# pip install .\n",
+    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes, Vec2Box"
    ]
   },
   {
@@ -23,9 +27,8 @@
     "CONFIG_NAME = \"config\"\n",
     "\n",
     "DEVICE = 'cuda:0'\n",
-    "# If weight not exist there, it will try to download it\n",
-    "\n",
-    "WEIGHT_PATH = '../weights/v9-cnw.pt' \n",
+    "CLASS_NUM = 80\n",
+    "WEIGHT_PATH = '../weights/v9-c.pt' \n",
     "IMAGE_PATH = '../demo/images/inference/image.png'\n",
     "\n",
     "custom_logger()\n",
@@ -40,9 +43,10 @@
    "outputs": [],
    "source": [
     "with initialize(config_path=CONFIG_PATH, version_base=None, job_name=\"notebook_job\"):\n",
-    "    cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\"])\n",
-    "    model = create_model(cfg.model, WEIGHT_PATH).to(device)\n",
-    "    transform = AugmentationComposer([], cfg.image_size)"
+    "    cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\", \"model=v9-c-deploy\"])\n",
+    "    model = create_model(cfg.model, class_num=CLASS_NUM, weight_path=WEIGHT_PATH).to(device)\n",
+    "    transform = AugmentationComposer([], cfg.image_size)\n",
+    "    vec2box = Vec2Box(model, cfg.image_size, device)"
    ]
   },
   {
@@ -63,7 +67,9 @@
    "source": [
     "with torch.no_grad():\n",
     "    predict = model(image)\n",
-    "predict_box = bbox_nms(predict[1][0], cfg.task.nms)\n",
+    "    predict = vec2box(predict[\"Main\"])\n",
+    "\n",
+    "predict_box = bbox_nms(predict[0], predict[2], cfg.task.nms)\n",
     "draw_bboxes(image, predict_box, save_path='../demo/images/output/', idx2label=cfg.class_list)"
    ]
   },

yolo/config/model/v9-c-deploy.yaml

@@ -0,0 +1,77 @@
+anchor:
+  reg_max: 16
+
+model:
+  backbone:
+    - Conv:
+        args: {out_channels: 64, kernel_size: 3, stride: 2}
+        source: 0
+    - Conv:
+        args: {out_channels: 128, kernel_size: 3, stride: 2}
+    - RepNCSPELAN:
+        args: {out_channels: 256, part_channels: 128}
+
+    - ADown:
+        args: {out_channels: 256}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 256}
+        tags: B3
+
+    - ADown:
+        args: {out_channels: 512}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: B4
+
+    - ADown:
+        args: {out_channels: 512}
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: B5
+
+  neck:
+    - SPPELAN:
+        args: {out_channels: 512}
+        tags: N3
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B4]
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: N4
+
+    - UpSample:
+        args: {scale_factor: 2, mode: nearest}
+    - Concat:
+        source: [-1, B3]
+
+  head:
+    - RepNCSPELAN:
+        args: {out_channels: 256, part_channels: 256}
+        tags: P3
+
+    - ADown:
+        args: {out_channels: 256}
+    - Concat:
+        source: [-1, N4]
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: P4
+
+    - ADown:
+        args: {out_channels: 512}
+    - Concat:
+        source: [-1, N3]
+    - RepNCSPELAN:
+        args: {out_channels: 512, part_channels: 512}
+        tags: P5
+
+  detection:
+    - MultiheadDetection:
+        source: [P3, P4, P5]
+        tags: Main
+        args:
+            reg_max: ${model.anchor.reg_max}
+        output: True

yolo/utils/deploy_utils.py

@@ -35,8 +35,13 @@ class FastModelLoader:
 
         def onnx_forward(self: InferenceSession, x: Tensor):
             x = {self.get_inputs()[0].name: x.cpu().numpy()}
-            x = [torch.from_numpy(y) for y in self.run(None, x)]
-            return [x]
+            model_outputs, layer_output = [], []
+            for idx, predict in enumerate(self.run(None, x)):
+                layer_output.append(torch.from_numpy(predict))
+                if idx % 3 == 2:
+                    model_outputs.append(layer_output)
+                    layer_output = []
+            return {"Main": model_outputs}
 
         InferenceSession.__call__ = onnx_forward
         try: