🚸 [Add] converter function, auto choose conveter

examples/notebook_inference.ipynb

@@ -16,8 +16,15 @@
     "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"
+    "from yolo import (\n",
+    "    AugmentationComposer,\n",
+    "    Config,\n",
+    "    PostProccess,\n",
+    "    create_converter,\n",
+    "    create_model,\n",
+    "    custom_logger,\n",
+    "    draw_bboxes,\n",
+    ")"
    ]
   },
   {
@@ -48,8 +55,7 @@
     "    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",
-    "    converter = Anc2Box(model, cfg.model.anchor, cfg.image_size, device)\n",
-    "    # converter = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)\n",
+    "    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)\n",
     "    post_proccess = PostProccess(converter, cfg.task.nms)"
    ]
   },
@@ -86,23 +92,6 @@
     "\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": {

yolo/__init__.py

@@ -3,7 +3,7 @@ from yolo.model.yolo import create_model
 from yolo.tools.data_loader import AugmentationComposer, create_dataloader
 from yolo.tools.drawer import draw_bboxes
 from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
-from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms
+from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import custom_logger
 from yolo.utils.model_utils import PostProccess
@@ -18,6 +18,7 @@ all = [
     "Vec2Box",
     "Anc2Box",
     "bbox_nms",
+    "create_converter",
     "AugmentationComposer",
     "create_dataloader",
     "FastModelLoader",

yolo/lazy.py

@@ -10,7 +10,7 @@ from yolo.config.config import Config
 from yolo.model.yolo import create_model
 from yolo.tools.data_loader import create_dataloader
 from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
-from yolo.utils.bounding_box_utils import Vec2Box
+from yolo.utils.bounding_box_utils import create_converter
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import ProgressLogger
 from yolo.utils.model_utils import get_device
@@ -27,13 +27,14 @@ def main(cfg: Config):
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
         model = model.to(device)
 
-    vec2box = Vec2Box(model, cfg.image_size, device)
+    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
+
     if cfg.task.task == "train":
-        solver = ModelTrainer(cfg, model, vec2box, progress, device, use_ddp)
+        solver = ModelTrainer(cfg, model, converter, progress, device, use_ddp)
     if cfg.task.task == "validation":
-        solver = ModelValidator(cfg.task, cfg.dataset, model, vec2box, progress, device)
+        solver = ModelValidator(cfg.task, cfg.dataset, model, converter, progress, device)
     if cfg.task.task == "inference":
-        solver = ModelTester(cfg, model, vec2box, progress, device)
+        solver = ModelTester(cfg, model, converter, progress, device)
     progress.start()
     solver.solve(dataloader)
 

yolo/utils/bounding_box_utils.py

@@ -364,7 +364,15 @@ class Anc2Box:
         return preds_cls, None, preds_box, preds_cnf.sigmoid()
 
 
-def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor]):
+def create_converter(model_version: str = "v9-c", *args, **kwargs):
+    if "v7" in model_version:  # check model if v7
+        converter = Anc2Box(*args, **kwargs)
+    else:
+        converter = Vec2Box(*args, **kwargs)
+    return converter
+
+
+def bbox_nms(cls_dist: Tensor, bbox: Tensor, nms_cfg: NMSConfig, confidence: Optional[Tensor] = None):
     cls_dist = cls_dist.sigmoid() * (1 if confidence is None else confidence)
 
     # filter class by confidence