From 054a1c787369182e367c3c0d9b0d98ecfd119cc0 Mon Sep 17 00:00:00 2001 From: henrytsui000 Date: Wed, 31 Jul 2024 12:18:23 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=F0=9F=90=9B=20[Fix]=20#56=20bugs,=20create?= =?UTF-8?q?=5Fconverter=20->=20Vec2Box?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- demo/hf_demo.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/demo/hf_demo.py b/demo/hf_demo.py index 749e345..73ce1c5 100644 --- 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) From f1bfd742d210522981f21a889ee80fb97ee5c752 Mon Sep 17 00:00:00 2001 From: henrytsui000 Date: Wed, 31 Jul 2024 12:35:40 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=90=9B=20[Fix]=20#56=20bugs,=20the=20?= =?UTF-8?q?using=20converter=20in=20ipynb?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/notebook_TensorRT.ipynb | 22 ++++++++++----- examples/notebook_inference.ipynb | 21 +++++++++++++-- examples/notebook_smallobject.ipynb | 20 ++++++++++---- examples/sample_inference.py | 42 ++++++++++++++++++----------- examples/sample_train.py | 38 +++++++++++++++----------- yolo/__init__.py | 3 ++- 6 files changed, 99 insertions(+), 47 deletions(-) diff --git a/examples/notebook_TensorRT.ipynb b/examples/notebook_TensorRT.ipynb index 87ff346..b75e031 100644 --- 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 index 8bb3cc6..ddb7dcf 100644 --- 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 index c025b52..c2f75ab 100644 --- 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 index 2574515..7b2b245 100644 --- 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 index 2c86209..dc91718 100644 --- 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 index 8641398..900de97 100644 --- 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",