🔀 [Merge] branch 'DEPLOY' into TEST

WongKinYiu · Jul 31, 2024 · 5c8e313 · 5c8e313
2 parents 8a2b98c + f1bfd74
commit 5c8e313
Show file tree

Hide file tree

Showing 7 changed files with 108 additions and 56 deletions.
diff --git a/demo/hf_demo.py b/demo/hf_demo.py
@@ -11,7 +11,7 @@
     AugmentationComposer,
     NMSConfig,
     PostProccess,
-    Vec2Box,
+    create_converter,
     create_model,
     draw_bboxes,
 )
@@ -25,22 +25,22 @@ def load_model(model_name, device):
     model_cfg.model.auxiliary = {}
     model = create_model(model_cfg, True)
     model.to(device).eval()
-    return model
+    return model, model_cfg
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = load_model(DEFAULT_MODEL, device)
-v2b = Vec2Box(model, IMAGE_SIZE, device)
-class_list = OmegaConf.load("yolo/config/general.yaml").class_list
+model, model_cfg = load_model(DEFAULT_MODEL, device)
+converter = create_converter(model_cfg.name, model, model_cfg.anchor, IMAGE_SIZE, device)
+class_list = OmegaConf.load("yolo/config/dataset/coco.yaml").class_list
 
 transform = AugmentationComposer([])
 
 
 def predict(model_name, image, nms_confidence, nms_iou):
-    global DEFAULT_MODEL, model, device, v2b, class_list, post_proccess
+    global DEFAULT_MODEL, model, device, converter, class_list, post_proccess
     if model_name != DEFAULT_MODEL:
-        model = load_model(model_name, device)
-        v2b = Vec2Box(model, IMAGE_SIZE, device)
+        model, model_cfg = load_model(model_name, device)
+        converter = create_converter(model_cfg.name, model, model_cfg.anchor, IMAGE_SIZE, device)
         DEFAULT_MODEL = model_name
 
     image_tensor, _, rev_tensor = transform(image)
@@ -49,7 +49,7 @@ def predict(model_name, image, nms_confidence, nms_iou):
     rev_tensor = rev_tensor.to(device)[None]
 
     nms_config = NMSConfig(nms_confidence, nms_iou)
-    post_proccess = PostProccess(v2b, nms_config)
+    post_proccess = PostProccess(converter, nms_config)
 
     with torch.no_grad():
         predict = model(image_tensor)

diff --git a/examples/notebook_TensorRT.ipynb b/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,

diff --git a/examples/notebook_inference.ipynb b/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": {

diff --git a/examples/notebook_smallobject.ipynb b/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)"
    ]
   },
   {

diff --git a/examples/sample_inference.py b/examples/sample_inference.py
@@ -2,29 +2,39 @@
 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__":

diff --git a/examples/sample_train.py b/examples/sample_train.py
@@ -2,29 +2,35 @@
 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__":

diff --git a/yolo/__init__.py b/yolo/__init__.py
@@ -5,12 +5,13 @@
 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",