🔀 [Merge] branch 'main' into INFERENCE

yolo/model/yolo.py

@@ -126,5 +126,5 @@ def get_model(cfg: Config) -> YOLO:
     model = YOLO(cfg.model, cfg.hyper.data.class_num)
     logger.info("✅ Success load model")
     log_model_structure(model.model)
-    # draw_model(model=model)
+    draw_model(model=model)
     return model

yolo/tools/dataset_preparation.py

@@ -61,7 +61,7 @@ def prepare_dataset(cfg):
             extract_to = os.path.join(data_dir, data_type) if data_type != "annotations" else data_dir
             final_place = os.path.join(extract_to, dataset_type)
 
-            os.makedirs(extract_to, exist_ok=True)
+            os.makedirs(final_place, exist_ok=True)
             if check_files(final_place, dataset_args.get("file_num")):
                 logger.info(f"✅ Dataset {dataset_type: <12} already verified.")
                 continue
@@ -78,7 +78,7 @@ if __name__ == "__main__":
     import sys
 
     sys.path.append("./")
-    from tools.log_helper import custom_logger
+    from utils.logging_utils import custom_logger
 
     custom_logger()
     prepare_dataset()

yolo/tools/drawer.py

@@ -95,6 +95,8 @@ def draw_model(*, model_cfg=None, model=None, v7_base=False):
         for jdx in range(idx, model_size):
             if model_mat[idx, jdx]:
                 dot.edge(str(idx), str(jdx))
-
-    dot.render("Model-arch", format="png", cleanup=True)
+    try:
+        dot.render("Model-arch", format="png", cleanup=True)
+    except:
+        logger.info("Warning: Could not find graphviz backend, continue without drawing the model architecture")
     logger.info("🎨 Drawing Model Architecture at Model-arch.png")