Add yolov5 v6.0 pretrained model weights

hubconf.py

@@ -1,4 +1,13 @@
 # Optional list of dependencies required by the package
 dependencies = ["torch", "torchvision"]
 
-from yolort.models import yolov5n, yolov5s, yolov5m, yolov5l, yolov5ts
+from yolort.models import (
+    yolov5n,
+    yolov5n6,
+    yolov5s,
+    yolov5s6,
+    yolov5m,
+    yolov5m6,
+    yolov5l,
+    yolov5ts,
+)

test/test_models.py

@@ -357,7 +357,12 @@ def test_torchscript(arch):
 
 @pytest.mark.parametrize(
     "arch, version, upstream_version, hash_prefix",
-    [("yolov5s", "r4.0", "v4.0", "9ca9a642")],
+    [
+        ("yolov5s", "r4.0", "v4.0", "9ca9a642"),
+        ("yolov5n", "r6.0", "v6.0", "649e089f"),
+        ("yolov5s", "r6.0", "v6.0", "c3b140f3"),
+        ("yolov5n6", "r6.0", "v6.0", "beecbbae"),
+    ],
 )
 def test_load_from_yolov5(
     arch: str,
@@ -376,16 +381,25 @@ def test_load_from_yolov5(
         checkpoint_path,
         hash_prefix=hash_prefix,
     )
+    score_thresh = 0.25
 
-    model_yolov5 = YOLOv5.load_from_yolov5(checkpoint_path, version=version)
+    model_yolov5 = YOLOv5.load_from_yolov5(
+        checkpoint_path,
+        score_thresh=score_thresh,
+        version=version,
+    )
     model_yolov5.eval()
     out_from_yolov5 = model_yolov5.predict(img_path)
     assert isinstance(out_from_yolov5[0], dict)
     assert isinstance(out_from_yolov5[0]["boxes"], Tensor)
     assert isinstance(out_from_yolov5[0]["labels"], Tensor)
     assert isinstance(out_from_yolov5[0]["scores"], Tensor)
 
-    model = models.__dict__[arch](pretrained=True, score_thresh=0.25)
+    model = models.__dict__[arch](
+        upstream_version=version,
+        pretrained=True,
+        score_thresh=score_thresh,
+    )
     model.eval()
     out = model.predict(img_path)
 

tools/yolov5_to_yolort.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
+import argparse
+from pathlib import Path
+
+from yolort.utils import convert_yolov5_to_yolort
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(
+        "Convert checkpoints from yolov5 to yolort", add_help=True
+    )
+
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        required=True,
+        help="Path of the checkpoint weights",
+    )
+    parser.add_argument(
+        "--version",
+        type=str,
+        default="r6.0",
+        help="Upstream version released by the ultralytics/yolov5, Possible "
+        "values are ['r3.1', 'r4.0', 'r6.0']. Default: 'r6.0'.",
+    )
+    # Dataset Configuration
+    parser.add_argument(
+        "--image_path",
+        type=str,
+        default="./test/assets/zidane.jpg",
+        help="Path of the test image",
+    )
+
+    parser.add_argument(
+        "--output_path", type=str, default=None, help="Path where to save"
+    )
+    return parser
+
+
+def cli_main():
+    parser = get_parser()
+    args = parser.parse_args()
+    print(f"Command Line Args: {args}")
+    checkpoint_path = Path(args.checkpoint_path)
+    assert checkpoint_path.exists(), f"Not found checkpoint file at '{checkpoint_path}'"
+
+    if args.output_path is None:
+        args.output_path = checkpoint_path.parent
+    output_path = Path(args.output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    convert_yolov5_to_yolort(checkpoint_path, output_path, version=args.version)
+
+
+if __name__ == "__main__":
+    cli_main()

yolort/models/__init__.py

@@ -8,7 +8,18 @@
 from .yolo import YOLO
 from .yolo_module import YOLOv5
 
-__all__ = ["YOLO", "YOLOv5", "yolov5n", "yolov5s", "yolov5m", "yolov5l", "yolov5ts"]
+__all__ = [
+    "YOLO",
+    "YOLOv5",
+    "yolov5n",
+    "yolov5n6",
+    "yolov5s",
+    "yolov5s6",
+    "yolov5m",
+    "yolov5m6",
+    "yolov5l",
+    "yolov5ts",
+]
 
 
 def yolov5n(

yolort/models/yolo.py

@@ -20,6 +20,13 @@
     "yolov5_darknet_pan_s_r40",
     "yolov5_darknet_pan_m_r40",
     "yolov5_darknet_pan_l_r40",
+    "yolov5_darknet_pan_n_r60",
+    "yolov5_darknet_pan_n6_r60",
+    "yolov5_darknet_pan_s_r60",
+    "yolov5_darknet_pan_s6_r60",
+    "yolov5_darknet_pan_m_r60",
+    "yolov5_darknet_pan_m6_r60",
+    "yolov5_darknet_pan_l_r60",
     "yolov5_darknet_tan_s_r40",
     "build_model",
 ]
@@ -221,21 +228,6 @@ def load_from_yolov5(
         return model
 
 
-model_urls_root = "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0"
-
-model_urls = {
-    # Path Aggregation Network
-    "yolov5_darknet_pan_s_r31_coco": f"{model_urls_root}/yolov5_darknet_pan_s_r31_coco-eb728698.pt",
-    "yolov5_darknet_pan_m_r31_coco": f"{model_urls_root}/yolov5_darknet_pan_m_r31_coco-670dc553.pt",
-    "yolov5_darknet_pan_l_r31_coco": f"{model_urls_root}/yolov5_darknet_pan_l_r31_coco-4dcc8209.pt",
-    "yolov5_darknet_pan_s_r40_coco": f"{model_urls_root}/yolov5_darknet_pan_s_r40_coco-e3fd213d.pt",
-    "yolov5_darknet_pan_m_r40_coco": f"{model_urls_root}/yolov5_darknet_pan_m_r40_coco-d295cb02.pt",
-    "yolov5_darknet_pan_l_r40_coco": f"{model_urls_root}/yolov5_darknet_pan_l_r40_coco-4416841f.pt",
-    # Tranformer Attention Network
-    "yolov5_darknet_tan_s_r40_coco": f"{model_urls_root}/yolov5_darknet_tan_s_r40_coco-fe1069ce.pt",
-}
-
-
 def build_model(
     backbone_name: str,
     depth_multiple: float,
@@ -280,6 +272,34 @@ def build_model(
     return model
 
 
+model_urls_root_r30 = (
+    "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0"
+)
+model_urls_root_r52 = (
+    "https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.5.2-alpha"
+)
+
+model_urls = {
+    # Path Aggregation Network 3.1 and 4.0
+    "yolov5_darknet_pan_s_r31_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_s_r31_coco-eb728698.pt",
+    "yolov5_darknet_pan_m_r31_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_m_r31_coco-670dc553.pt",
+    "yolov5_darknet_pan_l_r31_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_l_r31_coco-4dcc8209.pt",
+    "yolov5_darknet_pan_s_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_s_r40_coco-e3fd213d.pt",
+    "yolov5_darknet_pan_m_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_m_r40_coco-d295cb02.pt",
+    "yolov5_darknet_pan_l_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_l_r40_coco-4416841f.pt",
+    # Path Aggregation Network 6.0
+    "yolov5_darknet_pan_n_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n_r60_coco-bc15659e.pt",
+    "yolov5_darknet_pan_n6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n6_r60_coco-4e823e0f.pt",
+    "yolov5_darknet_pan_s_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s_r60_coco-9f44bf3f.pt",
+    "yolov5_darknet_pan_s6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s6_r60_coco-b4ff1fc2.pt",
+    "yolov5_darknet_pan_m_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m_r60_coco-58d32352.pt",
+    "yolov5_darknet_pan_m6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m6_r60_coco-cc010533.pt",
+    "yolov5_darknet_pan_l_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_l_r60_coco-321d8dcd.pt",
+    # Tranformer Attention Network
+    "yolov5_darknet_tan_s_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_tan_s_r40_coco-fe1069ce.pt",
+}
+
+
 def yolov5_darknet_pan_s_r31(
     pretrained: bool = False,
     progress: bool = True,

yolort/utils/__init__.py

@@ -7,14 +7,15 @@
 
 from .hooks import FeatureExtractor
 from .image_utils import cv2_imshow, get_image_from_url, read_image_to_tensor
-from .update_module_state import load_from_ultralytics
+from .update_module_state import convert_yolov5_to_yolort, load_from_ultralytics
 
 
 __all__ = [
     "FeatureExtractor",
     "cv2_imshow",
     "get_image_from_url",
     "get_callable_dict",
+    "convert_yolov5_to_yolort",
     "load_from_ultralytics",
     "load_state_dict_from_url",
     "read_image_to_tensor",

yolort/utils/update_module_state.py

@@ -2,14 +2,45 @@
 from functools import reduce
 from typing import List, Dict, Optional
 
+import torch
 from torch import nn
 
 from yolort.models import yolo
 from yolort.v5 import load_yolov5_model, get_yolov5_size
 from .image_utils import to_numpy
 
 
-def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"):
+def convert_yolov5_to_yolort(
+    checkpoint_path: str,
+    output_path: str,
+    version: str = "r6.0",
+    prefix: str = "yolov5_darknet_pan",
+    postfix: str = "custom.pt",
+):
+    """
+    Convert model checkpoint trained with ultralytics/yolov5 to yolort.
+
+    Args:
+        checkpoint_path (str): Path of the YOLOv5 checkpoint model.
+        output_path (str): Path of the converted yolort checkpoint model.
+        version (str): upstream version released by the ultralytics/yolov5, Possible
+            values are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0".
+        prefix (str): The prefix string of the saved model. Default: "yolov5_darknet_pan".
+        postfix (str): The postfix string of the saved model. Default: "custom.pt".
+    """
+    model_info = load_from_ultralytics(checkpoint_path, version=version)
+    model_state_dict = model_info["state_dict"]
+
+    size = model_info["size"]
+    use_p6 = "6" if model_info["use_p6"] else ""
+    output_postfix = f"{prefix}_{size}{use_p6}_{version.replace('.', '')}_{postfix}"
+    torch.save(model_state_dict, output_path / output_postfix)
+
+
+def load_from_ultralytics(
+    checkpoint_path: str,
+    version: str = "r6.0",
+):
     """
     Allows the user to load model state file from the checkpoint trained from
     the ultralytics/yolov5.
@@ -85,7 +116,7 @@ def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"):
         use_p6=use_p6,
     )
     module_state_updater.updating(checkpoint_yolov5)
-    state_dict = module_state_updater.model.state_dict()
+    state_dict = module_state_updater.model.half().state_dict()
 
     size = get_yolov5_size(depth_multiple, width_multiple)