Rename load_model to load_yolov5_model

yolort/utils/update_module_state.py

@@ -5,6 +5,7 @@
 from torch import nn
 
 from yolort.models import yolo
+from yolort.v5 import load_yolov5_model
 
 
 ARCHITECTURE_MAPS = {
@@ -41,9 +42,7 @@ def update_module_state_from_ultralytics(
             Default: True.
         verbose (bool): print all information to screen. Default: True.
     """
-
-    from yolort.v5 import load_model
-    model = load_model(model_path, autoshape=False, verbose=verbose)
+    model = load_yolov5_model(model_path, autoshape=False, verbose=verbose)
 
     key_arch = f'{arch}_{feature_fusion_type.lower()}_v4.0'
     if key_arch not in ARCHITECTURE_MAPS:
@@ -141,7 +140,7 @@ def attach_parameters_heads(self, state_dict, name):
 def rgetattr(obj, attr, *args):
     """
     Nested version of getattr.
-    See <https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects>
+    Ref: https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects
     """
     def _getattr(obj, attr):
         return getattr(obj, attr, *args)

yolort/v5/__init__.py

@@ -1,3 +1,3 @@
 from .models import *
 from .utils import *
-from .helper import add_yolov5_context, load_model
+from .helper import add_yolov5_context, load_yolov5_model

yolort/v5/helper.py

@@ -7,7 +7,7 @@
 from .models.yolo import Model
 from .utils import attempt_download, set_logging
 
-__all__ = ['add_yolov5_context', 'load_model']
+__all__ = ['add_yolov5_context', 'load_yolov5_model']
 
 
 @contextlib.contextmanager
@@ -27,7 +27,7 @@ def add_yolov5_context():
         sys.path.remove(path_ultralytics_yolov5)
 
 
-def load_model(model_path: str, autoshape: bool = False, verbose: bool = True):
+def load_yolov5_model(model_path: str, autoshape: bool = False, verbose: bool = True):
     """
     Creates a specified YOLOv5 model
 

yolort/v5/utils/general.py

@@ -87,7 +87,8 @@ def methods(instance):
 def set_logging(rank=-1, verbose=True):
     logging.basicConfig(
         format="%(message)s",
-        level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
+        level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN,
+    )
 
 
 def init_seeds(seed=0):