LaiEthanLai commited on
Commit
b18186e
·
unverified ·
1 Parent(s): a0cddf7

✅ [Pass] tests, skip drawing if graphviz not found (#21)

Browse files

Hi 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

Files changed (2) hide show
  1. yolo/model/yolo.py +1 -1
  2. 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
- # draw_model(model=model)
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
- dot.render("Model-arch", format="png", cleanup=True)
 
 
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")