🐛 [Fix] #56 bugs, the using converter in ipynb

examples/notebook_TensorRT.ipynb

@@ -18,7 +18,15 @@
     "project_root = Path().resolve().parent\n",
     "sys.path.append(str(project_root))\n",
     "\n",
-    "from yolo import AugmentationComposer, bbox_nms, create_model, custom_logger, draw_bboxes, Vec2Box\n",
+    "from yolo import (\n",
+    "    AugmentationComposer, \n",
+    "    bbox_nms, \n",
+    "    create_model, \n",
+    "    custom_logger, \n",
+    "    create_converter,\n",
+    "    draw_bboxes, \n",
+    "    Vec2Box\n",
+    ")\n",
     "from yolo.config.config import NMSConfig"
    ]
   },
@@ -49,6 +57,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "with open(MODEL_CONFIG) as stream:\n",
+    "    cfg_model = OmegaConf.load(stream)\n",
     "if os.path.exists(TRT_WEIGHT_PATH):\n",
     "    from torch2trt import TRTModule\n",
     "\n",
@@ -57,8 +67,6 @@
     "else:\n",
     "    from torch2trt import torch2trt\n",
     "\n",
-    "    with open(MODEL_CONFIG) as stream:\n",
-    "        cfg_model = OmegaConf.load(stream)\n",
     "\n",
     "    model = create_model(cfg_model, weight_path=WEIGHT_PATH)\n",
     "    model = model.to(device).eval()\n",
@@ -70,7 +78,7 @@
     "    logger.info(f\"📥 TensorRT model saved to oonx.pt\")\n",
     "\n",
     "transform = AugmentationComposer([], IMAGE_SIZE)\n",
-    "vec2box = Vec2Box(model_trt, IMAGE_SIZE, device)\n"
+    "converter = create_converter(cfg_model.name, model_trt, cfg_model.anchor, IMAGE_SIZE, device)\n"
    ]
   },
   {
@@ -79,7 +87,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "image, bbox = transform(image, torch.zeros(0, 5))\n",
+    "image, bbox, rev_tensor = transform(image, torch.zeros(0, 5))\n",
     "image = image.to(device)[None]"
    ]
   },
@@ -91,7 +99,7 @@
    "source": [
     "with torch.no_grad():\n",
     "    predict = model_trt(image)\n",
-    "    predict = vec2box(predict[\"Main\"])\n",
+    "    predict = converter(predict[\"Main\"])\n",
     "predict_box = bbox_nms(predict[0], predict[2], NMSConfig(0.5, 0.5))\n",
     "draw_bboxes(image, predict_box)"
    ]
@@ -122,7 +130,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.1.undefined"
+   "version": "3.10.14"
   }
  },
  "nbformat": 4,

examples/notebook_inference.ipynb

@@ -1,5 +1,15 @@
 {
  "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -35,7 +45,7 @@
    "source": [
     "CONFIG_PATH = \"../yolo/config\"\n",
     "CONFIG_NAME = \"config\"\n",
-    "MODEL = \"v7-base\"\n",
+    "MODEL = \"v9-c\"\n",
     "\n",
     "DEVICE = 'cuda:0'\n",
     "CLASS_NUM = 80\n",
@@ -54,7 +64,9 @@
     "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}\", f\"model={MODEL}\"])\n",
     "    model = create_model(cfg.model, class_num=CLASS_NUM).to(device)\n",
+    "\n",
     "    transform = AugmentationComposer([], cfg.image_size)\n",
+    "\n",
     "    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)\n",
     "    post_proccess = PostProccess(converter, cfg.task.nms)"
    ]
@@ -81,7 +93,7 @@
     "    predict = model(image)\n",
     "    pred_bbox = post_proccess(predict, rev_tensor)\n",
     "\n",
-    "draw_bboxes(pil_image, pred_bbox, idx2label=cfg.class_list)"
+    "draw_bboxes(pil_image, pred_bbox, idx2label=cfg.dataset.class_list)"
    ]
   },
   {
@@ -92,6 +104,11 @@
     "\n",
     "![image](../demo/images/output/visualize.png)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {

examples/notebook_smallobject.ipynb

@@ -30,7 +30,17 @@
     "project_root = Path().resolve().parent\n",
     "sys.path.append(str(project_root))\n",
     "\n",
-    "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes, Vec2Box, NMSConfig, PostProccess"
+    "from yolo import (\n",
+    "    AugmentationComposer, \n",
+    "    Config, \n",
+    "    NMSConfig, \n",
+    "    PostProccess,\n",
+    "    bbox_nms, \n",
+    "    create_model, \n",
+    "    create_converter, \n",
+    "    custom_logger, \n",
+    "    draw_bboxes, \n",
+    ")"
    ]
   },
   {
@@ -62,8 +72,8 @@
     "    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)\n",
-    "    post_proccess = PostProccess(vec2box, NMSConfig(0.5, 0.9))\n",
+    "    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)\n",
+    "    post_proccess = PostProccess(converter, NMSConfig(0.5, 0.9))\n",
     "    "
    ]
   },
@@ -112,7 +122,7 @@
     "with torch.no_grad():\n",
     "    total_image, total_shift = slide_image(image)\n",
     "    predict = model(total_image)\n",
-    "    pred_class, _, pred_bbox = vec2box(predict[\"Main\"])\n",
+    "    pred_class, _, pred_bbox = converter(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)\n",
@@ -126,7 +136,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "draw_bboxes(pil_image, predict_box, idx2label=cfg.class_list)"
+    "draw_bboxes(pil_image, predict_box, idx2label=cfg.dataset.class_list)"
    ]
   },
   {

examples/sample_inference.py

@@ -2,29 +2,39 @@ import sys
 from pathlib import Path
 
 import hydra
-import torch
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
-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
-from yolo.utils.logging_utils import custom_logger, validate_log_directory
 
+from yolo import (
+    Config,
+    FastModelLoader,
+    ModelTester,
+    ProgressLogger,
+    create_converter,
+    create_dataloader,
+    create_model,
+)
+from yolo.utils.model_utils import get_device
 
-@hydra.main(config_path="../yolo/config", config_name="config", version_base=None)
-def main(cfg: Config):
-    custom_logger()
-    save_path = validate_log_directory(cfg, cfg.name)
-    dataloader = create_dataloader(cfg)
-
-    device = torch.device(cfg.device)
-    model = create_model(cfg).to(device)
 
-    tester = ModelTester(cfg, model, save_path, device)
-    tester.solve(dataloader)
+@hydra.main(config_path="config", config_name="config", version_base=None)
+def main(cfg: Config):
+    progress = ProgressLogger(cfg, exp_name=cfg.name)
+    device, use_ddp = get_device(cfg.device)
+    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
+    if getattr(cfg.task, "fast_inference", False):
+        model = FastModelLoader(cfg).load_model(device)
+    else:
+        model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
+        model = model.to(device)
+
+    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
+
+    solver = ModelTester(cfg, model, converter, progress, device)
+    progress.start()
+    solver.solve(dataloader)
 
 
 if __name__ == "__main__":

examples/sample_train.py

@@ -2,29 +2,35 @@ import sys
 from pathlib import Path
 
 import hydra
-import torch
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
-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 ModelTrainer
-from yolo.utils.logging_utils import custom_logger, validate_log_directory
 
+from yolo import (
+    Config,
+    ModelTrainer,
+    ProgressLogger,
+    create_converter,
+    create_dataloader,
+    create_model,
+)
+from yolo.utils.model_utils import get_device
 
-@hydra.main(config_path="../yolo/config", config_name="config", version_base=None)
+
+@hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    custom_logger()
-    save_path = validate_log_directory(cfg, cfg.name)
-    dataloader = create_dataloader(cfg)
-    # TODO: get_device or rank, for DDP mode
-    device = torch.device(cfg.device)
-    model = create_model(cfg).to(device)
-
-    trainer = ModelTrainer(cfg, model, save_path, device)
-    trainer.solve(dataloader, cfg.task.epoch)
+    progress = ProgressLogger(cfg, exp_name=cfg.name)
+    device, use_ddp = get_device(cfg.device)
+    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
+    model = create_model(cfg.model, class_num=cfg.dataset.class_num, weight_path=cfg.weight)
+    model = model.to(device)
+
+    converter = create_converter(cfg.model.name, model, cfg.model.anchor, cfg.image_size, device)
+
+    solver = ModelTrainer(cfg, model, converter, progress, device)
+    progress.start()
+    solver.solve(dataloader)
 
 
 if __name__ == "__main__":

yolo/__init__.py

@@ -5,12 +5,13 @@ 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, create_converter
 from yolo.utils.deploy_utils import FastModelLoader
-from yolo.utils.logging_utils import custom_logger
+from yolo.utils.logging_utils import ProgressLogger, custom_logger
 from yolo.utils.model_utils import PostProccess
 
 all = [
     "create_model",
     "Config",
+    "ProgressLogger",
     "NMSConfig",
     "custom_logger",
     "validate_log_directory",