🔀 [Merge] branch 'DEPLOY' into SETUP

README.md

@@ -47,7 +47,7 @@ pip install -r requirements.txt
 | ------------------ | :---------: | :-------: | :-------: |
 | PyTorch            | v1.12      | v2.3+     | v1.12     |
 | ONNX               | ✅         | ✅        | -         |
-| TensorRT           | 🧪         | 🧪        | -         |
+| TensorRT           | ✅         | -        | -         |
 | OpenVINO           | -          | 🧪        | ❔        |
 
 </td></tr> </table>

examples/notebook_TensorRT.ipynb

@@ -0,0 +1,130 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import torch\n",
+    "from PIL import Image \n",
+    "from loguru import logger\n",
+    "from omegaconf import OmegaConf\n",
+    "\n",
+    "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.config.config import NMSConfig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MODEL = \"v9-c\"\n",
+    "DEVICE = \"cuda:0\"\n",
+    "\n",
+    "WEIGHT_PATH = f\"../weights/{MODEL}.pt\" \n",
+    "TRT_WEIGHT_PATH = f\"../weights/{MODEL}.trt\"\n",
+    "MODEL_CONFIG = f\"../yolo/config/model/{MODEL}.yaml\"\n",
+    "\n",
+    "IMAGE_PATH = \"../demo/images/inference/image.png\"\n",
+    "IMAGE_SIZE = (640, 640)\n",
+    "\n",
+    "custom_logger()\n",
+    "device = torch.device(DEVICE)\n",
+    "image = Image.open(IMAGE_PATH)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if os.path.exists(TRT_WEIGHT_PATH):\n",
+    "    from torch2trt import TRTModule\n",
+    "\n",
+    "    model_trt = TRTModule()\n",
+    "    model_trt.load_state_dict(torch.load(TRT_WEIGHT_PATH))\n",
+    "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",
+    "\n",
+    "    dummy_input = torch.ones((1, 3, 640, 640)).to(device)\n",
+    "    logger.info(f\"♻️ Creating TensorRT model\")\n",
+    "    model_trt = torch2trt(model, [dummy_input])\n",
+    "    torch.save(model_trt.state_dict(), TRT_WEIGHT_PATH)\n",
+    "    logger.info(f\"📥 TensorRT model saved to oonx.pt\")\n",
+    "\n",
+    "transform = AugmentationComposer([], IMAGE_SIZE)\n",
+    "vec2box = Vec2Box(model_trt, IMAGE_SIZE, device)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image, bbox = transform(image, torch.zeros(0, 5))\n",
+    "image = image.to(device)[None]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with torch.no_grad():\n",
+    "    predict = model_trt(image)\n",
+    "    predict = vec2box(predict[\"Main\"])\n",
+    "predict_box = bbox_nms(predict[0], predict[2], NMSConfig(0.5, 0.5))\n",
+    "draw_bboxes(image, predict_box)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Sample Output:\n",
+    "\n",
+    "![image](../demo/images/output/visualize.png)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "yolomit",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.1.undefined"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

yolo/config/model/v9-c.yaml

@@ -1,6 +1,6 @@
 anchor:
   reg_max: 16
-  anchors: [8, 16, 32]
+  strides: [8, 16, 32]
 
 model:
   backbone:

yolo/lazy.py

@@ -25,7 +25,7 @@ def main(cfg: Config):
         model = FastModelLoader(cfg).load_model()
     else:
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
-        model = model.to(device)
+    model = model.to(device)
 
     vec2box = Vec2Box(model, cfg.image_size, device)
 

yolo/model/module.py

@@ -105,7 +105,7 @@ class Anchor2Vec(nn.Module):
     def forward(self, anchor_x: Tensor) -> Tensor:
         anchor_x = rearrange(anchor_x, "B (P R) h w -> B R P h w", P=4)
         vector_x = anchor_x.softmax(dim=1)
-        vector_x = self.anc2vec(vector_x).squeeze(1)
+        vector_x = self.anc2vec(vector_x)[:, 0]
         return anchor_x, vector_x
 
 

yolo/model/yolo.py

@@ -26,13 +26,15 @@ class YOLO(nn.Module):
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.model: List[YOLOLayer] = nn.ModuleList()
         self.build_model(model_cfg.model)
+        self.strides = getattr(model_cfg.anchor, "strides", None)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
         output_dim, layer_idx = [3], 1
         logger.info(f"🚜 Building YOLO")
         for arch_name in model_arch:
-            logger.info(f"  🏗️  Building {arch_name}")
+            if model_arch[arch_name]:
+                logger.info(f"  🏗️  Building {arch_name}")
             for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=layer_idx):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
@@ -45,7 +47,6 @@ class YOLO(nn.Module):
                     layer_args["in_channels"] = output_dim[source]
                 if "Detection" in layer_type:
                     layer_args["in_channels"] = [output_dim[idx] for idx in source]
-                if "Detection" in layer_type or "Anchor2Box" in layer_type:
                     layer_args["num_classes"] = self.num_classes
 
                 # create layers
@@ -135,6 +136,7 @@ def create_model(model_cfg: ModelConfig, weight_path: Optional[str], class_num:
         if os.path.exists(weight_path):
             # TODO: fix map_location
             model.model.load_state_dict(torch.load(weight_path), strict=False)
-            logger.info("✅ Success load model weight")
-
+            logger.info("✅ Success load model & weight")
+    else:
+        logger.info("✅ Success load model")
     return model

yolo/tools/drawer.py

@@ -13,7 +13,7 @@ def draw_bboxes(
     img: Union[Image.Image, torch.Tensor],
     bboxes: List[List[Union[int, float]]],
     *,
-    idx2label: Optional[list],
+    idx2label: Optional[list] = None,
 ):
     """
     Draw bounding boxes on an image.
@@ -47,7 +47,7 @@ def draw_bboxes(
         draw.rounded_rectangle(bbox, outline=(*color_map, 200), radius=5, width=2)
         draw.rounded_rectangle(bbox, fill=(*color_map, 100), radius=5)
 
-        class_text = str(idx2label[int(class_id)] if idx2label else class_id)
+        class_text = str(idx2label[int(class_id)] if idx2label else int(class_id))
         label_text = f"{class_text}" + (f" {conf[0]: .0%}" if conf else "")
 
         text_bbox = font.getbbox(label_text)

yolo/utils/bounding_box_utils.py

@@ -9,6 +9,7 @@ from torch import Tensor
 from torchvision.ops import batched_nms
 
 from yolo.config.config import MatcherConfig, ModelConfig, NMSConfig
+from yolo.model.yolo import YOLO
 
 
 def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
@@ -264,8 +265,8 @@ class BoxMatcher:
 
 
 class Vec2Box:
-    def __init__(self, model, image_size, device, anchors: list = None):
-        if anchors is None:
+    def __init__(self, model: YOLO, image_size, device):
+        if model.strides is None:
             logger.info("🧸 Found no anchor, Make a dummy test for auto-anchor size")
             dummy_input = torch.zeros(1, 3, *image_size).to(device)
             dummy_output = model(dummy_input)
@@ -274,7 +275,8 @@ class Vec2Box:
                 _, _, *anchor_num = predict_head[2].shape
                 anchors_num.append(anchor_num)
         else:
-            anchors_num = [[image_size[0] / anchor, image_size[0] / anchor] for anchor in anchors]
+            logger.info(f"🈶 Found anchor {model.strides}")
+            anchors_num = [[image_size[0] // stride, image_size[0] // stride] for stride in model.strides]
         anchor_grid, scaler = generate_anchors(image_size, anchors_num)
         self.anchor_grid, self.scaler = anchor_grid.to(device), scaler.to(device)
         self.anchor_norm = (anchor_grid / scaler[:, None])[None].to(device)

yolo/utils/deploy_utils.py

@@ -30,9 +30,7 @@ class FastModelLoader:
             return self._load_trt_model()
         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
-        )
+        return create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight)
 
     def _load_onnx_model(self):
         from onnxruntime import InferenceSession
@@ -91,9 +89,9 @@ class FastModelLoader:
         from torch2trt import torch2trt
 
         model = create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).eval()
-        dummy_input = torch.ones((1, 3, *self.cfg.image_size))
+        dummy_input = torch.ones((1, 3, *self.cfg.image_size)).cuda()
         logger.info(f"♻️ Creating TensorRT model")
-        model_trt = torch2trt(model, [dummy_input])
+        model_trt = torch2trt(model.cuda(), [dummy_input])
         torch.save(model_trt.state_dict(), self.model_path)
         logger.info(f"📥 TensorRT model saved to {self.model_path}")
         return model_trt