🎨 [Add] Visualization of YOLO model

examples/example_train.py

@@ -13,6 +13,7 @@ from yolo.model.yolo import get_model
 from yolo.tools.log_helper import custom_logger
 from yolo.tools.trainer import Trainer
 from yolo.utils.dataloader import get_dataloader
+from yolo.utils.drawer import draw_model
 from yolo.utils.get_dataset import prepare_dataset
 
 
@@ -23,6 +24,7 @@ def main(cfg: Config):
 
     dataloader = get_dataloader(cfg)
     model = get_model(cfg.model)
+    draw_model(model=model)
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 

yolo/utils/drawer.py

@@ -1,5 +1,6 @@
 from typing import List, Union
 
+import numpy as np
 import torch
 from loguru import logger
 from PIL import Image, ImageDraw, ImageFont
@@ -39,3 +40,56 @@ def draw_bboxes(img: Union[Image.Image, torch.Tensor], bboxes: List[List[Union[i
 
     img.save("visualize.jpg")  # Save the image with annotations
     logger.info("Saved visualize image at visualize.png")
+
+
+def draw_model(*, model_cfg=None, model=None):
+    from graphviz import Digraph
+
+    if model_cfg:
+        from yolo.model.yolo import get_model
+
+        model = get_model(model_cfg)
+    elif model is None:
+        raise ValueError("Drawing Object is None")
+
+    model_size = len(model.model)
+    model_mat = np.zeros((model_size, model_size), dtype=bool)
+
+    layer_name = []
+    for idx, layer in enumerate(model.model):
+        layer_name.append(str(type(layer)).split(".")[-1][:-2])
+        if isinstance(layer.source, int):
+            source = layer.source + (layer.source < 0) * idx
+            model_mat[source, idx] = True
+        else:
+            for source in layer.source:
+                source = source + (source < 0) * idx
+                model_mat[source, idx] = True
+
+    pattern_list = [("ELAN", 8, 3), ("ELAN", 8, 55), ("MP", 5, 11)]
+    pattern_mat = []
+    for name, size, position in pattern_list:
+        pattern_mat.append(
+            (name, size, model_mat[position : position + size, position + 1 : position + 1 + size].copy())
+        )
+
+    dot = Digraph(comment="Model Flow Chart")
+    node_idx = 0
+
+    for idx in range(model_size):
+        for jdx in range(idx, model_size - 7):
+            for name, size, pattern in pattern_mat:
+                if (model_mat[idx : idx + size, jdx : jdx + size] == pattern).all():
+                    layer_name[idx] = name
+                    model_mat[idx : idx + size, jdx : jdx + size] = False
+                    model_mat[idx, idx + size] = True
+
+        if model_mat[idx].any():
+            dot.node(str(idx), f"{node_idx}-{layer_name[idx]}")
+            node_idx += 1
+        for jdx in range(idx, model_size):
+            if model_mat[idx, jdx] == 1:
+                dot.edge(str(idx), str(jdx))
+
+    dot.render("Model-arch", format="png", cleanup=True)
+    logger.info("🎨 Drawing Model Architecture at Model-arch.png")