LaiEthanLai
commited on
✅ [Pass] tests, skip drawing if graphviz not found (#21)
Browse filesHi Ethan,
Thank you for your contribution! Your changes look great, and I appreciate the addition of the condition to skip drawing if Graphviz is not found. This will definitely help in environments where Graphviz isn't available.
I have reviewed the commit, and everything seems to be in order. All tests have passed as well, which is great. I'll go ahead and merge this PR.
Thanks again for your efforts!
Best regards,
Henry Tsui
- yolo/model/yolo.py +1 -1
- yolo/tools/drawer.py +4 -2
yolo/model/yolo.py
CHANGED
@@ -126,5 +126,5 @@ def get_model(cfg: Config) -> YOLO:
|
|
126 |
model = YOLO(cfg.model, cfg.hyper.data.class_num)
|
127 |
logger.info("✅ Success load model")
|
128 |
log_model_structure(model.model)
|
129 |
-
|
130 |
return model
|
|
|
126 |
model = YOLO(cfg.model, cfg.hyper.data.class_num)
|
127 |
logger.info("✅ Success load model")
|
128 |
log_model_structure(model.model)
|
129 |
+
draw_model(model=model)
|
130 |
return model
|
yolo/tools/drawer.py
CHANGED
@@ -95,6 +95,8 @@ def draw_model(*, model_cfg=None, model=None, v7_base=False):
|
|
95 |
for jdx in range(idx, model_size):
|
96 |
if model_mat[idx, jdx]:
|
97 |
dot.edge(str(idx), str(jdx))
|
98 |
-
|
99 |
-
|
|
|
|
|
100 |
logger.info("🎨 Drawing Model Architecture at Model-arch.png")
|
|
|
95 |
for jdx in range(idx, model_size):
|
96 |
if model_mat[idx, jdx]:
|
97 |
dot.edge(str(idx), str(jdx))
|
98 |
+
try:
|
99 |
+
dot.render("Model-arch", format="png", cleanup=True)
|
100 |
+
except:
|
101 |
+
logger.info("Warning: Could not find graphviz backend, continue without drawing the model architecture")
|
102 |
logger.info("🎨 Drawing Model Architecture at Model-arch.png")
|