🚀 [Add] deploy option, auto remove aux head

examples/notebook_inference.ipynb

@@ -43,8 +43,8 @@
    "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}\", \"model=v9-c-deploy\"])\n",
-    "    model = create_model(cfg.model, class_num=CLASS_NUM, weight_path=WEIGHT_PATH).to(device)\n",
+    "    cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\", \"model=v9-m\"])\n",
+    "    model = create_model(cfg.model, class_num=CLASS_NUM, weight_path=WEIGHT_PATH, device = device)\n",
     "    transform = AugmentationComposer([], cfg.image_size)\n",
     "    vec2box = Vec2Box(model, cfg.image_size, device)"
    ]
@@ -70,7 +70,7 @@
     "    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)"
+    "draw_bboxes(image, predict_box, idx2label=cfg.class_list)"
    ]
   },
   {
@@ -81,13 +81,6 @@
     "\n",
     "![image](../demo/images/output/visualize.png)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

yolo/utils/deploy_utils.py

@@ -9,14 +9,15 @@ from yolo.model.yolo import create_model
 
 
 class FastModelLoader:
-    def __init__(self, cfg: Config):
+    def __init__(self, cfg: Config, device):
         self.cfg = cfg
+        self.device = device
         self.compiler = cfg.task.fast_inference
         self._validate_compiler()
         self.model_path = f"{os.path.splitext(cfg.weight)[0]}.{self.compiler}"
 
     def _validate_compiler(self):
-        if self.compiler not in ["onnx", "trt"]:
+        if self.compiler not in ["onnx", "trt", "deploy"]:
             logger.warning(f"⚠️ Compiler '{self.compiler}' is not supported. Using original model.")
             self.compiler = None
         if self.cfg.device == "mps" and self.compiler == "trt":
@@ -28,7 +29,11 @@ class FastModelLoader:
             return self._load_onnx_model()
         elif self.compiler == "trt":
             return self._load_trt_model()
-        return create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight)
+        elif self.compiler == "deploy":
+            self.cfg.model.model.auxiliary = {}
+        return create_model(
+            self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight, device=self.device
+        )
 
     def _load_onnx_model(self):
         from onnxruntime import InferenceSession