|
|
from typing import List, Union |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
from loguru import logger |
|
|
from PIL import Image, ImageDraw, ImageFont |
|
|
from torchvision.transforms.functional import to_pil_image |
|
|
|
|
|
|
|
|
def draw_bboxes(img: Union[Image.Image, torch.Tensor], bboxes: List[List[Union[int, float]]]): |
|
|
""" |
|
|
Draw bounding boxes on an image. |
|
|
|
|
|
Args: |
|
|
- img (PIL Image or torch.Tensor): Image on which to draw the bounding boxes. |
|
|
- bboxes (List of Lists/Tensors): Bounding boxes with [class_id, x_min, y_min, x_max, y_max], |
|
|
where coordinates are normalized [0, 1]. |
|
|
""" |
|
|
|
|
|
if isinstance(img, torch.Tensor): |
|
|
if img.dim() > 3: |
|
|
logger.info("Multi-frame tensor detected, using the first image.") |
|
|
img = img[0] |
|
|
bboxes = bboxes[0] |
|
|
img = to_pil_image(img) |
|
|
|
|
|
draw = ImageDraw.Draw(img) |
|
|
width, height = img.size |
|
|
font = ImageFont.load_default(30) |
|
|
|
|
|
for bbox in bboxes: |
|
|
class_id, x_min, y_min, x_max, y_max = bbox |
|
|
x_min = x_min * width |
|
|
x_max = x_max * width |
|
|
y_min = y_min * height |
|
|
y_max = y_max * height |
|
|
shape = [(x_min, y_min), (x_max, y_max)] |
|
|
draw.rectangle(shape, outline="red", width=3) |
|
|
draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue") |
|
|
|
|
|
img.save("visualize.jpg") |
|
|
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") |
|
|
|
